[GH-ISSUE #2150] kt-sft: chunked_prefill_size sized from cutoff_len, never scaled by batch — any per_device_train_batch_size>1 hard-errors #1293

New issue

Open

opened 2026-10-29 14:50:55 +00:00 by bsakaguchi · 1 comment

bsakaguchi commented

2026-10-29 14:50:55 +00:00

Owner

Originally created by @PUMA-SPz on GitHub (Aug 8, 2026).
Original GitHub issue: https://github.com/kvcache-ai/ktransformers/issues/2150

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

Stop clobbering the env var — respect a user-supplied ACCELERATE_KT_MODEL_MAX_LENGTH instead
of overwriting it from cutoff_len.
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.

Originally created by @PUMA-SPz on GitHub (Aug 8, 2026). Original GitHub issue: https://github.com/kvcache-ai/ktransformers/issues/2150 # `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`.

bsakaguchi commented

2026-10-29 14:50:56 +00:00

Author

Owner

@yyj6666667 commented on GitHub (Aug 8, 2026):

We will soon deliver the cookbook . please wait.

@yyj6666667 commented on GitHub (Aug 8, 2026): We will soon deliver the cookbook . please wait.