feat: add BaseLogitsProcessor core interface

[bhuvan002]

Overview:

This PR adds a base logits processor class to use with different dynamo backends.

Details:

This PR adds a base logits processor class that is backend (vllm/trtllm) agnostic and can be easily extended to implement custom logit processing logic. Additional backend specific utilities and wrappers will be added in later PRs.

Where should the reviewer start?

Related Issues: (use one of the action keywords Closes / Fixes / Resolves / Relates to)

• closes GitHub issue: #xxx

Summary by CodeRabbit

• New Features
  • Added a public BaseLogitsProcessor interface for logits processing, enabling users to implement and plug in custom processors across supported backends.
  • Simplified imports by exposing the interface at the package level for direct access.

[grahamking]

Looks reasonable to me.

[bhuvan002]

Looks reasonable to me.

Thank you. I was going to add the TRT-LLM plumbing to this PR too, but I will just merge this first and add backend specific things in follow up PRs in the interest of keeping PRs small.

[coderabbitai]

Walkthrough

Adds a new logits_processing package initializer and a base protocol defining a standard call interface for logits processors. Exposes BaseLogitsProcessor at the package level via all and re-export. No runtime logic added.

Changes

Cohort / File(s)	Summary
Package init for logits processing lib/bindings/python/src/dynamo/logits_processing/__init__.py	New package initializer with license header, docstring, re-exports BaseLogitsProcessor from .base, and sets all accordingly.
Base protocol definition lib/bindings/python/src/dynamo/logits_processing/base.py	Introduces BaseLogitsProcessor Protocol with signature __call__(input_ids: List[int], logits: torch.Tensor) -> torch.Tensor; includes docstring and type imports.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

Poem

A bunny taps the keys with grace,
A Protocol finds its rightful place;
Logits line up, neat and clear,
One interface to volunteer.
From base to init, paths align—
Hop, hop, shipped on Dynamo time! 🐇✨

Tip

🔌 Remote MCP (Model Context Protocol) integration is now available!

Pro plan users can now connect to remote MCP servers from the Integrations page. Connect with popular remote MCPs such as Notion and Linear to add more context to your reviews and chats.

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

🪧 Tips

Chat

There are 3 ways to chat with CodeRabbit:

• Review comments: Directly reply to a review comment made by CodeRabbit. Example:
  • I pushed a fix in commit <commit_id>, please review it.
  • Open a follow-up GitHub issue for this discussion.
• Files and specific lines of code (under the "Files changed" tab): Tag @coderabbitai in a new review comment at the desired location with your query.
• PR comments: Tag @coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
  • @coderabbitai gather interesting stats about this repository and render them as a table. Additionally, render a pie chart showing the language distribution in the codebase.
  • @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.

Support

Need help? Create a ticket on our support page for assistance with any issues or questions.

CodeRabbit Commands (Invoked using PR/Issue comments)

Type @coderabbitai help to get the list of available commands.

Other keywords and placeholders

• Add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.
• Add @coderabbitai summary to generate the high-level summary at a specific location in the PR description.
• Add @coderabbitai anywhere in the PR title to generate the title automatically.

Status, Documentation and Community

• Visit our Status Page to check the current availability of CodeRabbit.
• Visit our Documentation for detailed information on how to use CodeRabbit.
• Join our Discord Community to get help, request features, and share feedback.
• Follow us on X/Twitter for updates and announcements.

[coderabbitai]

Actionable comments posted: 0

🧹 Nitpick comments (4)

lib/bindings/python/src/dynamo/logits_processing/base.py (4)

24-29: Broaden type/shape contract to cover batched decode and strengthen invariants.

Many backends surface logits as either (vocab_size,) or (batch_size, vocab_size). Also, List is stricter than necessary for tokens and can hinder reuse. Recommend:

• Accept Sequence[int] for input_ids.
• Document both unbatched and batched logits shapes.
• State that return must preserve shape, dtype, and device, and avoid in-place mutation.

Apply this diff:

@@
-    def __call__(
+    def __call__(
         self,
-        input_ids: List[int],
+        input_ids: Sequence[int],
         logits: torch.Tensor,
     ) -> torch.Tensor:
@@
-        Args:
-            input_ids: The input token IDs generated so far.
-            logits: The raw logits for the next token. Shape: (vocab_size,)
+        Args:
+            input_ids: Token IDs generated so far (unbatched sequence).
+            logits: Raw next-token logits.
+                Shape: (vocab_size,) for single-request decode or
+                (batch_size, vocab_size) for batched decode.
@@
-        Returns:
-            The modified logits tensor with same shape as input.
+        Returns:
+            A tensor with the same shape, dtype, and device as `logits`.
+            Implementations should not modify `logits` in-place.

You’ll also need the import tweak below (see next comment).

Also applies to: 32-38

---

11-13: Loosen token container type: import Sequence instead of List.

Follow-up to the previous comment—switch the import to use Sequence:

-from typing import List, Protocol
+from typing import Protocol, Sequence

---

39-39: Prefer ellipsis over raising in Protocol methods.

Protocols are structural; the method body is never executed and raising here is unnecessary noise. Use an ellipsis for clarity:

-        raise NotImplementedError
+        ...

---

16-16: Optional: make the Protocol runtime-checkable if you’ll use isinstance/issubclass.

If adapters intend to perform isinstance(x, BaseLogitsProcessor) at runtime, add runtime_checkable:

-from typing import Protocol, Sequence
+from typing import Protocol, Sequence, runtime_checkable
@@
-class BaseLogitsProcessor(Protocol):
+@runtime_checkable
+class BaseLogitsProcessor(Protocol):

Skip if you won’t do runtime checks; type checkers don’t require this.

Also applies to: 11-11

📜 Review details

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

📥 Commits

Reviewing files that changed from the base of the PR and between 6a358f7 and c5a02b2.

📒 Files selected for processing (2)

• lib/bindings/python/src/dynamo/logits_processing/__init__.py (1 hunks)
• lib/bindings/python/src/dynamo/logits_processing/base.py (1 hunks)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (3)

• GitHub Check: Build and Test - dynamo
• GitHub Check: pre-merge-rust (.)
• GitHub Check: pre-merge-rust (lib/bindings/python)

🔇 Additional comments (3)

lib/bindings/python/src/dynamo/logits_processing/base.py (1)

16-22: Good foundation: clear, minimal Protocol for backend-agnostic processors.

The Protocol-based, callable interface is a sensible starting point and should compose well with adapter layers. No runtime coupling introduced.

Also applies to: 24-29

lib/bindings/python/src/dynamo/logits_processing/__init__.py (2)

11-13: Public re-export looks good.

Clean package surface: re-export via all and relative import keeps API stable.

---

1-14: The packaging configuration under lib/bindings/python has not yet been inspected for the new logits_processing module. Let’s confirm whether the existing pyproject.toml there includes the src root and picks up dynamo/logits_processing. Once we have that, we can determine if any adjustments are needed to packages or include settings to ensure downstream imports work.

I’ve run a quick inspection of the first 200 lines and relevant patterns—awaiting those results now.