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:

```python
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:

| batch | cutoff_len (600-token samples) |
|---|---|
| 2 | 1536 |
| 4 | 3072 |
| 8 | 6144 |

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`.

Start with kt_kernel/sft/wrapper.py around chunked_prefill_size and trace get_kt_config_dict into model_args.py and llamafactory/hparams/parser.py. Reproduce the reported batch-size failure, then inspect kt_kernel/sft/layer.py and base.py to confirm qlen scaling. Done means supported per-device batch sizes no longer exceed the buffer unexpectedly and the relevant configuration behavior is covered.

Written by the indexing model from the issue text.