Initial commit

Author: jrplatin

.github/CODEOWNERS

@@ -18,9 +18,9 @@ src/MaxText/elastic_train.py @lukebaumann @shauryagup @richjames0 @shralex
 src/MaxText/layers/quantizations.py @khatwanimohit @jshin1394 @liudangyi @richjames0 @shralex
 
 # Inference
-src/MaxText/tests/inference @vipannalla @mitalisi @gpolovets1 @mailvijayasingh @jrplatin @patemotter @lumosis @richjames0
-src/MaxText/inference @vipannalla @mitalisi @gpolovets1 @mailvijayasingh @jrplatin @patemotter @lumosis @richjames0
-src/MaxText/inference_mlperf @vipannalla @mitalisi @gpolovets1 @mailvijayasingh @jrplatin @patemotter @lumosis @richjames0
+src/maxtext/tests/inference @vipannalla @mitalisi @gpolovets1 @mailvijayasingh @jrplatin @patemotter @lumosis @richjames0
+src/maxtext/inference @vipannalla @mitalisi @gpolovets1 @mailvijayasingh @jrplatin @patemotter @lumosis @richjames0
+src/maxtext/inference_mlperf @vipannalla @mitalisi @gpolovets1 @mailvijayasingh @jrplatin @patemotter @lumosis @richjames0
 
 # Dockerfiles and dependencies
 *.Dockerfile @bvandermoon @parambole @richjames0 @shralex

.vscode/launch.json

@@ -37,7 +37,7 @@
       "python": "python3",
       "module": "MaxText.decode",
       "args": ["src/MaxText/configs/base.yml",
-               "run_name=runner_$(date +%Y-%m-%d-%H-%M)", 
+               "run_name=runner_$(date +%Y-%m-%d-%H-%M)",
                "base_output_directory=gs://test-maxtext-output",
                "dataset_path=gs://test-maxtext-dataset",
                "steps=2",
@@ -53,7 +53,7 @@
       "python": "python3",
       "module": "MaxText.train",
       "args": ["src/MaxText/configs/base.yml",
-               "run_name=runner_$(date +%Y-%m-%d-%H-%M)", 
+               "run_name=runner_$(date +%Y-%m-%d-%H-%M)",
                "base_output_directory=gs://test-maxtext-output",
                "dataset_path=gs://test-maxtext-dataset",
                "steps=2",
@@ -66,7 +66,7 @@
       "console": "integratedTerminal",
       "justMyCode": false,
       "python": "python3",
-      "module": "MaxText.inference_microbenchmark",
+      "module": "maxtext.inference_microbenchmark",
       "args": [
         "src/MaxText/configs/base.yml",
         "model_name=llama2-7b",
@@ -82,7 +82,7 @@
         "inference_microbenchmark_prefill_lengths=32,64,128,256,512,1024",
         "inference_microbenchmark_stages=generate",
         "inference_microbenchmark_loop_iters=1",
-        "run_name=runner_$(date +%Y-%m-%d-%H-%M)", 
+        "run_name=runner_$(date +%Y-%m-%d-%H-%M)",
         "base_output_directory=gs://test-maxtext-output",
         "prefill_cache_axis_order=0,2,1,3",
         "ar_cache_axis_order=0,2,1,3",

codecov.yml

@@ -24,7 +24,7 @@
 # During scheduled runs, the 'regular' flag is carried forward from the last PR.
 
 # Exclude non-source code, deprecated and experimental folders from coverage tracking
-codecov: 
+codecov:
  token: 35742a22-fb1f-4839-97ff-b54da5588689
 # By default file names in the coverage report will have their path in the file system, which in our
 # runners would be /__w/maxtext/maxtext/src/MaxText/* but Codecov expects src/MaxText/* so we need to fix the path
@@ -36,8 +36,8 @@ ignore:
   - "src/MaxText/configs"
   - "src/MaxText/examples"
   - "src/MaxText/experimental"
-  - "src/MaxText/inference"
-  - "src/MaxText/inference_mlperf"
+  - "src/maxtext/inference"
+  - "src/maxtext/inference_mlperf"
   - "src/MaxText/scratch_code"
   - "src/MaxText/distillation" # code moved to src/maxtext/trainers/post_train/distillation
   - "src/MaxText/sft" # code moved to src/maxtext/trainers/post_train/sft

docs/guides/optimization/pallas_kernels_performance.md

@@ -103,7 +103,7 @@
 # for converting PyTorch code to JAX
 block_for_rag = [
     "src/MaxText/layers",  # Neural network layers and building blocks
-    "src/MaxText/inference",  # Inference and prediction code
+    "src/maxtext/inference",  # Inference and prediction code
     "src/MaxText/common_types.py",  # Common data types and structures
     "src/MaxText/maxtext_utils.py",  # Utility functions and helpers
 ]

src/MaxText/experimental/agent/integrative_rag_agent/config.py

@@ -93,7 +93,7 @@ def visit_Attribute(self, node):
       if base_name in self.git_aliases:
         # It's an external dependency. We need to format it with the attribute path.
         # Example: base_name='page_manager', attr_chain='PageState'
-        # self.git_dependencies['page_manager'] might be 'src/MaxText/inference/page_manager.py#page_manager'
+        # self.git_dependencies['page_manager'] might be 'src/maxtext/inference/page_manager.py#page_manager'
         path, obj = self.git_dependencies[base_name].split("#", 1)
 
         # As per the user request, we append the attribute access to the object name.
@@ -197,9 +197,9 @@ def convert_package_to_path(self, path):
     """Convert an absolute import line to a mapping of names to file anchors.
 
     Example:
-        "from MaxText.inference import page_manager, utils" ->
-        {"page_manager": "src/MaxText/inference.py#page_manager",
-         "utils": "src/MaxText/inference.py#utils"}
+        "from maxtext.inference import page_manager, utils" ->
+        {"page_manager": "src/maxtext/inference.py#page_manager",
+         "utils": "src/maxtext/inference.py#utils"}
 
     Args:
         path (str): A normalized absolute import string.
@@ -215,8 +215,8 @@ def convert_package_to_path(self, path):
       # or a module 'pkg' corresponds to 'path_form/pkg.py'
       # The logic in get_absolute_imports should ideally resolve this ambiguity.
       # A heuristic could be used here (e.g., checking casing) but we stick to the current logic.
-      # The user's example `from MaxText.inference import page_manager` creates a path
-      # `src/MaxText/inference.py#page_manager`, which is what the new visitor expects to correct.
+      # The user's example `from maxtext.inference import page_manager` creates a path
+      # `src/maxtext/inference.py#page_manager`, which is what the new visitor expects to correct.
       import_dict[pkg.strip()] = path_form + ".py#" + pkg.strip()
     return import_dict
 

src/MaxText/experimental/agent/orchestration_agent/split_python_file.py

@@ -75,7 +75,7 @@
 from MaxText import train_utils
 from MaxText import pyconfig
 from MaxText.utils import gcs_utils
-from MaxText.inference import offline_engine
+from maxtext.inference import offline_engine
 from MaxText.experimental.rl import grpo_input_pipeline
 from MaxText.experimental.rl import grpo_utils
 from MaxText.globals import EPS

src/MaxText/experimental/rl/grpo_trainer.py

@@ -24,7 +24,7 @@
 from MaxText import max_logging
 from MaxText import max_utils
 from MaxText.common_types import DecoderBlockType
-from MaxText.inference.offline_engine import InputData
+from maxtext.inference.offline_engine import InputData
 
 from pathwaysutils.experimental import reshard as experimental_reshard
 from pathwaysutils.experimental import split_by_mesh_axis

src/MaxText/experimental/rl/grpo_utils.py

@@ -61,10 +61,10 @@
     AttentionType,
     DEFAULT_MASK_VALUE,
 )
-from MaxText.inference import kvcache
-from MaxText.inference import page_manager
-from MaxText.inference import paged_attention
-from MaxText.inference.kvcache import KVQuant
+from maxtext.inference import kvcache
+from maxtext.inference import page_manager
+from maxtext.inference import paged_attention
+from maxtext.inference.kvcache import KVQuant
 from MaxText.sharding import create_sharding
 from MaxText.layers import nnx_wrappers
 from MaxText.layers.attentions import Attention

src/MaxText/layers/attention_mla.py

@@ -69,8 +69,8 @@
     Q_LENGTH,
     Q_LENGTH_NO_EXP,
 )
-from MaxText.inference import page_manager
-from MaxText.inference.kvcache import KVQuant, KVTensor
+from maxtext.inference import page_manager
+from maxtext.inference.kvcache import KVQuant, KVTensor
 from MaxText.kernels import jax_flash_attention
 from MaxText.kernels.ragged_attention import ragged_gqa
 from MaxText.kernels.ragged_attention import ragged_mha