[fix](kt-sft): scale chunked_prefill_size by the per-device batch size

[devangpratap]

What does this PR do?

Fixes #2150

chunked_prefill_size is sized from kt_model_max_length, which LLaMA-Factory populates from cutoff_len — a per-sequence bound. The rank-0 buffer already scales by distributed_world_size, but the trainer flattens a batch before the kernel ever sees it (qlen = batch_size * seq_len, sft/layer.py), and batch is precisely what grows qlen on each rank. So the buffer is undersized by exactly the batch factor and _validate_forward_inputs raises for any per_device_train_batch_size > 1:

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

The chain on current main:

sft/wrapper.py:774-780   rank0_chunked_prefill_size = kt_model_max_length * distributed_world_size
sft/layer.py             qlen = batch_size * seq_len
sft/base.py              if qlen > self.chunked_prefill_size: raise

This adds a kt_train_batch_size field to KTConfig and multiplies the rank-0 buffer by it, alongside the existing world-size factor. It follows the existing config idiom, so it is settable through model_args or ACCELERATE_KT_TRAIN_BATCH_SIZE, and defaults to 1 — current behaviour is unchanged.

As the issue notes, ACCELERATE_KT_MODEL_MAX_LENGTH is not a usable lever: LLaMA-Factory overwrites it from cutoff_len before kt reads it, so config.py's if ... is None fallback never fires. ACCELERATE_KT_TRAIN_BATCH_SIZE is a separate name and is not overwritten.

The AMX consequence is the substantive part. At batch 1, with 256 experts and num_experts_per_tok: 8, each expert's GEMM gets roughly tokens * 8 / 256 ≈ 6 rows against _AMX_M_STEP = 32, so the tile runs mostly padded. Batch 8 brings that to ~48 and fills it.

What sets kt_train_batch_size

Worth being upfront: nothing in this repo populates it automatically, so as of this PR it is a lever the user sets, not an auto-derived value. That is deliberate — per_device_train_batch_size lives on training_args, and neither training_args nor get_kt_config_dict exists anywhere in kt-kernel (_build_kt_plugin_from_args only receives model_args and finetuning_args), so auto-deriving it is a LLaMA-Factory-side change, not one I can make here.

What this PR does give you:

• a working manual lever, which today does not exist at all, since ACCELERATE_KT_MODEL_MAX_LENGTH is overwritten from cutoff_len before kt reads it and the only remaining knob is cutoff_len itself
• the hook for the automatic version. LLaMA-Factory already does os.environ["ACCELERATE_KT_MODEL_MAX_LENGTH"] = str(cutoff_len); setting ACCELERATE_KT_TRAIN_BATCH_SIZE next to it is one line and needs nothing further here

Happy to send that LLaMA-Factory patch too if you point me at the fork you build against.

Memory

The buffer grows by the batch factor, so it is worth a sanity check. At cutoff_len 1024 and batch 8 the rank-0 buffer goes to 8192, still comfortably under the chunked_prefill_size=25600 used in this repo's own SFT example in python/experts.py. Default stays 1, so nothing changes for anyone who does not opt in.

Scope

Kept to the sizing bug. The issue also notes that a YAML model_max_length: is clobbered at llamafactory/hparams/parser.py:609 — that is in LLaMA-Factory, not this repo, so it is out of scope here.

Testing

Added kt-kernel/test/test_sft_chunked_prefill_size.py, covering the default (batch 1, unchanged), the exact case from the issue (cutoff_len 1024 at batch 2, 4 and 8), composition with distributed_world_size, the env-var override, and the max_position_embeddings fallback.

It loads sft/config.py under a synthetic parent package rather than importing kt_kernel, whose __init__ loads the compiled extension — so it runs without a build and without AMX hardware:

$ python3 kt-kernel/test/test_sft_chunked_prefill_size.py
All chunked_prefill_size tests passed.

Verified it is a real regression test: with the two source files reverted it exits 1, and exits 0 once they are restored.

I do not have a Sapphire Rapids machine, so I have not run an end-to-end fine-tune at batch > 1. The reporter's traceback plus the three files above are the evidence for the sizing itself.

Before submitting

• Did you read the contributor guideline?
• Did you write any new necessary tests?

[devangpratap]

@yyj6666667 could I get the run-ci label on this one? check-changes exits 1 without it and pr-test-kt-kernel-finish follows from that. The test and deploy jobs are green on all three platforms.

Tagging you rather than @jdai0 since you picked up #2150.

The change is small: chunked_prefill_size is sized from kt_model_max_length (per-sequence), while each rank's qlen is batch_size * seq_len, so the rank-0 buffer needs the batch factor alongside the distributed_world_size one it already has. Defaults to 1, so nothing changes unless a batch size is set.

The new test runs without a build or AMX hardware, it loads sft/config.py under a synthetic parent package instead of importing kt_kernel.