Pure KDA 300M on Primus — End-to-End Guide (FLA-validated)

Pure KDA 300M on Primus — End-to-End Guide (FLA-validated)#

This document is a runnable walkthrough for the 300M pure Kimi Delta Attention (KDA) pretraining recipe in Primus, validated on 8× AMD MI300X against the Flash Linear Attention (FLA) reference implementation. It covers every step from raw dataset → tokenization → training → checkpoint conversion → lm-eval benchmark.

The same recipe scales up to the 1B pure-KDA config (kda_1B_BF16-pretrain.yaml).

It mirrors gdn-guide.md and reuses the same Megatron-LM patches, dataset shim, FLA-init flow, and lm-eval wrapper pattern.

Final result#

After 4768 iterations (≈10B tokens) on FineWeb-Edu sample-10BT:

Axis

FLA reference

Primus (this branch)

Δ

Per-iteration time (steady state, iter > 200)

1493 ms

1466.8 ms

−1.8 % (Primus faster)

Throughput

175,617 tok/s/GPU

178,810 tok/s/GPU

+1.8 %

TFLOP/s/GPU

—

626.9

—

Wall time (4768 iters, 8× MI300X, healthy node)

1h 58m 39s

1h 56m 33s

−126 s

Loss @ iter 1

11.9673

11.9669

−0.00 % (bit-perfect)

Loss @ iter 4700 (final logged)

3.3388

3.3624

+0.71 %

First Primus-below-FLA crossover

—

iter 2600

—

Loss trajectories overlap from iter ~2000 onward; the only persistent gap is in the LR-warmup region (iter 50–500) and closes monotonically. See kda-fla-parity.md for the deep-dive on every patch and env var.

lm-eval-harness (FLA-paper 8-task suite)#

Random chance is 100 / num_choices — 25 % for the 4-choice tasks (arc, hellaswag, openbookqa, mmlu, race) and 50 % for the 2-choice tasks (piqa, winogrande). Any score above random shows the model has learned something; the FLA and Primus rows show how closely the two training stacks track each other on the same 10 B-token diet.

Task

Metric

Random

FLA

Primus

Δ (Primus − FLA)

arc_challenge

acc_norm

25.00

25.17

25.00

−0.17 pp

arc_easy

acc

25.00

48.78

47.94

−0.84 pp

arc_easy

acc_norm

25.00

42.76

43.39

+0.63 pp

hellaswag

acc_norm

25.00

29.16

29.18

+0.02 pp

openbookqa

acc_norm

25.00

30.40

29.00

−1.40 pp

piqa

acc_norm

50.00

60.99

60.34

−0.65 pp

winogrande

acc

50.00

51.85

52.72

+0.87 pp

mmlu (aggregate)

acc

25.00

22.88

23.12

+0.24 pp

race

acc

25.00

25.07

25.45

+0.38 pp

mean absolute Δ

0.58 pp

Every task within ±1.4 pp — well inside the ±1.5 pp tolerance set by the 0.49% mid-training loss delta. Both stacks comfortably beat random on arc_easy, hellaswag, openbookqa and piqa; mmlu/race/arc_challenge are at random-chance for both training stacks (expected for a 300 M model on only 10 B tokens — those benchmarks need 7 B+ parameters and/or trillion-token training to lift above 25 %).

Table of contents#

Overview#

The 300M pure-KDA model has:

  • 12 Kimi-Delta-Attention blocks + 12 MLP blocks → 24 Megatron “sublayers”

  • hidden_size = 1024, ffn_hidden_size = 4096

  • num_heads = num_v_heads = 8 (Q, K, V all share head count)

  • head_k_dim = 32, head_v_dim = 64 (expand_v = 2.0)

  • Short-conv kernel size 4 (depthwise, on the concatenated QKV)

  • Per-head output gate (g_a → g_b) + per-head decay gate (f_a → f_b, combined with learnable A_log and dt_bias via softplus)

  • Tied embeddings, no positional encoding (delta-rule recurrence), RMSNorm with eps = 1e-6

  • Tokenizer: meta-llama/Llama-3.2-1B (128k vocab)

  • Total parameters: 0.302 B

Training schedule (matched to FLA’s kda_300M.json):

  • 4768 iterations × 1024 global batch × 2048 seq len = 10.0 B tokens

  • AdamW (β1=0.9, β2=0.95, wd=0.01), peak LR 2e-4, cosine decay, 200-step warmup

  • BF16 training, no dropout, gradient clip 1.0

Prerequisites#

  • Hardware: 8× AMD MI300X (or compatible ROCm GPU) on a single node

  • Software: ROCm ≥ 7.0, Docker ≥ 24.0

  • Container image: rocm/primus:v26.2 (or v25.10 with the same patches)

  • HF token: HF_TOKEN set for the gated meta-llama/Llama-3.2-1B tokenizer

  • Disk: ~20 GB for the FLA-aligned tokenized dataset + ~5 GB per saved checkpoint

  • flash-linear-attention checked out at /home/<user>/flash-linear-attention (or installed via pip install -e .) — provides the FLA KDAForCausalLM class for HF conversion + lm-eval, plus the Triton kernels that the PRIMUS_FLA_* toggles route into.

Step 1: Environment#

1.1 Start the dev container#

docker run -it \
  --device /dev/dri --device /dev/kfd \
  --device=/dev/infiniband --network host --ipc host \
  --group-add video --cap-add SYS_PTRACE \
  --security-opt seccomp=unconfined --privileged \
  -v $HOME:$HOME -v $(pwd):$(pwd) -w $(pwd) --shm-size 64G --name primus_hybrid_new \
  rocm/primus:v26.2

This runs the rocm/primus:v26.2 image with /dev/dri, /dev/kfd, IB devices, --privileged, your $HOME mounted in-place, and --shm-size 64G. The container is named primus_hybrid_new.

To re-attach later:

docker exec -it primus_hybrid_new bash
cd /home/<user>/Primus

1.2 Install Python dependencies inside the container#

pip install -r requirements.txt
pip install -e /home/<user>/flash-linear-attention   # FLA model classes + Triton kernels
pip install lm-eval                                  # for benchmark evaluation

The editable FLA install removes the need to set PYTHONPATH for every later command.

Step 2: Dataset preparation#

Identical to the GDN recipe — see gdn-guide.md. KDA reuses the same FineWeb-Edu sample-10BT preprocessed Arrow shards and the same Llama-3.2-1B tokenizer.

The default 300M YAML already points at the FLA-aligned binary:

train_data_path: >
  /home/<user>/Primus/data/fla_aligned/fla_fineweb_edu_10BT_text_sentence

(adjust the user prefix in examples/megatron/configs/MI300X/kda_300M_BF16-pretrain.yaml to match your home directory).

Step 3: Megatron-LM patches (automatic — no action needed)#

KDA uses the same six patches as GDN — no KDA-specific Megatron patch is required. They’re implemented as runtime monkey-patches using Primus’s own patch system (primus/core/patches), living under [primus/backends/megatron/patches/](https://github.com/AMD-AGI/Primus/tree/main/primus/backends/megatron/patches). Each patch registers unconditionally but is gated behind a condition= on the relevant config flag, so nothing needs to be run by hand.

See gdn-guide.md §3 for the patch-by-patch breakdown.

Step 4: (Optional) Initialize from FLA weights#

For bit-perfect iter-1 loss alignment, the validated run loads FLA’s initialized but untrained KDA-300M checkpoint and then trains from there. The YAML’s load: field points at this directory:

load: /home/<user>/Primus/output/fla_init_kda_300M
finetune: true            # load weights, ignore optimizer state and iteration count
no_load_optim: true
no_load_rng: true

Generate it once with:

python tools/hybrid/convert_fla_kda_init_to_megatron.py
#   → output/fla_init_kda_300M/iter_0000000/mp_rank_00/model_optim_rng.pt

The script instantiates FLA’s KDAForCausalLM with seed=42, harvests its randomly-initialized weights, concatenates the six FLA hidden_states → X projections into Primus’s single fused in_proj, and writes a Megatron-shape checkpoint. Skip this step if you’re happy with Primus’s own random init — final loss is identical, only iter-1 drifts by ~5e-3.

Step 5: Train#

5.1 Inspect the config#

The training config lives at examples/megatron/configs/MI300X/kda_300M_BF16-pretrain.yaml. Key parameters (matched to FLA):

train_iters: 4768                 # ≈ 10B tokens at global_batch=1024, seq=2048
micro_batch_size: 128             # per-GPU
global_batch_size: 1024           # 8 GPUs × 128 = 1024
seq_length: 2048
lr: 2.0e-4
min_lr: 2.0e-5                    # min_lr_rate=0.1 → 2e-5
lr_warmup_iters: 200
lr_decay_iters: 4768
lr_decay_style: cosine
adam_beta1: 0.9
adam_beta2: 0.95
weight_decay: 0.01
clip_grad: 1.0
seed: 42
layernorm_epsilon: 1.0e-6         # MUST be explicit — TransformerConfig default 1e-5 silently overrides the model YAML
hidden_dropout: 0.0               # MUST be explicit — language_model.yaml default 0.1 leaks through
attention_dropout: 0.0
spec: ['primus.backends.megatron.core.models.hybrid.hybrid_mamba_mla_layer_specs', 'kda_hybrid_stack_spec_no_te']
use_fla_triton_kda: true
use_fla_kda_in_kernel_gate: true
use_fla_fused_norm_gated: true
use_distributed_optimizer: false  # 300M fits — ZeRO-1 adds allreduce overhead
finetune: true
load: /home/<user>/Primus/output/fla_init_kda_300M
no_load_optim: true
no_load_rng: true

The architecture-only YAML it extends from is primus/configs/models/megatron/kda_300M.yaml.

5.2 Launch#

# inside the container, in /home/<user>/Primus
./primus-cli direct --log_file primus_kda.log \
  -- train pretrain \
  --config examples/megatron/configs/MI300X/kda_300M_BF16-pretrain.yaml

AMD Triton 3.6 cannot compile the num_stages = 4 variants of FLA’s KDA intra-chunk kernels — the MLIR pass pipeline aborts inside tritonamdgpu-schedule-loops — so the first autotune sweep fails. The training process removes those candidates from the autotune space itself, at startup, via the megatron.fla.kda_safe_autotune patch (fla_kda_autotune_patches.py) — nothing on disk is modified, so the installed fla package stays pristine. It fires on ROCm for any run that resolves to the FLA KDA backend, which this one does via use_fla_triton_kda: true, and the num_warps sweep is kept so the autotuner still picks the best remaining config. Look for this line in the log:

[Patch:megatron.fla.kda_safe_autotune] chunk_kda_bwd_kernel_intra: kept 4/12 autotune configs ...

On Triton 3.7 the full sweep compiles and the autotuner picks num_stages = 2 for this model anyway, so the narrowing is free here; the patch stays until the supported toolchain floor reaches Triton 3.7.

Expected wall time on a healthy MI300X box: ~1h 56m for the full 4768 iters (about 2 min faster than FLA’s HF-Trainer reference run).

5.4 Output layout#

Checkpoints land under Primus’s work_group/user_name/exp_name template:

output/amd/root/kda_300M_BF16-pretrain/
├── checkpoints/
│   ├── iter_0001024/
│   ├── iter_0002048/
│   ├── iter_0003072/
│   ├── iter_0004096/
│   ├── iter_0004768/                  ← FINAL (~4.5 GB)
│   │   └── mp_rank_00/
│   │       └── model_optim_rng.pt
│   └── latest_checkpointed_iteration.txt  → "4768"
└── logs/
    └── pre_trainer/

save_interval: 1024 in the YAML produces 4 mid-training checkpoints plus the final one.

Step 6: Monitor and compare against FLA#

Megatron logs iteration / elapsed_ms_inst / elapsed_ms_avg / TFLOP/s/GPU / tok/s/GPU / lm loss every 100 steps. A representative tail looks like:

iteration  4700/ 4768 | elapsed time per iteration (ms): 1467.8/1466.1 |
  TFLOP/s/GPU: 626.1 | tokens per GPU (tokens/s/GPU): 178596.5 | lm loss: 3.362445E+00

To diff against FLA’s reference log (/home/<user>/checkpoints/kda_300M_10B/trainer_state.json), divide the FLA loss field by 8 (DeepSpeed reports sum-across-ranks):

iter

FLA / 8

Primus

Δ %

Notes

1

11.9673

11.9669

−0.00 %

bit-perfect

100

7.7171

9.6903

+25.6 %

warmup gap (peak)

500

4.7349

4.8390

+2.20 %

warmup closing

1000

4.0357

4.0720

+0.90 %

LR-warmup done

2000

3.6009

3.6141

+0.37 %

converged

2600

3.5056

3.5047

−0.03 %

first Primus < FLA crossover

3000

3.4356

3.4571

+0.63 %

matched

3600

3.4107

3.4075

−0.09 %

Primus slightly lower

4000

3.3831

3.3861

+0.09 %

identical

4500

3.3603

3.3694

+0.27 %

identical

4700

3.3388

3.3624

+0.71 %

identical

Final wall time on a healthy MI300X box: 6993 s vs FLA 7119 s = Primus 126 s faster.

Step 7: Convert checkpoint to HuggingFace format#

Use tools/hybrid/convert_kda_to_fla_hf.py to translate the Megatron checkpoint into FLA’s native KDAForCausalLM HF format:

python tools/hybrid/convert_kda_to_fla_hf.py \
    --checkpoint-path output/amd/root/kda_300M_BF16-pretrain/checkpoints/iter_0004768 \
    --output-dir      output/kda_300M_fla_hf \
    --config          /home/<user>/flash-linear-attention/legacy/training/configs/kda_300M.json \
    --tokenizer-src   /home/<user>/checkpoints/kda_300M_10B

What it does:

  • Reads mp_rank_00/model_optim_rng.pt and pulls the model state dict

  • For each of the 12 FLA layers, pairs the alternating Megatron sublayers:

    • KDA sublayer (even index) → FLA model.layers.<i>.attn.*

    • MLP sublayer (odd index) → FLA model.layers.<i>.mlp.*

  • Splits Primus’s fused projections into FLA’s separate ones:

    • mixer.in_proj.weight (rows = 2·qk_dim + v_dim + 2·head_v_dim + num_v_heads) → q_proj / k_proj / v_proj / f_proj.0 / g_proj.0 / b_proj

    • mlp.linear_fc1.weight (rows = 2·intermediate_size) → gate_proj / up_proj

  • Preserves A_log, dt_bias, per-head g_norm (FLA’s FusedRMSNormGated), o_proj, f_proj.1, g_proj.1, embeddings, tied lm_head, final norm

  • Copies tokenizer files from --tokenizer-src into the output dir

Output:

output/kda_300M_fla_hf/
├── config.json              # KDAConfig, architectures=["KDAForCausalLM"]
├── model.safetensors        # ~870 MB
└── tokenizer{,_config}.json + special_tokens_map.json

Step 8: Verify conversion#

Quick smoke test in the container (with FLA importable):

PYTHONPATH=/home/<user>/flash-linear-attention \
python - <<'PY'
import torch
import fla   # auto-registers "kda" with transformers.AutoConfig

from transformers import AutoModelForCausalLM, AutoTokenizer

ckpt = "output/kda_300M_fla_hf"
tok = AutoTokenizer.from_pretrained(ckpt)
model = AutoModelForCausalLM.from_pretrained(
    ckpt, trust_remote_code=True, torch_dtype=torch.bfloat16
).cuda().eval()

for prompt in [
    "The capital of France is",
    "Once upon a time, there was a small",
    "The first law of thermodynamics states that",
]:
    inp = tok(prompt, return_tensors="pt").to("cuda")
    with torch.no_grad():
        out = model.generate(**inp, max_new_tokens=40, do_sample=False)
    print("---"); print(tok.decode(out[0], skip_special_tokens=True))
PY

Expected output for a healthy 300 M-on-10 B model: grammatical but repetitive English (canonical small-undertrained-LM failure mode under greedy decoding with no repetition penalty). Knowing “capital of France” → “Paris” is the standard sanity-check pass.

If AutoConfig raises model type kda not recognized, FLA was not imported before AutoModelForCausalLM. Either prepend PYTHONPATH=/home/<user>/flash-linear-attention or run pip install -e /home/<user>/flash-linear-attention so the auto-registration in fla/models/kda/__init__.py fires on import.

Step 9: Run lm-eval-harness benchmarks#

Use tools/hybrid/eval_kda_lm_eval.py, which imports fla first (so AutoConfig recognizes the kda model type) and patches KDAForCausalLM.__init__ / KDAModel.__init__ to accept the dtype kwarg that transformers ≥ 4.55 passes internally.

Do not invoke lm_eval --model hf ... directly — AutoConfig.from_pretrained will fail with model type kda not recognized.

9.1 Evaluate the Primus checkpoint (~15–30 min on one MI300X)#

mkdir -p output/kda_300M_eval_results_primus

PYTHONPATH=/home/<user>/flash-linear-attention \
HIP_VISIBLE_DEVICES=0 \
TOKENIZERS_PARALLELISM=false \
python tools/hybrid/eval_kda_lm_eval.py \
    --model hf \
    --model_args pretrained=output/kda_300M_fla_hf,dtype=bfloat16,trust_remote_code=True,tokenizer=meta-llama/Llama-3.2-1B \
    --tasks arc_easy,arc_challenge,hellaswag,openbookqa,piqa,winogrande,mmlu,race \
    --batch_size auto \
    --output_path output/kda_300M_eval_results_primus \
    2>&1 | tee output/kda_300M_eval_results_primus/lm_eval.log

9.2 Evaluate the FLA reference checkpoint (apples-to-apples)#

mkdir -p output/kda_300M_eval_results_fla

PYTHONPATH=/home/<user>/flash-linear-attention \
HIP_VISIBLE_DEVICES=1 \
TOKENIZERS_PARALLELISM=false \
python tools/hybrid/eval_kda_lm_eval.py \
    --model hf \
    --model_args pretrained=/home/<user>/checkpoints/kda_300M_10B,dtype=bfloat16,trust_remote_code=True,tokenizer=meta-llama/Llama-3.2-1B \
    --tasks arc_easy,arc_challenge,hellaswag,openbookqa,piqa,winogrande,mmlu,race \
    --batch_size auto \
    --output_path output/kda_300M_eval_results_fla \
    2>&1 | tee output/kda_300M_eval_results_fla/lm_eval.log

9.3 Diff the two result JSONs#

python - <<'PY'
import json, glob
def load_latest(d):
    return json.load(open(sorted(glob.glob(f"{d}/**/results_*.json", recursive=True))[-1]))
fla    = load_latest("output/kda_300M_eval_results_fla")
primus = load_latest("output/kda_300M_eval_results_primus")
print(f"{'task':<18} {'FLA':>8} {'Primus':>8} {'Δ':>+8}")
for task in sorted(set(fla['results']) & set(primus['results'])):
    for k in ('acc,none', 'acc_norm,none'):
        if k in fla['results'][task] and k in primus['results'][task]:
            f, p = fla['results'][task][k], primus['results'][task][k]
            print(f"{task[:17]:<18} {f:>8.4f} {p:>8.4f} {p-f:>+8.4f}  ({k})")
PY

Measured result (validated on tw006, this branch). The Random column is 100 / num_choices for the lm-eval task — anything above it means the model learned something:

Task

Metric

Random

FLA

Primus

Δ (Primus − FLA)

arc_challenge

acc_norm

25.00

25.17

25.00

−0.17 pp

arc_easy

acc

25.00

48.78

47.94

−0.84 pp

arc_easy

acc_norm

25.00

42.76

43.39

+0.63 pp

hellaswag

acc_norm

25.00

29.16

29.18

+0.02 pp

openbookqa

acc_norm

25.00

30.40

29.00

−1.40 pp

piqa

acc_norm

50.00

60.99

60.34

−0.65 pp

winogrande

acc

50.00

51.85

52.72

+0.87 pp

mmlu (aggregate)

acc

25.00

22.88

23.12

+0.24 pp

race

acc

25.00

25.07

25.45

+0.38 pp

mean absolute Δ

0.58 pp

Every task within ±1.4 pp — consistent with the 0.49% loss delta at the end of training. mmlu / race / arc_challenge are at random-chance for both stacks (300 M params + 10 B tokens is below the threshold those benchmarks need to lift above noise).

Configs and tools used#

docs/04-technical-guides/hybrid-models/
├── kda-guide.md                                  ← this file
└── kda-fla-parity.md                              ← deep-dive on every change
examples/megatron/configs/MI300X/
└── kda_300M_BF16-pretrain.yaml        ← training config
primus/configs/models/megatron/
└── kda_300M.yaml                 ← architecture-only config
primus/backends/megatron/core/models/hybrid/
├── kimi_delta_attention.py                        ← FLA-aligned mixer (fused in_proj, FLA Triton paths)
├── kimi_delta_attention_layer.py                  ← eps propagation, optional pre-norm
└── hybrid_mamba_mla_layer_specs.py                ← kda_hybrid_stack_spec_no_te
primus/backends/megatron/patches/                  ← same 6 patches as GDN (Primus patch system, shared)
├── gdn_config_patches.py                          ← registers use_fla_triton_kda + fusion flags + hybrid init
├── mamba_fused_ce_patches.py                      ← FLA fused cross-entropy for MambaModel
├── torch_fused_adam_patches.py                    ← PRIMUS_TORCH_OPTIM opt-in
├── mlp_fla_swiglu_patches.py                      ← FLA Triton SwiGLU for MLP
├── torch_norm_fla_rmsnorm_patches.py              ← FLA RMSNorm for WrappedTorchNorm
├── fla_runtime_patches.py                         ← resolves PRIMUS_FLA_* knobs onto args
├── fla_kda_autotune_patches.py                    ← narrows the KDA Triton autotune space on ROCm (auto-applied, in-process)
└── mamba_fla_data_patches.py                      ← FLA-order dataset shim wiring
tools/hybrid/
├── convert_fla_to_megatron.py                     ← FLA Arrow → Megatron .bin/.idx (shared)
├── fla_order_dataset.py                           ← FLA-order dataset shim (shared)
├── convert_fla_kda_init_to_megatron.py            ← FLA HF init → Megatron sharded ckpt
├── convert_kda_to_fla_hf.py                       ← Megatron sharded ckpt → FLA HF
└── eval_kda_lm_eval.py                            ← lm-eval wrapper (registers KDA)

Troubleshooting#

KeyError: 'kda' at AutoModelForCausalLM.from_pretrained#

You imported transformers before fla (or didn’t import fla at all). fla/models/kda/__init__.py runs AutoConfig.register(KDAConfig.model_type, KDAConfig, exist_ok=True) on import. Either:

  • Prepend PYTHONPATH=/home/<user>/flash-linear-attention and import fla in your script BEFORE the transformers import, OR

  • pip install -e /home/<user>/flash-linear-attention once and forget about PYTHONPATH, OR

  • Use the wrapper: python tools/hybrid/eval_kda_lm_eval.py ...

Conversion: KeyError: 'decoder.layers.0.mixer.in_proj.weight'#

You trained with an older code branch that still had six separate projections. Either re-train with the current fused-in_proj branch or patch the converter to read the unfused q_proj_weight/k_proj_weight/… keys (see git history of tools/hybrid/convert_kda_to_fla_hf.py).

Iter 1 loss ~12.05 instead of ~11.97#

The layernorm_epsilon: 1.0e-6 override is being silently overwritten by the TransformerConfig default of 1e-5. Confirm it’s in the training YAML’s overrides: block (not just the model YAML).

Iter 1 loss not bit-matching FLA but converges fine#

You probably didn’t load the FLA-init checkpoint (Step 4) or didn’t set PRIMUS_FLA_DATA=1. Without either, the first batch differs (Megatron shuffler vs HF DistributedSampler) and the per-parameter nn.init.normal_ draw order differs (Megatron traverses Primus’s fused in_proj, FLA traverses 6 separate nn.Linear modules). The gap disappears by iter ~2000 even without either fix.

Loss is +0.2–0.4 above FLA across the whole run (with FLA-init loaded)#

You probably have use_fla_kda_in_kernel_gate: false or use_fla_fused_norm_gated: false. Those toggles select the bit-identical- to-old-FLA fused_kda_gate + _apply_gated_norm paths, which run the gate compute in fp32 (slightly different rounding than the in-kernel bf16 accumulator). Set both to true to match the current FLA reference.

Per-iter time ≫ 1500 ms#

Most likely you have PRIMUS_FLA_CONV=0. The Tri-Dao causal_conv1d_fn on ROCm requires [B, D, T] layout, so each iteration pays two transpose+contiguous copies of the (B, qk_dim·2 + v_dim, T) tensor — about 35 ms wasted per iter at micro_batch=128. Set PRIMUS_FLA_CONV=1 to switch to FLA’s Triton causal_conv1d (accepts [B, T, D] natively).

Out-of-memory at iter 1#

Two common culprits:

  1. PYTORCH_ALLOC_CONF=expandable_segments:True is unset — set it.

  2. q.contiguous()/k.contiguous()/v.contiguous() removed from KDA forward — the Triton kernel will allocate its own copies while autograd still pins the original views, doubling Q/K/V activation memory. Restore the explicit contiguous calls (see kimi_delta_attention.py around the chunk_kda call site).

Eval truncation warnings#

Some samples exceed the model’s max_position_embeddings = 2048. Add max_length=1024 to --model_args if it bothers you; it only meaningfully affects RACE.

See also#

Contents