Jfacevedo kimi k3 - #4967
Conversation
|
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for the Kimi K3 model architecture in MaxText, adding configuration files, the Kimi Decoupled Attention (KDA) layer, custom activations, parameter mapping, and several diagnostic scripts and unit tests. The review feedback highlights critical correctness and compatibility issues that need to be addressed. Specifically, the ShortConv1D layer lacks a caching mechanism for autoregressive decoding, which will break generation equivalence. Additionally, changes to shared modules like MlpBlock and RoutedMoE introduce potential regressions for other models by forcing float32 activations and altering gate input dimensions. Other issues include hardcoded local paths in tests, concurrency risks in the stateful MXFP4DequantizeHook, a potential IndexError in KimiLinearModel, and missing PyTorch import guards that could crash the test suite.
| # Load naive.py directly without triggering fla.ops.__init__ (which requires triton) | ||
| spec = importlib.util.spec_from_file_location( | ||
| "kda_naive", | ||
| "/Users/jfacevedo/.gemini/jetski/brain/0487c2aa-4e99-434c-b4e2-9147cc01875b/scratch/venv/lib/python3.12/site-packages/fla/ops/kda/naive.py", |
There was a problem hiding this comment.
This test contains a hardcoded absolute path to a local virtual environment directory (/Users/jfacevedo/...). This will fail on any other machine or in CI/CD pipelines. Please resolve this path dynamically or package the reference implementation/dependency properly.
| def __call__(self, x: jax.Array) -> jax.Array: | ||
| """x: [B, T, features] -> [B, T, features]""" | ||
| # Depthwise 1D conv along sequence dimension T | ||
| # Pad left by (kernel_size - 1) to maintain causal alignment | ||
| padded = jnp.pad(x, ((0, 0), (self.kernel_size - 1, 0), (0, 0))) |
There was a problem hiding this comment.
The ShortConv1D layer does not implement any caching mechanism for autoregressive decoding (where sequence length T = 1). During decoding, padding left with zeros (jnp.pad(x, ((0, 0), (self.kernel_size - 1, 0), (0, 0)))) on every step means the convolution will ignore the actual history of past tokens and assume they are all zeros. This breaks mathematical equivalence with the training path during generation, leading to severe model quality degradation. To support correct autoregressive generation, ShortConv1D must accept and return a conv state/cache (storing the last kernel_size - 1 inputs), and this state must be managed by the decoder layer and model.
| if cfg.activations_in_float32: | ||
| x = x.astype(jnp.float32) | ||
| x = _convert_to_activation_function(act_fn)(x) | ||
| x = _convert_to_activation_function(act_fn)(x.astype(jnp.float32)) |
There was a problem hiding this comment.
Forcing all activations to be computed in float32 by hardcoding x.astype(jnp.float32) in the shared MlpBlock affects all models in MaxText, not just Kimi K3. This can cause performance regressions (slower execution) and numerical changes for other models (like Llama or Gemma) that expect activations to run in bfloat16. Instead of hardcoding this, please revert to using the conditional cfg.activations_in_float32 check, and enable activations_in_float32: true in the Kimi K3 configuration files (kimi-k3.yml, kimi-k3-tiny.yml, etc.).
| x = _convert_to_activation_function(act_fn)(x.astype(jnp.float32)) | |
| if cfg.activations_in_float32: | |
| x = x.astype(jnp.float32) | |
| x = _convert_to_activation_function(act_fn)(x) |
| except ValueError as e: | ||
| if "not found in HF checkpoint index" in str(e): | ||
| return np.zeros(shape, dtype=np.float32) | ||
| raise e |
There was a problem hiding this comment.
Silently returning zeros when a key is not found in the HF checkpoint index can mask parameter mapping typos or missing weights, leading to silent model corruption. It is highly recommended to log a warning when this fallback occurs.
| except ValueError as e: | |
| if "not found in HF checkpoint index" in str(e): | |
| return np.zeros(shape, dtype=np.float32) | |
| raise e | |
| except ValueError as e: | |
| if "not found in HF checkpoint index" in str(e): | |
| max_logging.log(f"Warning: key {key} not found in HF checkpoint index. Initializing with zeros.") | |
| return np.zeros(shape, dtype=np.float32) | |
| raise e |
| class MXFP4DequantizeHook: | ||
| """Hook to dequantize MXFP4 weight_packed & weight_scale to bfloat16 for a single expert in MaxText. | ||
| Pads in_features or out_features to 7168 with zeros and caches scales for 100x speedup. | ||
| """ |
There was a problem hiding this comment.
The MXFP4DequantizeHook class relies on a mutable state (self.current_expert_idx) to track which expert is currently being dequantized. This assumes a strict, sequential, single-threaded execution order during checkpoint loading. If parallel loading (e.g., via a thread pool) is ever introduced, or if any expert is skipped or re-ordered, this stateful counter will get out of sync, leading to silent weight corruption (applying incorrect scales to experts). Consider making this hook stateless or adding explicit safeguards/assertions to ensure sequential consistency.
| # 2. Sequential Decoder Layers | ||
| kda_states = [] | ||
| for i, layer in enumerate(self.layers): | ||
| init_state = initial_kda_states[i] if initial_kda_states is not None else None |
There was a problem hiding this comment.
Potential IndexError if initial_kda_states is shorter than self.layers. Adding a safe length check prevents potential crashes if a partial list of states is passed.
| init_state = initial_kda_states[i] if initial_kda_states is not None else None | |
| init_state = initial_kda_states[i] if (initial_kda_states is not None and i < len(initial_kda_states)) else None |
| import os | ||
| import sys | ||
| import unittest |
There was a problem hiding this comment.
PyTorch import without pytest.importorskip('torch') can crash the entire test collection suite if PyTorch is not installed in the environment. Adding the import guard ensures the test is gracefully skipped.
| import os | |
| import sys | |
| import unittest | |
| import numpy as np | |
| import pytest | |
| pytest.importorskip("torch") | |
| import torch | |
| import torch.nn as nn |
| import pytest | ||
| import torch |
There was a problem hiding this comment.
PyTorch import without pytest.importorskip('torch') can crash the entire test collection suite if PyTorch is not installed in the environment. Adding the import guard ensures the test is gracefully skipped.
| import pytest | |
| import torch | |
| import pytest | |
| pytest.importorskip("torch") | |
| import torch |
| import pytest | ||
| import torch |
There was a problem hiding this comment.
PyTorch import without pytest.importorskip('torch') can crash the entire test collection suite if PyTorch is not installed in the environment. Adding the import guard ensures the test is gracefully skipped.
| import pytest | |
| import torch | |
| import pytest | |
| pytest.importorskip("torch") | |
| import torch |
| @classmethod | ||
| def setUpClass(cls): | ||
| cls.checkpoint_dir = "/Users/jfacevedo/apps/maxtext/scratch/kimi_k3_orbax_checkpoint" | ||
| if not os.path.exists(cls.checkpoint_dir): |
There was a problem hiding this comment.
The checkpoint directory path is hardcoded to a local path (/Users/jfacevedo/...). Although the test skips gracefully if the path is not found, it is better to make this path configurable via an environment variable or a command-line argument so other developers or CI/CD can run this test with their own checkpoints.
…pydantic schema - Register DecoderBlockType.KIMI_K3 in common_types.py - Add Kimi K3 & KDA specific fields to ModelArchitecture in types.py - Exempt KIMI_K3 from base_mlp_dim == base_moe_mlp_dim validation check in types.py - Create kimi-k3.yml (full 93-layer hybrid KDA/MLA + 896-expert MoE config) - Create kimi-k3-tiny.yml (4-layer tiny config for local testing) - Add KIMI_K3_CONFIGS test suite to configs_test.py - Decouple optional Google-internal/uninstalled dependencies across maxtext
…ests - Implement situ_and_mul in linears.py with beta (4.0) and linear_beta (25.0) parameters - Register 'situ' activation in _convert_to_activation_function in linears.py - Add situ_activation_test.py with 7 test cases verifying numerical parity against PyTorch Kimi-K3 reference across float32 and bfloat16
…unit tests - Implement kda_recurrent_kernel in JAX using jax.lax.scan with log-space decay - Implement ShortConv1D in JAX using jax.lax.conv_general_dilated (depthwise 1D conv with silu) - Implement KimiDecoupledAttention NNX module in kda.py with Q/K/V projections, 1D convs, L2-normalization, gate/beta projections, A_log/dt_bias parameters, and FusedRMSNormGated - Add kda_test.py with 6 test cases verifying ShortConv1D causality, kda_recurrent_kernel parity against fla naive_recurrent_kda across sequence lengths T=1..128, and KimiDecoupledAttention NNX module forward pass
… tests - Add mla_use_output_gate support in MLA.__init__ to initialize g_a_proj (emb_dim -> head_dim), g_b_proj (head_dim -> (num_query_heads, v_head_dim)), and o_norm (RMSNorm) - Add mla_use_output_gate forward pass in MLA.__call__ applying RMSNorm(out) * sigmoid(g) before out_projection - Add mla_output_gate_test.py verifying MLA initialization and forward pass with mla_use_output_gate=True
…nit tests - Update mlp_activations in kimi-k3.yml and kimi-k3-tiny.yml to ["situ", "linear_beta_tanh"] to cleanly support gated activations in RoutedMoE - Add situ and linear_beta_tanh single-input activations in linears.py - Add latent_moe_use_norm support in RoutedAndSharedMoE in moe.py to apply RMSNorm to routed experts when enabled - Add pure JAX jax.lax.scan fallback for gmm in megablox/backend.py for CPU/macOS execution without qwix.pallas - Add qpl, tokamax, and drjax None checks for optional Google-internal dependencies - Add kimi_moe_test.py verifying RoutedAndSharedMoE initialization and forward pass with 896-expert MoE config
…-to-end unit tests - Create KimiDecoderLayer in src/maxtext/layers/kimi_decoder_layer.py to dynamically select between KDA and MLA attention layers per layer_idx and pair with RoutedAndSharedMoE - Create KimiLinearModel in src/maxtext/models/kimi_linear.py to assemble the full Kimi K3 text-only backbone in NNX (Embed -> nnx.List[KimiDecoderLayer] -> RMSNorm -> DenseGeneral) - Add kimi_linear_model_test.py with 3 test cases verifying KimiDecoderLayer, KimiLinearModel end-to-end forward pass, and initial KDA state handling
…n, and loading verification
…ix NNXDecoder _apply_embedding for ToLinen
…or clean unit testing
…orous layer-by-layer logit parity unit tests (KL < 1e-4)
a8d79f0 to
2b10f89
Compare
…e cache, stateless MXFP4 dequantization, import guards, and path cleanup
|
🤖 Hi @entrpn, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
📋 Review Summary
This Pull Request introduces comprehensive support for Kimi K3, a hybrid architecture featuring Decoupled Attention (KDA) and Multi-Head Latent Attention (MLA), along with dynamic activations (Situ / Linear Beta Tanh) and Routed/Shared MoE. The overall quality of the implementation is very high, and the inclusion of extensive unit and JAX-PyTorch parity tests is exemplary. However, there are a few critical correctness and flexibility issues that must be addressed to support standard configurations (like GQA) and prevent state-loss during sequential generation.
🔍 General Feedback
- Exemplary Test Coverage: The addition of detailed parity and functional tests (e.g.,
kimi_k3_logit_parity_test.py,kimi_linear_model_test.py,kimi_moe_test.py, andsitu_activation_test.py) is a major highlight, ensuring mathematical correctness and preventing regressions. - Stateless/Stateful Design: The Kimi Decoder layer and Kimi Linear Model backbones are beautifully modular. Standardizing state keys and names will make integration into production pipelines (like MaxEngine and vLLM) even cleaner.
- Dynamic Dimension Routing: Avoiding hardcoded constants (like
7168) in the checkpoint conversion pipelines will make the codebase more resilient to scaled-down configurations or customized Kimi-based experiments.
| B, T, H, K = q.shape | ||
| HV, V = v.shape[2], v.shape[3] | ||
| G = HV // H | ||
| if scale is None: | ||
| scale = K**-0.5 | ||
|
|
||
| # Repeat interleave q, k to HV if HV != H | ||
| if G > 1: | ||
| q = jnp.repeat(q, G, axis=2) | ||
| k = jnp.repeat(k, G, axis=2) |
There was a problem hiding this comment.
🔴 The KDA recurrent kernel will crash with a shape mismatch error if GQA/MQA is configured (where query heads outnumber key-value heads).
| B, T, H, K = q.shape | |
| HV, V = v.shape[2], v.shape[3] | |
| G = HV // H | |
| if scale is None: | |
| scale = K**-0.5 | |
| # Repeat interleave q, k to HV if HV != H | |
| if G > 1: | |
| q = jnp.repeat(q, G, axis=2) | |
| k = jnp.repeat(k, G, axis=2) | |
| B, T, H, K = q.shape | |
| HV, V = v.shape[2], v.shape[3] | |
| if scale is None: | |
| scale = K**-0.5 | |
| # Support standard Grouped Query Attention (H > HV) and other configurations (HV > H) | |
| if H > HV: | |
| G = H // HV | |
| k = jnp.repeat(k, G, axis=2) | |
| v = jnp.repeat(v, G, axis=2) | |
| g = jnp.repeat(g, G, axis=2) | |
| beta = jnp.repeat(beta, G, axis=2) | |
| HV = H | |
| elif HV > H: | |
| G = HV // H | |
| q = jnp.repeat(q, G, axis=2) | |
| k = jnp.repeat(k, G, axis=2) |
| w_transposed = np.transpose(w_dequant, (1, 0)) | ||
|
|
||
| w_padded = np.pad(w_transposed, ((0, 0), (0, 7168 - 3584)), mode="constant") | ||
| return w_padded.astype(ml_dtypes.bfloat16) |
There was a problem hiding this comment.
🟡 Hardcoding the expert MLP sizes (7168 and 3584) in top-level dequantization hooks breaks weight conversion for any non-standard model sizes. We can compute the padding dynamically based on target_shape.
| return w_padded.astype(ml_dtypes.bfloat16) | |
| if target_shape: | |
| pad_0 = max(0, target_shape[0] - w_transposed.shape[0]) | |
| pad_1 = max(0, target_shape[1] - w_transposed.shape[1]) | |
| w_padded = np.pad(w_transposed, ((0, pad_0), (0, pad_1)), mode="constant") | |
| else: | |
| w_padded = np.pad(w_transposed, ((0, 0), (0, 7168 - 3584)), mode="constant") |
| if self.is_kda: | ||
| attn_out, kda_state = self.self_attention( | ||
| normed_inputs, | ||
| initial_state=initial_kda_state, | ||
| ) |
There was a problem hiding this comment.
🔴 The initial_kda_state parameter defaults to None, but NNXDecoder always passes the recurrent/KV cache state as kv_cache. Since kv_cache is ignored here, the state is completely lost during generation, causing the KDA layers to use a zero recurrent state and zero convolution history at every step of generation!
| if self.is_kda: | |
| attn_out, kda_state = self.self_attention( | |
| normed_inputs, | |
| initial_state=initial_kda_state, | |
| ) | |
| if self.is_kda: | |
| attn_out, kda_state = self.self_attention( | |
| normed_inputs, | |
| initial_state=initial_kda_state if initial_kda_state is not None else kv_cache, | |
| ) |
|
|
||
| w_dequant = w_fp * scales | ||
| w_transposed = np.transpose(w_dequant, (1, 0)) | ||
|
|
There was a problem hiding this comment.
🟡 Hardcoding the expert MLP sizes (7168 and 3584) in top-level dequantization hooks breaks weight conversion for any non-standard model sizes. We can compute the padding dynamically based on target_shape.
| if target_shape: | |
| pad_0 = max(0, target_shape[0] - w_transposed.shape[0]) | |
| pad_1 = max(0, target_shape[1] - w_transposed.shape[1]) | |
| w_padded = np.pad(w_transposed, ((0, pad_0), (0, pad_1)), mode="constant") | |
| else: | |
| w_padded = np.pad(w_transposed, ((0, 7168 - 3584), (0, 0)), mode="constant") |
|
|
||
| def routed_expert_norm_hook(x, target_shape=None): | ||
| if hasattr(x, "shape") and len(x.shape) == 1 and x.shape[0] < 7168: | ||
| return np.pad(x, (0, 7168 - x.shape[0]), mode="constant", constant_values=1.0).astype(np.float32) |
There was a problem hiding this comment.
🟡 Hardcoding the model's hidden dimension (7168) in checkpoint conversion hooks prevents converting/loading checkpoints for custom or scaled-down configs (like the tiny test config, where emb_dim is 256). Let's calculate the dimensions dynamically.
| def routed_expert_norm_hook(x, target_shape=None): | |
| if hasattr(x, "shape") and len(x.shape) == 1 and x.shape[0] < 7168: | |
| return np.pad(x, (0, 7168 - x.shape[0]), mode="constant", constant_values=1.0).astype(np.float32) | |
| def routed_expert_norm_hook(x, target_shape=None): | |
| target_dim = target_shape[0] if target_shape else maxtext_config.emb_dim | |
| if hasattr(x, "shape") and len(x.shape) == 1 and x.shape[0] < target_dim: | |
| return np.pad(x, (0, target_dim - x.shape[0]), mode="constant", constant_values=1.0).astype(np.float32) | |
| return x |
2 participants
Description
Start with a short description of what the PR does and how this is a change from
the past.
The rest of the description includes relevant details and context, examples:
If the change fixes a bug or a Github issue, please include a link, e.g.,:
FIXES: b/123456
FIXES: #123456
You can also provide a comma-separated list. If you don't want to close a bug but
simply to reference it, use BUGS, e.g.:
BUGS: b/123456
Notice 1: Once all tests pass, the "pull ready" label will automatically be assigned.
This label is used for administrative purposes. Please do not add it manually.
Notice 2: For external contributions, our settings currently require an approval from a MaxText maintainer to trigger CI tests.
Tests
Please describe how you tested this change, and include any instructions and/or
commands to reproduce.
Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.