fix(dsa): align FA3 metadata with unpadded queries

[Jzz1943]

Motivation

GLM-5.3-Flash fails during breakable prefill CUDA Graph capture when using FA3 with MegaMoE:

RuntimeError: batch_size must be equal to batch_size_k

MoE synchronization can pad the DSA length metadata, while the prefill attention bridge narrows the query tensors and page table to the real token count. FA3 then receives inconsistent batch dimensions.

For example, 19 real query tokens can be passed with 20 entries in cache_seqlens and 21 entries in the cumulative sequence-length arrays.

#41311 addressed the pooled-indexer output row mismatch. This PR fixes the remaining FA3 metadata mismatch.

Modifications

Update DeepseekSparseAttnBackend._forward_fa3() to:

• Detect trailing metadata padding for the one-token-per-sequence sparse DSA layout.
• Validate that the page table matches the query count and that both cumulative-length arrays match the padded metadata count.
• Slice cache_seqlens, cu_seqlens_q, and cu_seqlens_k to the effective query count before calling FA3.

The slices are local views and do not modify the original metadata tensors. The change adds no tensor-value reads or CPU/GPU synchronization. Inputs without excess metadata remain unchanged.

Accuracy Tests

Server launch command

CUDA_VISIBLE_DEVICES=0,1,2,3 \
sglang serve \
  --model-path /path/to/GLM-5.3-Flash \
  --served-model-name llm \
  --trust-remote-code \
  --tp-size 4 \
  --ep-size 4 \
  --mem-fraction-static 0.75 \
  --max-mamba-cache-size 240 \
  --chunked-prefill-size 8192 \
  --cuda-graph-max-bs-prefill 8192 \
  --cuda-graph-backend-prefill breakable \
  --dsa-prefill-backend fa3 \
  --dsa-decode-backend fa3 \
  --kv-cache-dtype bfloat16 \
  --moe-runner-backend deep_gemm \
  --moe-a2a-backend megamoe \
  --speculative-algorithm EAGLE \
  --speculative-num-steps 5 \
  --speculative-eagle-topk 1 \
  --speculative-num-draft-tokens 6 \
  --reasoning-parser glm45 \
  --tool-call-parser glm47 \
  --host 0.0.0.0 \
  --port 30000

The same command was used before and after applying the fix.

End-to-end reproduction

Tested on 4 × NVIDIA H200 with GLM-5.3-Flash, based on upstream commit 49bafb58cd6.

Configuration:

• TP4 / EP4
• Breakable prefill CUDA Graph, chunk size and maximum capture size: 8192
• FA3 for both DSA prefill and decode
• BF16 KV cache
• DeepGEMM MoE runner with MegaMoE all-to-all
• EAGLE: 5 speculative steps, top-k 1, 6 draft tokens

Before: startup fails during prefill graph capture with the error above.

After: graph capture and server startup complete successfully with the same configuration.

GSM8K

Evaluated the patched server using:

• 100 examples, 5-shot
• 8 concurrent evaluation threads
• 2 repetitions
• Temperature 0, maximum output length 4096
• reasoning_effort=low

Run	Accuracy
1	99/100
2	99/100

The unpatched configuration cannot complete startup, so no end-to-end accuracy baseline is available for this configuration.

Isolated FA3 validation

Separate H200 kernel-level validation passed eight cases covering padded and unpadded inputs, NoPE and RoPE, eager execution, and CUDA Graph replay. Patched outputs matched the explicitly trimmed reference exactly and passed comparison against a mathematical attention reference.

These validation cases are not included as repository unit tests in this PR.

Speed Tests and Profiling

No dedicated performance benchmark or profiling comparison was performed. This is a correctness fix for a configuration that previously failed at startup; no throughput improvement is claimed.

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.

---

CI States

Latest PR Test (Base): ❌ Run #37512517398
Latest PR Test (Extra): ❌ Run #37512516796
Latest PR Test (AMD ROCm 10): ❌ Run #37512517169

[Jzz1943]

@merrymercy @Fridge003 @ispobock Could you please take a look at this PR when you have a chance? Thanks!