Reuse host lengths to avoid KDA prefill synchronization

[ormandj]

Motivation

This change targets a shorter wait for the first response token, especially for long prompts. KDA prefill currently waits for a sequence-length read from the GPU in every layer of every chunk, even though the scheduler already has that length on the CPU.

Reusing the host lengths removes those repeated waits without changing the attention computation or generation settings. The profile below shows 1,122 device scalar reads eliminated for a 131,158-token prompt. That establishes reduced synchronization work; an isolated time-to-first-token improvement has not yet been measured.

Modifications

Use the sum of extend_seq_lens_cpu when available. Keep the device-offset fallback when host lengths are absent, and preserve input trimming, output padding, cached-state handling and target verification.

Accuracy Tests

Author CPU validation at c0a43f8c19: 24 tests passed, with CUDA hidden. Its Python runtime and test trees are unchanged at refreshed head eb71d35924 on main afe90a8bc9. Earlier GPU and serving results retain their stated source scope.

Author-run result: 24 CPU cases pass in test/registered/unit/layers/attention/test_kda_extend_host_lengths.py. They cover host lengths and the device fallback, padding and state tracking. The host-metadata cases raise if the code reads device offsets, and check exact output values, shape and padding with stubbed numerical kernels. These tests validate the metadata change, not GPU kernel numerics.

python -m pytest -q test/registered/unit/layers/attention/test_kda_extend_host_lengths.py

Speed Tests and Profiling

Author-reported CPU/CUDA profiling of GLM-5.3-Flash W4A16 with FP8 KV on two RTX PRO 6000 Blackwell Max-Q 96 GB GPUs at 300 W, TP2 over PCIe. One cold 131,158-token prompt was traced for each path, using 4,096-token prefill chunks and stopping after one generated token. The table counts KDA sequence-length scalar reads on rank 0 across the 33 target prefill chunks.

KDA sequence-length reads from device	Device-offset path	Host-length path
Per target prefill chunk	34	0
Full prompt	1,122	0

These are operation counts from the profiled paths. The surrounding builds also differed in KV layout, so they do not establish an isolated serving speedup for this PR. The device-offset fallback still synchronizes when host lengths are absent.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

Developed with AI assistance.

---

CI States

Latest PR Test (Base): ❌ Run #34282262308
Latest PR Test (Extra): ❌ Run #34282261977
Latest PR Test (AMD ROCm 7.2): ❌ Run #34282262293