[GLM-5.3] Fuse unquantized KDA projections with attention TP

[HanHan009527]

Motivation

GLM KDA projection fusion is disabled when a checkpoint has a quantization config or when attention TP differs from global TP. The GLM-5.3-Flash native FP8 checkpoint keeps the KDA attention projections in BF16, so the global quantization gate excludes otherwise eligible projections. DP attention also needs the fused projections to use the attention group's rank and size.

This draft adapts the per-prefix eligibility design from linkedlist771's #38254 and extends it to attention-TP sharding. The changes overlap directly: #38254 addresses precision eligibility while retaining the TP-equality gate. If it lands first, this branch should be rebased and the overlapping eligibility changes reconciled, retaining the additional attention-TP support and applicable safeguards/tests.

Modifications

• Check the original qkv/b/f_a/f_b/g_a/g_b projection prefixes before selecting fusion. Mixed-precision eligibility is limited to the exact native Fp8Config; other nonempty configs retain the original projections because a generic LinearBase probe cannot represent class-sensitive matching.
• Construct eligible fused projections without quantization and pass the attention TP rank and size to both fused Linear modules. Keep f_a/g_a replicated and the other outputs sharded by attention TP.
• Preserve the original modules when LoRA is enabled or LoRA paths are supplied.
• Add projection weight-loading, eager/CUDA Graph, topology/rank metadata, eligibility and LoRA regression coverage. Shared Linear implementations are unchanged.

Accuracy Tests

Measured on clean wheels from baseline 12771786f23190b1845db33366eba09cb5eacf41 and candidate a5fe51d52e3b4cf50b0890fb7215625336829a12, with installed source hashes checked. These results compare those fixed commits, not the current moving main branch.

• 18 CUDA projection tests passed in 8.37 s. The topology/rank unit cases simulate parallel metadata; real distributed coverage comes from the separate full-model PD run.
• Real checkpoint loading and short/65,536-token PD smoke requests passed. All 34 KDA layers in the actual checkpoint configuration satisfy the candidate eligibility predicate.
• Fixed short GSM8K: 63/64 for both.
• Full GSM8K: 1258/1319 → 1266/1319, 14 paired wins and 6 losses. The conservative paired 95% interval is [-0.544, +1.744] percentage points, passing the predefined 1 pp noninferiority margin. All full-set responses finished without error or truncation.
• Exact-65,536-token GSM8K: 62/64 for both, but with two paired wins and two new losses. Manual review confirms new content errors: one answer substitutes 15 for a stated 25; another omits a 4-dollar tip. Candidate responses all stopped normally; one baseline response hit the unchanged 2,048-token output cap. Equal totals do not resolve these new losses, and a single matched run does not establish a deterministic fusion-caused regression.

This remains a draft: long-context quality differences and the transport error below are unresolved.

Speed Tests and Profiling

Full-model A/B on 2 × 8 H20, P TP8/EP8/PP1 and D TP8/EP8/DP8/PP1. Both roles use Triton / no MoE A2A, Mooncake PD, TileLang BF16 KV, and no MTP or MegaMoE. Prefill CUDA Graph is disabled; decode uses full CUDA Graph with global maximum batch size 128. Both arms use identical launch commands and dependencies.

Every formal request has 65,536 input and 1,536 output tokens. Cold points flush cache and use distinct inputs; warm points use a shared prefilled prefix. Request-set hashes, actual token counts, SSE completion/usage, native cache counters, all P TP/D DP rank coverage, and at least 8 GiB free memory per GPU were checked.

Point	Baseline	Candidate	Change
Cold C1, N5: median TTFT	4,873.50 ms	4,823.95 ms	-1.02%
Cold C32, N32: output throughput	272.73 tok/s	274.69 tok/s	+0.72%
Warm C32, N128: output throughput	1,527.11 tok/s	1,671.56 tok/s	+9.46%
Warm C128, N512: output throughput	4,276.97 tok/s	4,383.95 tok/s	+2.50%

The candidate C128 result is one unchanged replay after the first attempt completed 511/512 requests with one router-to-prefill HTTP 502. The replay completed 512/512; the original failure is retained and its connection-error cause remains unknown. These are observed gains for the complete patch in this configuration, not an isolated attribution to attention-TP support or a universal/statistically established speedup. No deep kernel profiling was performed.

Checklist

• Format code with pre-commit; targeted checks and diff checks passed.
• Add regression tests in the registered kernel test suite.
• Update user documentation if maintainers consider it necessary; no new public flag or API is introduced.
• Provide accuracy and speed benchmark results, including failures and limitations.
• Follow existing SGLang Linear, quantization and model patterns.

---

CI States

Latest PR Test (Base): ❌ Run #34954199239
Latest PR Test (Extra): ❌ Run #34954199252
Latest PR Test (AMD ROCm 10): ❌ Run #34954199225