test(kimi-k3): guards for the unreachable export and inference paths, plus the tests upstream lacks - #36
Open
XiaohanZhangCMU wants to merge 6 commits into
Open
Conversation
XiaohanZhangCMU
force-pushed
the
xiaohan/kimi-k3-bridge
branch
from
August 19, 2026 21:02
229f223 to
b2dde36
Compare
XiaohanZhangCMU
force-pushed
the
k3/k3-harness-and-guards
branch
from
August 19, 2026 21:05
6b2c51c to
57c6207
Compare
XiaohanZhangCMU
changed the title
XiaohanZhangCMU
force-pushed
the
xiaohan/kimi-k3-bridge
branch
from
August 19, 2026 22:59
881b97b to
b06a057
Compare
XiaohanZhangCMU
force-pushed
the
k3/k3-harness-and-guards
branch
from
August 19, 2026 22:59
57c6207 to
8c7a222
Compare
XiaohanZhangCMU
force-pushed
the
xiaohan/kimi-k3-bridge
branch
from
August 19, 2026 23:24
b06a057 to
8d24365
Compare
XiaohanZhangCMU
force-pushed
the
k3/k3-harness-and-guards
branch
from
August 19, 2026 23:24
8c7a222 to
3fc8bfc
Compare
XiaohanZhangCMU
force-pushed
the
xiaohan/kimi-k3-bridge
branch
from
August 20, 2026 02:52
8d24365 to
b06a057
Compare
XiaohanZhangCMU
force-pushed
the
k3/k3-harness-and-guards
branch
from
August 20, 2026 02:52
3fc8bfc to
8c7a222
Compare
XiaohanZhangCMU
force-pushed
the
xiaohan/kimi-k3-bridge
branch
from
August 20, 2026 05:41
b06a057 to
6071541
Compare
bank_num_rows and pack/unpack are pure-tensor and CPU-testable, and they are what breaks silently under pipeline parallelism: a wrong row count still produces a payload of the right width, so the snapshot bank shifts with nothing raised. Signed-off-by: Xiaohan Zhang <xiaohan.zhang@baseten.co>
Two Kimi-K3 paths were unsupported but did not say so. Config-only export (F4). build_conversion_tasks dereferenced hf_pretrained.state.source unconditionally, so the repository's standard examples/conversion/convert_checkpoints.py export path -- which passes a PretrainedConfig with no .state -- died with an opaque AttributeError. K3 export is inherently source-backed: the virtual dense expert keys are derived from the checkpoint's MXFP4 .weight_packed/.weight_scale pairs. Now rejected up front with an error that names the reason, not just the symptom. Cached incremental inference (F7). The attention forward accepted inference_context and then del'd it. MLA never populates the KV cache and KDA never carries recurrent or convolution state, so MCore's static and dynamic cached-inference engines silently returned wrong tokens after prefill -- no exception, just a model that forgot its prefix. Now raises NotImplementedError. The documented greedy script recomputes the whole growing prefix and passes no context, so it is unaffected. Both guards run before their function touches self, so the tests drive them through the unbound functions with lightweight stand-ins -- no model construction, no CUDA, no checkpoint. The cached-inference suite also asserts the None case does NOT trip the guard, so the fix cannot regress into blocking the supported prefix-recompute path. Signed-off-by: Xiaohan Zhang <xiaohan.zhang@baseten.co>
…nfig truncation Upstream ships create_hf_toy_model.py with no test. This exercises the two helpers a K3 toy model depends on -- _select_hub_files and _truncate_config -- against upstream's own implementation. The 523-line generator rewrite and the K3 bridge tests that used to ride along in this commit are dropped: the new base carries upstream's generator, and upstream's test_kimi_k3_bridge.py is byte-identical to the copy this PR was carrying.
Two fixes to the vendored page. Conversion: it invoked `./scripts/conversion/convert.sh`, which does not exist here. The fork's entrypoint is `examples/conversion/convert_checkpoints_multi_gpu.py`, which is also the one that takes the `--tp/--pp/--ep/--etp` flags the example passes. The toy-model link now points at the copy in this repo rather than NVIDIA's main. Limitations: four behaviours a reader would otherwise have to find by running into them. - Full-weight training with TP > 1 is unsupported. K3 marks replicated norms and AttnRes projections `sum_gradients_across_tp_domain` and Megatron-LM has no consumer for that marker, so the replicas drift. Only the frozen-base LoRA path, which does not train them, is supported. - Plain-dtype export of the routed experts is unverified. - Config-only export and cached incremental inference now raise instead of producing wrong output; the page says so. Signed-off-by: Xiaohan Zhang <xiaohan.zhang@baseten.co>
…yout `KimiK3DecoderLayer.__init__` builds the set of every stage's first layer by calling `get_transformer_layer_offset` once per pipeline rank. Under a custom `pipeline_model_parallel_layout` that call used to answer for the local rank whatever rank it was given, so the set collapsed to one element and the sending and receiving stages disagreed about how many AttnRes bank rows the payload carries. Megatron-LM#31 forwards the rank. This pins the K3 side of it: two custom layouts, one even and one uneven, asserting one distinct start per stage and the bank width that follows from each boundary. Signed-off-by: Xiaohan Zhang <xiaohan.zhang@baseten.co>
The page pointed readers at a machine-readable verification card three times without saying it is NVIDIA's and not vendored here, so the guidance read as in-repo and non-actionable. Signed-off-by: Xiaohan Zhang <xiaohan.zhang@baseten.co>
XiaohanZhangCMU
force-pushed
the
k3/k3-harness-and-guards
branch
from
August 20, 2026 05:42
8c7a222 to
2117df9
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The Kimi-K3 pieces that are genuinely this fork's: two guards on paths upstream leaves unguarded, and tests upstream does not have. Stacked on #27, which is stacked on #34's base bump.
What changed since the last review
This PR was 13 files. Rebasing onto #34's new base - NVIDIA main
e89eef9dinstead of a 2026-07-09 snapshot - removed more than half of it, because upstream now ships what those commits were backporting:docs(kimi): vendor upstream's K3 model documentationdocs/models/kimi/{index,kimi-k3}.mdfix(conversion): backportHFWeightTuple.iter_finalizediter_finalizednowdocs(conversion): explain the E8M0 exponent decodecreate_hf_toy_model.pyrewritetest_kimi_k3_bridge.pyWhat is left
Two guards, on paths upstream does not guard (
fix(kimi): fail loudly on config-only export and cached inference). Neither is reachable from production training - the trainer does no MCore generation and never takes the config-only export path - so these turn two silent-wrong-answer paths into loud failures for anyone who does reach them:Tests upstream has none of:
test_kimi_k3_guards.py- pins both guards abovetest_kimi_k3_pipeline.py- AttnRes stage boundaries under a custom pipeline layout, which is the arrangement our trainer usestest_create_hf_toy_model.py-_select_hub_filesand_truncate_configin upstream's generator, verified to still exist on this baseOne doc delta on
docs/models/kimi/kimi-k3.md: points at this fork's tooling and records the limits we measured, and says plainly that the model verification card lives upstream.Testing
Not run - the suite needs a GPU. Static only: no conflict markers, every changed file parses, and the symbols the tests import resolve against the new base. These run in the validation matrix on #34.