[xsuler]

What does this PR do?

When a prompt spans multiple prefill chunks, later queries must attend to the cached prefix as well as their current chunk. Carry prefix and total-KV lengths through scheduler metadata and use paged extend attention with bottom-right causal alignment. The zero-score regression now produces [1, 2, 3, 4] across chunks rather than [1, 2, 5, 6].

• Use one paged FlashAttention varlen call for mixed prefix/chunk lengths; support native attention through per-query paged metadata.
• Preserve Ling KDA short-convolution history across chunks and isolate fresh rows in mixed batches.
• Continue prefill inside the last allocated KV page, and explicitly reject cached-prefix prefill on NPU.
• Add CPU/CUDA regressions and reproducible checkpoint/performance validation scripts. Keep the existing scheduler token budget unchanged.

Related issue

Fixes #635

Type of change

• 🐛 Bug fix

How was it tested?

Real NVIDIA H20-3e CUDA validation with PyTorch 2.8 and FlashAttention 2.7.4:

NVIDIA_TF32_OVERRIDE=0 pytest -q \
  tests/test_chunked_prefill_cpu.py tests/test_native_attention_gpu.py \
  tests/test_inference_scheduler_cpu.py tests/test_sampling_cpu.py

The initial combined run passed 62 tests. After adding three Ling KDA cases, the CUDA suite passed all 20 tests, covering:

• The exact zero-score reproduction, one-token chunks, GQA, padded values, mixed sequence/prefix lengths, page boundaries and sliding windows.
• Dense Qwen3 adapter logits, fixed-token logprobs, subsequent decode and CUDA-graph replay.
• Ling BF16 KDA attention outputs and FP32 recurrent/conv continuation state for 128-, 65- and 1-token partitions.
• CPU scheduler checks for final-chunk sampling, decode interleaving and slot lifecycle.

Ling-3.0-tiny checkpoint diagnostics also record full-prefill batch-shape sensitivity alongside chunk-partition measurements. These measurements are separate from the attention/state regressions; bitwise equality of BF16 MoE logits is not assumed.

Performance

Paired alternating measurements on the same Ling checkpoint, one ongoing 32-token prompt plus one incoming long prompt, 32 generated tokens per request, eager decode and fixed tokens. Model/scheduler timings exclude HTTP, tokenization and sampling.

Incoming prompt	Repeats	Original full prefill	Fixed full prefill	Throughput ratio
1025 tokens	5	51.25 tok/s	50.79 tok/s	99.1%
4097 tokens	3	47.29 tok/s	47.69 tok/s	100.9%

Both workloads meet the predefined unchanged-budget throughput floor of 90% of baseline. For the 4097-token workload, a 256-token budget reduced the longest ongoing-decode gap from 189 ms to 115 ms, while throughput changed from 47.69 to 20.76 tok/s and incoming-prompt TTFT from 152 ms to 1792 ms. This PR preserves the default budget and makes no aggregate throughput-improvement claim for smaller chunks.

Reproduce with scripts/validation/chunked_prefill_perf.py, passing the original backend file from e17794c as --baseline.

A Tic-Tac-Toe GSPO run used Ling-3.0-tiny, LR 1e-5, batch size 32, 8 samples per prompt, concurrency 256, TP 4, world size 8, and mini batch size 8. It was stopped at the user's request. TensorBoard recorded reward means -0.903125, -0.710938, and -0.275000 for steps 0–2; the requested 10-step limit was not reached.

Long-prompt selected-token logprobs

Tested the original BF16 Ling checkpoint with TP=4, eager decode, temperature 1, and the same 128 fixed continuation tokens for every comparison. Means exclude prompt tokens and cover selected next-token logprobs, in nats. The train comparison uses packed model.train() forward without gradients or optimizer updates, with sequence parallelism disabled and the same one-token LM-head projection shape as rollout.

Prompt tokens	Prefill budget	Full vs chunked rollout mean abs diff	Chunked rollout vs train mean abs diff
4097	2048	0.022061	0.016831
4097	512	0.016846	0.012143
4097	128	0.020449	0.016980
8193	2048	0.016886	0.014880
8193	512	0.013455	0.012511
8193	128	0.015992	0.016355

Full-rollout vs train mean abs diff: 0.015572 (4097 prompt tokens), 0.012157 (8193 prompt tokens). Full-prefill repeat controls had zero difference for both prompt lengths. Individual-token differences are also recorded: the largest full-vs-chunked value was 0.983284 nats.

Reproduce with torchrun --standalone --nproc-per-node=4 scripts/validation/ling_long_prompt_logprobs.py --model /path/to/Ling-3.0-tiny --output results.json. The output preserves token-level logprobs, token hashes, mean, median, p95 and maximum deltas.

Checklist

• The PR title summarizes the contribution.
• Linked the related issue.
• Relevant scheduler, sampling and attention tests pass.
• New behavior is covered by tests.
• Described validation commands, hardware and performance tradeoffs.
• Public API / CLI behavior is unchanged.