kt-sft: chunked_prefill_size sized from cutoff_len, never scaled by batch — any per_device_train_batch_size>1 hard-errors

chunked_prefill_size is sized from cutoff_len and never scaled by batch size

Component: kt_kernel.sft (kt-kernel 0.6.4) + LLaMA-Factory 0.9.6.dev0 Severity: medium — silently caps kt-sft at per_device_train_batch_size: 1 unless the user happens to inflate cutoff_len, leaving most of the AMX tile width unused.

Summary

kt sizes its prefill buffer from kt_model_max_length, which LLaMA-Factory populates with cutoff_len — a per-sequence bound. The trainer flattens a batch into a single qlen, so the buffer is undersized by exactly the batch factor and any per_device_train_batch_size > 1 raises:

ValueError: qlen (1168) exceeds chunked_prefill_size (1024).
Increase chunked_prefill_size or reduce qlen to avoid buffer overrun.

with cutoff_len: 1024, per_device_train_batch_size: 2. Reproduced identically at batch 2, 4 and 8.

Chain

llamafactory/hparams/parser.py:609    model_max_length := cutoff_len
llamafactory/hparams/parser.py:623 -> model_args.py:501
                                      "kt_model_max_length": model_max_length
kt_kernel/sft/wrapper.py:376-382      chunked_prefill_size = kt_model_max_length * distributed_world_size
kt_kernel/sft/layer.py:348,535        qlen = batch_size * seq_len
kt_kernel/sft/base.py:502-506         if qlen > chunked_prefill_size: raise

wrapper.py:382 scales by distributed_world_size but not by per-device batch size — yet batch is precisely what grows qlen on each rank.

The env-var override does not work

ACCELERATE_KT_MODEL_MAX_LENGTH looks like the escape hatch and is not one:

• model_args.py:~525-527 unconditionally does os.environ["ACCELERATE_KT_MODEL_MAX_LENGTH"] = str(cutoff_len)
• ~532 then does hf_kt._kt_config.update(...)
• config.py:228-229 only consults the env var when the field is still None

So the user's value is overwritten before kt reads it. A YAML model_max_length: is likewise clobbered at parser.py:609. The only working lever is cutoff_len itself.

(Noting this explicitly because it is the natural thing to try, and it fails silently — the run proceeds with the old buffer and raises the same error, which reads as if the override did nothing rather than as if it was overwritten.)

Impact

The AMX consequence is the substantive one. With 256 experts and num_experts_per_tok: 8, a batch of 1 gives each expert's GEMM roughly

M = tokens * 8 / 256  ~=  6 rows

against _AMX_M_STEP = 32 (amx.py:46) — so the tile runs ~80% padded. The hardware is being fed 6 rows of a 32-row unit. Batch 8 brings M to ~48 and fills it.

Measured on the reporting system (Ornith-1.0-35B, 40 layers, 256 experts, 1x RTX 3090): batch 1 is 0.142 samples/s with the GPU at 9 GB of 24 GB. The idle VRAM is not the constraint; this buffer is.

Suggested fix

Scale by the per-device batch size:

rank0_chunked_prefill_size = (
    int(chunked_prefill_size)
    * distributed_world_size
    * int(getattr(cfg, "kt_train_batch_size", 1) or 1)
)

passing per_device_train_batch_size through get_kt_config_dict alongside model_max_length.

Failing that, two cheaper improvements:

1. Stop clobbering the env var — respect a user-supplied ACCELERATE_KT_MODEL_MAX_LENGTH instead of overwriting it from cutoff_len.
2. Name the cause in the error. The current text sends users to look at sequence length, which is within bounds. Something like: "qlen (1168) exceeds chunked_prefill_size (1024); this buffer is sized from cutoff_len and does not account for per_device_train_batch_size=2 — set cutoff_len >= 2048".

Workaround

Raise cutoff_len to at least batch_size * longest_padded_sample, rounded up to a multiple of 32 so _aligned_max_len (amx.py:443) is a no-op:

Sizing at batch * cutoff_len also works but over-allocates the kt host pools by roughly 50% when samples are shorter than the cap.

Environment

kt-kernel 0.6.4 · transformers-kt 5.6.0 · accelerate-kt 1.14.0 · torch 2.9.1 · LLaMA-Factory 0.9.6.dev0 FSDP2, 1x RTX 3090, CPU experts on AMX (Xeon Max 9480), kt_threadpool_count: 2.