add_axk1
Author: kmswin1
Pull request: yes
Labels: name: New model

<!-- ci-dashboard-badge:start -->
[![CI](https://transformers-ci.lor-e.huggingface.cool/badge/pr?pr=46867)](https://transformers-ci.lor-e.huggingface.cool/d/pytest-observability-by-pr/pytest-observability-branch?var-pr=46867)
<!-- ci-dashboard-badge:end -->

Add A.X K1 model architecture
What does this PR do?
This PR adds support for A.X K1, a large-scale Mixture-of-Experts (MoE) language model developed by [SK Telecom](https://huggingface.co/skt). A.X K1 contains 519B total parameters with 33B active parameters per token, enabling strong reasoning performance with practical inference efficiency.

Paper: [A.X K1 Technical Report (arXiv:2601.09200)](https://arxiv.org/abs/2601.09200)
Model: [skt/A.X-K1](https://huggingface.co/skt/A.X-K1)
License: Apache 2.0

Architecture overview
A.X K1 is a decoder-only Transformer with the following key design choices:

Sparse MoE: 61 layers (1 dense + 60 MoE), with 192 routed experts + 1 shared expert per MoE layer. 8 experts + 1 shared expert are activated per token.
Multi-Latent Attention (MLA): Compressed KV representations for efficient long-context inference.
Post-MLP RMSNorm: An additional RMSNorm applied after the MLP/MoE block in each layer, improving training stability in large-scale sparse settings.
Multi-Token Prediction (MTP): An auxiliary training objective that also enables speculative decoding at inference time.
Hybrid Reasoning (Think / Non-Think): Supports user-controllable reasoning depth within a single unified model.
Context length: 131,072 tokens
Vocab size: 163,840 (BBPE-based tokenizer optimized for EN, KO, ZH, JA, ES and code)

This PR includes

Model implementation (AXK1ForCausalLM, AXK1Model, AXK1Config)
Tokenizer integration (BBPE-based, 163,840 vocab)
Documentation and model card
Tests

Before submitting

I confirm that this is not a pure code agent PR.
This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request), Pull Request section?
Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link to it if that's the case.
Did you make sure to update the documentation with your changes? Here are the [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
Did you write any new necessary tests?

Who can review?
@ArthurZucker @Cyrilvallez — text model architecture addition

Comments:
Thanks @kmswin1, we're taking a look

> Hey! Sorry for the delay! Very nice first implem in general! We need to make sure that we follow naming conventions, i.e. snake_case for model_type, filenames etc, and CamelCase for classes though, no uppercases. Missing tests as well!

Thank you for the review! I’ll go through the comments one by one and update the PR accordingly.

Thanks again for the thorough review! I addressed the comments in the latest commit:

* Removed the unrelated `sam3_lite_text` leftovers and the empty `ARCHIVE_MAP`.
* Moved the routing logic into `AXK1TopkRouter.forward`, which now returns the selected indices and weights. The `e_score_correction_bias` logic is also handled there, and `AXK1MoE` is now a thin wrapper around the router and experts.
* Inlined `yarn_get_mscale`.
* Imported `apply_rotary_pos_emb_interleave` from `deepseek_v3`. I kept the Cohere helper out because its interleaved RoPE convention is not numerically equivalent in our setting.
* Removed the unused non-interleaved RoPE path and the `rope_interleave` flag, since all released checkpoints use the interleaved path.
* Changed `post_mlp_layernorm` to `nn.Identity` for dense layers and removed the corresponding decoder forward branch.

I also verified that this refactor preserves behavior: generation with the released checkpoint is unchanged after the update.

The only remaining open point is the `AXK1` casing. I kept it for now because the released Hub checkpoints currently ship with `model_type: "AXK1"`, but I’m happy to migrate this to lowercase/snake_case in a follow-up if that is preferred.


Hey @kmswin1! Alright, let's go with the tradeoff of model_type in snake_case, while keeping ModuleClasses as AXK1Module then!
Also, please rebase/merge the main branch to fix conflicts
Finally, don't forget to add a test folder for the model, as well as a few integrations tests on real checkpoints! An example of this can be seen in `test_modeling_llama.py` (where the fast tests `class LlamaModelTest(CausalLMModelTest, unittest.TestCase)` inherits from CausalLMModelTest, and the IntegrationTests, or slow tests, are tests on real checkpoints)

Ping me once all this is done for a new review!

Thanks again! @Cyrilvallez 

I’ve updated the PR accordingly:

* Changed the internal `model_type` / registry names to snake_case, e.g. `axk1`, while keeping the public class names as `AXK1Config`, `AXK1Model`, and `AXK1ForCausalLM`.
* Rebased on the latest `main` and resolved the conflicts.
* Added the model test folder with fast modeling tests.
* Added slow integration tests on the released checkpoint, following the structure used in `test_modeling_llama.py`.

Please let me know if there is anything else you’d like me to adjust.

**[For maintainers]** Suggested jobs to run (before merge)

run-slow: auto

<!-- ci-dashboard-recap:start -->

---

### CI recap

**Dashboard:** [View test results in Grafana](https://transformers-ci.lor-e.huggingface.cool/d/pytest-observability-by-pr/pytest-observability-branch?var-pr=46867)
**Latest run:** [29005445833](https://github.com/huggingface/transformers/actions/runs/29005445833)
**Result:** `failure` | Grafana metrics are not available yet.

> **Code quality check failed**: test jobs were skipped. Fix the code quality issues and push again to run tests.

<!-- ci-dashboard-recap:end -->