Skip to content

Comments

Move maxengine for restructuring#3203

Open
hengtaoguo wants to merge 1 commit intomainfrom
hengtaoguo-re
Open

Move maxengine for restructuring#3203
hengtaoguo wants to merge 1 commit intomainfrom
hengtaoguo-re

Conversation

@hengtaoguo
Copy link
Collaborator

@hengtaoguo hengtaoguo commented Feb 20, 2026

Description

Restructure inference-related modules to align with the new maxtext directory. Moved four files and created two __init__.py as below:

  • src/maxtext/inference/maxengine/:
    • __init___.py.
    • maxengine.py
    • maxengine_config.py
    • maxengine_server.py
  • src/maxtext/inference/mlperf/microbenchmarks/
    • __init___.py
    • benchmark_chunked_prefill.py
  • Update imports/shell scripts/comments/docs throughout the repo.

Tests

CI.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

  • I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
  • I have necessary comments in my code, particularly in hard-to-understand areas.
  • I have run end-to-end tests tests and provided workload links above if applicable.
  • I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

@codecov
Copy link

codecov bot commented Feb 21, 2026

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

Copy link
Collaborator

@bvandermoon bvandermoon left a comment

Choose a reason for hiding this comment

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

Do we need shims to keep supporting the old decode and maxengine commands?

from MaxText.common_types import DECODING_ACTIVE_SEQUENCE_INDICATOR, MODEL_MODE_PREFILL
from maxtext.layers import quantizations
from MaxText.maxengine import MaxEngine
from maxtext.inference.maxengine import maxengine
Copy link
Collaborator

Choose a reason for hiding this comment

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

Was this extra import intended?

Comment on lines -25 to 31
from MaxText import maxengine
from MaxText import pyconfig
from maxtext.common import profiler
from maxtext.common.gcloud_stub import jetstream, is_decoupled
from maxtext.inference.maxengine import maxengine
from maxtext.multimodal import processor as mm_processor
from maxtext.multimodal import utils as mm_utils
from maxtext.utils import max_utils
Copy link
Collaborator

Choose a reason for hiding this comment

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

This isn't related to your PR, but can this be moved to inference/?

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants