Skip to content

Commit a2edf6c

Browse files
namgyu-younahg-g
andauthored
feat: video token estimation formula for prefix caching (#1408)
* add video/audio token extimation for prefix-caching approximation Signed-off-by: namgyu-youn <namgyu.dev@gmail.com> * revert: audio-related Signed-off-by: namgyu-youn <namgyu.dev@gmail.com> * update: token estimation formula Signed-off-by: namgyu-youn <namgyu.dev@gmail.com> * update Signed-off-by: namgyu-youn <namgyu.dev@gmail.com> * rename video config & drop dead-code Signed-off-by: namgyu-youn <namgyu.dev@gmail.com> * update Signed-off-by: namgyu-youn <namgyu.dev@gmail.com> * fix: video token estimiation Signed-off-by: namgyu-youn <namgyu.dev@gmail.com> * update: use metadata directly with fallback Signed-off-by: namgyu-youn <namgyu.dev@gmail.com> * update Signed-off-by: namgyu-youn <namgyu.dev@gmail.com> * video: replace duration*fps formula with Qwen3-VL-accurate estimate The simple formula is off by orders of magnitude past the 32-frame vLLM cap; this mirrors vLLM's frame sampling, smart_resize, and timestamp-token overhead instead. Signed-off-by: namgyu-youn <namgyu.dev@gmail.com> * video: trim redundant comments in config docs Signed-off-by: namgyu-youn <namgyu.dev@gmail.com> * video: drop tunable config, use built-in constants Video token estimation parameters are architecture constants, not per-deployment settings. Signed-off-by: namgyu-youn <namgyu.dev@gmail.com> * add UT for base64 test Signed-off-by: namgyu-youn <namgyu.dev@gmail.com> * doc: token estimation constant and formula Signed-off-by: namgyu-youn <namgyu.dev@gmail.com> * doc: token estimation formula Signed-off-by: namgyu-youn <namgyu.dev@gmail.com> --------- Signed-off-by: namgyu-youn <namgyu.dev@gmail.com> Co-authored-by: Abdullah Gharaibeh <40361897+ahg-g@users.noreply.github.com>
1 parent b5f8159 commit a2edf6c

7 files changed

Lines changed: 652 additions & 4 deletions

File tree

.typos.toml

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,8 @@ IST = "IST"
77
unmarshaling = "unmarshaling"
88
# image_grid_thw is a vLLM/CUDA tensor shape parameter
99
thw = "thw"
10+
# trak is the ISO/MP4 container's track atom name
11+
trak = "trak"
1012

1113
[files]
1214
extend-exclude = [

pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/estimate.go

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ const blockTypeText = "text"
4141
// so pairing this backend with the engine-correlated scorer yields misses, not bad routes.
4242
type estimateBackend struct {
4343
img imageEstimator
44+
vid videoEstimator
4445
}
4546

4647
func (b estimateBackend) produce(ctx context.Context, body *fwkrh.InferenceRequestBody) (*fwkrh.TokenizedPrompt, error) {
@@ -159,7 +160,7 @@ func (b estimateBackend) appendChatMessage(out []byte, features []fwkrh.MultiMod
159160
case "image_url":
160161
out, features = appendMMAsset(out, features, fwkrh.ModalityImage, block.ImageURL.URL, b.img.placeholderCount(block.ImageURL.URL))
161162
case "video_url":
162-
out, features = appendMMAsset(out, features, fwkrh.ModalityVideo, block.VideoURL.URL, assetPlaceholderCount(len(block.VideoURL.URL)))
163+
out, features = appendMMAsset(out, features, fwkrh.ModalityVideo, block.VideoURL.URL, b.vid.placeholderCount(block.VideoURL.URL))
163164
case "input_audio", "audio_url":
164165
data := block.InputAudio.Data + block.InputAudio.Format
165166
out, features = appendMMAsset(out, features, fwkrh.ModalityAudio, data, assetPlaceholderCount(len(data)))

pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/mm_estimate.go

Lines changed: 222 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ import (
2020
"bytes"
2121
"encoding/base64"
2222
"image"
23+
"math"
2324
"strings"
2425

2526
// Registers decoders so image.DecodeConfig can read dimensions.
@@ -148,3 +149,224 @@ func imageDimensionsFromBase64Payload(rawB64 string) (width, height int, ok bool
148149
}
149150
return cfg.Width, cfg.Height, true
150151
}
152+
153+
const (
154+
// Qwen3-VL / Gemma4 model defaults — also match vLLM VideoMediaIO defaults.
155+
// vLLM never feeds the model more than defaultVideoMaxFrames raw frames: a
156+
// video decoded to more frames than this is pre-sampled down to exactly
157+
// defaultVideoMaxFrames (see sampleFrames), so estimation must reproduce
158+
// that cap rather than the source frame count.
159+
defaultVideoMaxFrames = 32
160+
// defaultVideoSampleFPS is the resampling rate applied to videos that decode
161+
// to defaultVideoMaxFrames or fewer frames: vLLM re-derives the frame count
162+
// as roughly totalFrames/srcFPS*defaultVideoSampleFPS instead of using the
163+
// source fps directly (see sampleFrames).
164+
defaultVideoSampleFPS = 2.0
165+
// defaultVideoPatchSize is the ViT patch edge in pixels: smartResize rounds
166+
// each frame's height and width to a multiple of this before computing grids.
167+
defaultVideoPatchSize = 16
168+
// defaultVideoMergeSize is the spatial merge window: an NxN block of patches
169+
// (N = defaultVideoMergeSize) collapses into a single visual token.
170+
defaultVideoMergeSize = 2
171+
// defaultVideoTemporalPatch is the temporal merge window: this many
172+
// consecutive sampled frames collapse into a single time step (grid cell)
173+
// for both patch counting and timestamp averaging.
174+
defaultVideoTemporalPatch = 2
175+
// defaultVideoMaxPixels is the total pixel budget shared across all sampled
176+
// frames: smartResize divides it by nframes to get a per-frame budget, then
177+
// downscales height/width to fit before rounding to the patch grid.
178+
defaultVideoMaxPixels = 25_165_824
179+
180+
// Fallback values used when MP4 metadata (duration, fps, resolution)
181+
// cannot be extracted from the payload, so totalFrames = duration*fps and
182+
// the resolution passed to smartResize both fall back to model-typical values.
183+
defaultVideoFallbackDuration = 16.0
184+
defaultVideoFallbackFPS = 2.0
185+
defaultVideoFallbackWidth = 640
186+
defaultVideoFallbackHeight = 360
187+
188+
// minVideoFrames and maxVideoFrames clamp the resampled frame count
189+
// (see sampleFrames) so extremely short or long videos still produce a
190+
// sane number of temporal grid cells.
191+
minVideoFrames = 4
192+
maxVideoFrames = 768
193+
)
194+
195+
// videoEstimator implements the Qwen3-VL-accurate token estimation:
196+
// vLLM frame cap, smart_resize, temporal patch grouping, and per-grid timestamp tokens.
197+
// It has no tunable fields: every parameter is a Qwen3-VL/vLLM architecture
198+
// constant, not a per-deployment setting.
199+
// Validated against live Qwen3-VL-30B-A3B-Instruct API responses on the
200+
// Video-MME dataset: github.com/llm-d/llm-d-router/pull/1408#issuecomment-4697915077
201+
type videoEstimator struct{}
202+
203+
// newVideoEstimator returns a videoEstimator using built-in constants.
204+
func newVideoEstimator() videoEstimator {
205+
return videoEstimator{}
206+
}
207+
208+
// placeholderCount estimates placeholder tokens for a video URL. Base64 data URLs
209+
// are decoded for duration, fps, and resolution; other URLs fall back to defaults.
210+
func (e videoEstimator) placeholderCount(url string) int {
211+
duration, fps, width, height := defaultVideoFallbackDuration, defaultVideoFallbackFPS, defaultVideoFallbackWidth, defaultVideoFallbackHeight
212+
213+
if strings.HasPrefix(url, "data:video/") && strings.Contains(url, "base64,") {
214+
idx := strings.Index(url, "base64,")
215+
decoded, err := base64.StdEncoding.DecodeString(url[idx+len("base64,"):])
216+
if err == nil {
217+
if meta, err := parseMP4Metadata(decoded); err == nil {
218+
duration = meta.Duration
219+
fps = meta.FPS
220+
if meta.Width > 0 {
221+
width = meta.Width
222+
}
223+
if meta.Height > 0 {
224+
height = meta.Height
225+
}
226+
}
227+
}
228+
}
229+
230+
totalFrames := int(duration * fps)
231+
nframes := e.sampleFrames(totalFrames, fps)
232+
hBar, wBar := e.smartResize(nframes, height, width)
233+
return e.visualTokens(totalFrames, fps, nframes, hBar, wBar)
234+
}
235+
236+
// sampleFrames mirrors the vLLM VideoMediaIO frame-selection logic:
237+
// videos with more than maxFrames total frames are pre-sampled to exactly
238+
// maxFrames; shorter videos are resampled to sampleFPS and clamped.
239+
func (e videoEstimator) sampleFrames(totalFrames int, srcFPS float64) int {
240+
if totalFrames > defaultVideoMaxFrames {
241+
return defaultVideoMaxFrames
242+
}
243+
if srcFPS <= 0 {
244+
srcFPS = defaultVideoFallbackFPS
245+
}
246+
n := int(float64(totalFrames) / srcFPS * defaultVideoSampleFPS)
247+
if n < minVideoFrames {
248+
return minVideoFrames
249+
}
250+
if n > maxVideoFrames {
251+
return maxVideoFrames
252+
}
253+
return n
254+
}
255+
256+
// smartResize scales h and w to fit within the per-frame pixel budget and rounds
257+
// each dimension to the nearest patchSize multiple using banker's rounding
258+
// (round half to even) to match Python's round() and the HF transformers output.
259+
func (e videoEstimator) smartResize(nframes, h, w int) (int, int) {
260+
pixPerFrame := defaultVideoMaxPixels / nframes
261+
if pixPerFrame < 1 {
262+
pixPerFrame = 1
263+
}
264+
if h*w > pixPerFrame {
265+
scale := math.Sqrt(float64(pixPerFrame) / float64(h*w))
266+
h = int(math.Round(float64(h) * scale))
267+
w = int(math.Round(float64(w) * scale))
268+
}
269+
hBar := roundHalfEven(h, defaultVideoPatchSize)
270+
wBar := roundHalfEven(w, defaultVideoPatchSize)
271+
if hBar < defaultVideoPatchSize {
272+
hBar = defaultVideoPatchSize
273+
}
274+
if wBar < defaultVideoPatchSize {
275+
wBar = defaultVideoPatchSize
276+
}
277+
return hBar, wBar
278+
}
279+
280+
// roundHalfEven rounds n to the nearest multiple of factor using banker's
281+
// rounding (round half to nearest even quotient), matching Python's round().
282+
func roundHalfEven(n, factor int) int {
283+
q := n / factor
284+
rem := n % factor
285+
half := factor / 2
286+
switch {
287+
case rem < half:
288+
return q * factor
289+
case rem > half:
290+
return (q + 1) * factor
291+
default: // rem == half: round to nearest even quotient
292+
if q%2 == 0 {
293+
return q * factor
294+
}
295+
return (q + 1) * factor
296+
}
297+
}
298+
299+
// visualTokens computes the total visual placeholder count.
300+
//
301+
// Formula, with gridT = ceil(nframes/temporalPatch) and hBar, wBar the
302+
// smartResize output:
303+
//
304+
// tokens = floor(gridT*(hBar/patchSize)*(wBar/patchSize)/mergeSize^2) + 2*gridT + timestampTokens(gridT)
305+
func (e videoEstimator) visualTokens(totalFrames int, srcFPS float64, nframes, hBar, wBar int) int {
306+
tp := defaultVideoTemporalPatch
307+
ms := defaultVideoMergeSize
308+
ps := defaultVideoPatchSize
309+
tBar := ((nframes + tp - 1) / tp) * tp
310+
gridT := tBar / tp
311+
gridH := hBar / ps
312+
gridW := wBar / ps
313+
patchTokens := gridT * gridH * gridW / (ms * ms)
314+
return patchTokens + gridT*2 + e.timestampTokens(totalFrames, srcFPS, nframes, gridT)
315+
}
316+
317+
// timestampTokens computes the "<X.X seconds>" token overhead per temporal grid
318+
// cell. Each timestamp costs 6 tokens plus 1 for two-digit seconds and 1 more
319+
// for three-digit seconds, mirroring the Qwen3-VL chat-template formatting.
320+
func (e videoEstimator) timestampTokens(totalFrames int, srcFPS float64, nframes, gridT int) int {
321+
timestamps := make([]float64, 0, gridT)
322+
tp := defaultVideoTemporalPatch
323+
if totalFrames > defaultVideoMaxFrames {
324+
// vLLM pre-sampled nframes evenly from [0, totalFrames-1]; average pairs.
325+
frameIdx := videoLinspace(0, totalFrames-1, nframes)
326+
rawTS := make([]float64, nframes)
327+
for i, idx := range frameIdx {
328+
rawTS[i] = float64(idx) / srcFPS
329+
}
330+
for j := 0; j < gridT; j++ {
331+
base := j * tp
332+
var sum float64
333+
for k := 0; k < tp; k++ {
334+
sum += rawTS[base+k]
335+
}
336+
timestamps = append(timestamps, sum/float64(tp))
337+
}
338+
} else {
339+
duration := float64(totalFrames) / srcFPS
340+
for i := 0; i < gridT; i++ {
341+
timestamps = append(timestamps, (float64(i)+0.5)*duration/float64(gridT))
342+
}
343+
}
344+
total := 0
345+
for _, t := range timestamps {
346+
total += 6
347+
if t >= 10 {
348+
total++
349+
}
350+
if t >= 100 {
351+
total++
352+
}
353+
}
354+
return total
355+
}
356+
357+
// videoLinspace returns n evenly spaced integer indices from start to end inclusive,
358+
// matching numpy.linspace(start, end, n, dtype=int).
359+
func videoLinspace(start, end, n int) []int {
360+
if n <= 0 {
361+
return nil
362+
}
363+
result := make([]int, n)
364+
if n == 1 {
365+
result[0] = start
366+
return result
367+
}
368+
for i := range result {
369+
result[i] = start + int(math.Round(float64(i)*float64(end-start)/float64(n-1)))
370+
}
371+
return result
372+
}

0 commit comments

Comments
 (0)