Skip to content

fix: Make TokenizedPrompt.MultiModalFeatures per-prompt#1633

Open
albertoperdomo2 wants to merge 8 commits into
llm-d:mainfrom
albertoperdomo2:fix/per-prompt-mm-features
Open

fix: Make TokenizedPrompt.MultiModalFeatures per-prompt#1633
albertoperdomo2 wants to merge 8 commits into
llm-d:mainfrom
albertoperdomo2:fix/per-prompt-mm-features

Conversation

@albertoperdomo2

@albertoperdomo2 albertoperdomo2 commented Jun 12, 2026

Copy link
Copy Markdown
Contributor

What type of PR is this?

/kind bug
/kind feature

What this PR does / why we need it:

What this PR does

Introduces a PromptTokens struct that bundles token IDs and multimodal features per prompt, replacing the two parallel slices on TokenizedPrompt (PerPromptTokens [][]uint32 and MultiModalFeatures []MultiModalFeature). Renames TokenizedPrompt to TokenizedRequest to distinguish the request-level container from the new per-prompt type.

Before:

type TokenizedPrompt struct {
    PerPromptTokens    [][]uint32
    MultiModalFeatures []MultiModalFeature  // flat, wrong
    CacheSalt          string
}

After:

type TokenizedRequest struct {
    Prompts   []PromptTokens
    CacheSalt string
}

type PromptTokens struct {
    TokenIDs           []uint32
    MultiModalFeatures []MultiModalFeature
}

The association between tokens and their multimodal features is now structural, not positional. Consumers no longer need bounds checks or guards like if len(tp.PerPromptTokens) == 1 to pair features with the correct token array.

The InferenceRequestBody field is renamed from TokenizedPrompt to TokenizedRequest to match the type. The data-key identifiers (TokenizedPromptDataKey, tokenizedPromptKeyID) are intentionally unchanged since they are plugin registry keys, not type names.

Consumers updated:

  • blockkeys.go: Iterates tp.Prompts directly; the len(tp.PerPromptTokens) == 1 guard disappears because each PromptTokens carries its own features.
  • multimodal/producer.go (ExtractMMItems): Iterates tp.Prompts accessing p.MultiModalFeatures.
  • disagg/multimodal_helpers.go (hasMultimodalContent): Checks len(p.MultiModalFeatures) > 0 per prompt.
  • approximateprefix/hashing.go: Iterates tp.Prompts using p.TokenIDs.
  • vllmgrpc.go: Constructs TokenizedRequest with Prompts: []PromptTokens{{...}} for pre-tokenized gRPC inputs.

Why we need it

PerPromptTokens was [][]uint32 (per-prompt) but MultiModalFeatures was []MultiModalFeature (flat). MM feature offsets are only valid relative to one token array, but the flat shape forced consumers to assume which prompt the features belonged to. The blockkeys.go consumer had an explicit len(tp.PerPromptTokens) == 1 guard to avoid misapplying features to multi-prompt requests.

Bundling tokens and features into PromptTokens eliminates this assumption and the guard, and makes the per-prompt association structural rather than requiring lockstep indexing of parallel slices.

Known follow-up

flattenMMFeatures in scorer/preciseprefixcache/legacy_producer.go and convertMMFeaturesToUpstream in tokenizer/tokenizer.go are now identical (same input type, same output type, same logic). Deduplicating them is out of scope here since it crosses package boundaries and the legacy producer is deprecated.

Which issue(s) this PR fixes:

Fixes #1584

Release note (write NONE if no user-facing change):

NONE

@github-actions github-actions Bot added size/L Denotes a PR that changes 100-499 lines, ignoring generated files. kind/bug Categorizes issue or PR as related to a bug. kind/feature Categorizes issue or PR as related to a new feature. labels Jun 12, 2026
@albertoperdomo2

Copy link
Copy Markdown
Contributor Author

cc: @ahg-g @liu-cong

Comment on lines +115 to +119
PerPromptTokens [][]uint32
// MultiModalFeatures holds one entry per multimodal item in prompt order.
// Nil if the prompt contains no multimodal content. Offsets are relative
// to PerPromptTokens[0] (always single-prompt when multimodal content is
// present).
MultiModalFeatures []MultiModalFeature
// MultiModalFeatures holds multimodal items per prompt, indexed in
// lockstep with PerPromptTokens. Single-prompt requests use a length-1
// outer slice. Nil if the prompt contains no multimodal content.
MultiModalFeatures [][]MultiModalFeature

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

it is a bigger change, but creating a PromptTokens struct that encapsulates the tokens and the mm features is more idiomatic

type TokenizedPrompt struct {
... 
  PromptTokens []PromptToken
...
}

type  struct {
  Tokens []uint32
   MultiModalFeatures []MultiModalFeature
}

but we probably need better names, TokenizedPrompt and PromptToken are too close, perhaps ask Claude / Antigravity to suggest better names?

@github-actions github-actions Bot added size/XL Denotes a PR that changes 500-999 lines, ignoring generated files. kind/bug Categorizes issue or PR as related to a bug. kind/feature Categorizes issue or PR as related to a new feature. and removed size/L Denotes a PR that changes 100-499 lines, ignoring generated files. kind/feature Categorizes issue or PR as related to a new feature. kind/bug Categorizes issue or PR as related to a bug. labels Jun 15, 2026
@albertoperdomo2

Copy link
Copy Markdown
Contributor Author

@ahg-g anything blocking us on this one?

abdallahsamabd pushed a commit to abdallahsamabd/llm-d-inference-scheduler that referenced this pull request Jul 7, 2026
@albertoperdomo2
albertoperdomo2 force-pushed the fix/per-prompt-mm-features branch from 26dffe1 to e781181 Compare July 8, 2026 10:48
@albertoperdomo2
albertoperdomo2 requested a review from ahg-g July 8, 2026 10:48
Comment thread pkg/epp/framework/plugins/requestcontrol/dataproducer/prefixhash/hashing.go Outdated
@bongwoobak

Copy link
Copy Markdown
Member

burstprefix (not in this PR, pulled in by today's main-merge #1727) still uses the old TokenizedPrompt

@albertoperdomo2
albertoperdomo2 force-pushed the fix/per-prompt-mm-features branch 2 times, most recently from c6a1c9a to 4175bc7 Compare July 17, 2026 10:20
Signed-off-by: Alberto Perdomo <aperdomo@redhat.com>
Signed-off-by: Alberto Perdomo <aperdomo@redhat.com>
Signed-off-by: Alberto Perdomo <aperdomo@redhat.com>
Signed-off-by: Alberto Perdomo <aperdomo@redhat.com>
The auto-merge left one test using the old TokenizedPrompt type and
PerPromptTokens field in producer_test.go. Updated to use
TokenizedRequest with Prompts/PromptTokens.

Signed-off-by: Alberto Perdomo <aperdomo@redhat.com>
Signed-off-by: Alberto Perdomo <aperdomo@redhat.com>
@albertoperdomo2
albertoperdomo2 force-pushed the fix/per-prompt-mm-features branch from 18f3cdc to ee25953 Compare July 24, 2026 09:44
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

kind/bug Categorizes issue or PR as related to a bug. kind/feature Categorizes issue or PR as related to a new feature. size/XL Denotes a PR that changes 500-999 lines, ignoring generated files.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Make MultiModalFeatures per-prompt in TokenizedPrompt

3 participants