data/code-audit.json

{
  "$schema": "../schemas/code-audit.schema.json",
  "model": "moonshotai/Kimi-K3",
  "revision": "9f62e4e9fffbd0a83ddd60e1c209d828994b3569",
  "scope": "Public Hugging Face Transformers reference implementation",
  "caveat": "The audited Python code explains model semantics but is not Moonshot AI's complete training or production serving stack.",
  "findings": [
    {
      "id": "dual-kda-execution-regimes",
      "kind": "strength",
      "severity": "informational",
      "component": "Kimi Delta Attention",
      "evidence": {
        "file": "modeling_kimi_linear.py",
        "symbol": "KimiDeltaAttention.forward",
        "lineStart": 543,
        "lineEnd": 663
      },
      "observation": "The reference path selects chunked KDA for training and prefill, then fused recurrent KDA for one-token cached decoding.",
      "impact": "The code exposes the intended algorithmic split between parallel sequence processing and low-latency recurrent decoding.",
      "recommendation": "Keep both paths under differential logit and recurrent-state parity tests.",
      "confidence": "high"
    },
    {
      "id": "repeated-mla-cache-concatenation",
      "kind": "limitation",
      "severity": "high",
      "component": "MLA KV cache",
      "evidence": {
        "file": "modeling_kimi_linear.py",
        "symbol": "KimiDynamicCache.update",
        "lineStart": 157,
        "lineEnd": 173
      },
      "observation": "Each update appends keys and values with torch.cat.",
      "impact": "Long decoding repeatedly reallocates and copies prior cache content, creating avoidable memory traffic.",
      "recommendation": "Use indexed writes into preallocated or paged storage with copy-on-write prefix sharing.",
      "confidence": "high"
    },
    {
      "id": "single-rank-expert-layout",
      "kind": "limitation",
      "severity": "critical",
      "component": "Mixture of Experts",
      "evidence": {
        "file": "modeling_kimi_linear.py",
        "symbol": "KimiSparseMoeBlock.__init__",
        "lineStart": 768,
        "lineEnd": 813
      },
      "observation": "Expert parallel size is fixed to one and every expert is instantiated locally.",
      "impact": "The reference module cannot reproduce a practical 2.8T distributed deployment.",
      "recommendation": "Integrate a production expert-parallel runtime such as MoonEP or an equivalent GPU-only dispatcher.",
      "confidence": "high"
    },
    {
      "id": "moe-device-host-synchronization",
      "kind": "limitation",
      "severity": "critical",
      "component": "Mixture of Experts",
      "evidence": {
        "file": "modeling_kimi_linear.py",
        "symbol": "KimiSparseMoeBlock.moe_infer",
        "lineStart": 840,
        "lineEnd": 874
      },
      "observation": "Per-expert token counts move to CPU and NumPy before a Python loop invokes experts serially.",
      "impact": "Device synchronization, Python overhead, and serial launches prevent efficient use of hundreds of experts.",
      "recommendation": "Keep routing metadata on GPU and fuse permutation, grouped GEMM, weighting, and combination.",
      "confidence": "high"
    },
    {
      "id": "public-moe-training-disabled",
      "kind": "limitation",
      "severity": "critical",
      "component": "Training reproducibility",
      "evidence": {
        "file": "modeling_kimi_linear.py",
        "symbol": "KimiSparseMoeBlock.forward",
        "lineStart": 815,
        "lineEnd": 838
      },
      "observation": "The sparse MoE block raises NotImplementedError in training mode, while the router also asserts inference mode.",
      "impact": "Quantile Balancing, distributed MoE training, and the published training recipe cannot be reproduced from this module.",
      "recommendation": "Publish a tiny trainable fixture and the distributed routing, balancing, and gradient paths.",
      "confidence": "high"
    },
    {
      "id": "attnres-materialization",
      "kind": "limitation",
      "severity": "medium",
      "component": "Attention Residuals",
      "evidence": {
        "file": "modeling_kimi_linear.py",
        "symbol": "_apply_attn_res",
        "lineStart": 1075,
        "lineEnd": 1088
      },
      "observation": "Candidate residuals are concatenated, converted to FP32, normalized, scored, and mixed through separate PyTorch operations.",
      "impact": "Temporary allocation and memory traffic dominate a computation with low arithmetic intensity.",
      "recommendation": "Use online softmax and fuse normalization, scoring, mixing, partial-sum update, and adjacent RMSNorm.",
      "confidence": "high"
    },
    {
      "id": "situ-unfused",
      "kind": "limitation",
      "severity": "medium",
      "component": "SiTU-GLU",
      "evidence": {
        "file": "modeling_kimi_linear.py",
        "symbol": "SituAndMul.forward",
        "lineStart": 75,
        "lineEnd": 82
      },
      "observation": "Gate and up branches are promoted to FP32 and processed through separate tanh, sigmoid, and multiply operations.",
      "impact": "The stable formulation is preserved, but temporary tensors and repeated memory access reduce kernel efficiency.",
      "recommendation": "Fuse dequantization, both capped branches, multiplication, and output conversion.",
      "confidence": "high"
    },
    {
      "id": "multimodal-host-decisions",
      "kind": "limitation",
      "severity": "medium",
      "component": "Image-token merge",
      "evidence": {
        "file": "modeling_kimi_k3.py",
        "symbol": "KimiK3ForConditionalGeneration._merge_input_ids_with_image_features",
        "lineStart": 958,
        "lineEnd": 1090
      },
      "observation": "The merge path uses Python feature-length lists and a host-visible item operation to size the expanded sequence.",
      "impact": "Request preprocessing can synchronize with the device and is difficult to capture inside a stable CUDA Graph.",
      "recommendation": "Move ragged offset calculation into a dedicated GPU pass and bucket common visual shapes.",
      "confidence": "high"
    },
    {
      "id": "thin-reproducibility-scaffolding",
      "kind": "limitation",
      "severity": "high",
      "component": "Repository engineering",
      "evidence": {
        "file": "model repository",
        "symbol": "snapshot file inventory",
        "lineStart": 1,
        "lineEnd": 1
      },
      "observation": "The model snapshot provides custom code but no pinned environment, CI workflow, tests, or tiny deterministic fixture.",
      "impact": "Dependency drift and kernel substitutions are hard to validate independently.",
      "recommendation": "Publish a version matrix, golden tokenizer fixtures, tiny random weights, and cached-versus-uncached parity tests.",
      "confidence": "high"
    },
    {
      "id": "remote-code-integration",
      "kind": "limitation",
      "severity": "high",
      "component": "Deployment security",
      "evidence": {
        "file": "config.json",
        "symbol": "auto_map",
        "lineStart": 1,
        "lineEnd": 15
      },
      "observation": "Transformers loads custom model classes through trust_remote_code.",
      "impact": "Unpinned revisions can change executable Python code and dependency behavior.",
      "recommendation": "Pin the exact model revision, review the code, lock dependencies, and run it in an isolated environment.",
      "confidence": "high"
    }
  ]
}