[Perf] Improve GLM-5.3-Flash DFlash decoding on B300

[leonardHONG]

Motivation

Reduce synchronization and kernel overhead in GLM-5.3-Flash DFlash decoding. The DFlash/KDA-only comparison improves average decode throughput from 1091.85 to 1199.72 tokens/s (+9.88%) on 8×B300, with essentially unchanged memory usage.

Modifications

• Reuse CPU-side sequence lengths to avoid blocking GPU-to-CPU copies during DFlash input preparation.
• Build pooled page tables once per verification batch and reuse them across sparse attention layers.
• Tune KDA recurrence and gated norm launch configurations on SM103 while preserving the tested rounding behavior.
• Extend the sigmoid RMSNorm BT8 dispatch to BF16, D=128 workloads with 8–256 rows.

The FP8 MoE weight-layout optimization has been split out of this PR.

Accuracy Tests

Paired service checks matched output tokens and speculative acceptance statistics. Real-rejection cases also matched the checked generated-token logprobs.

Kernel checks cover KDA output/state parity and output/rstd parity for nearby norm shapes. Regression checks cover pooled metadata CUDA Graph replay and sequence-length bounds.

Speed Tests and Profiling

GLM-5.3-Flash + DFlash2, 8×B300, TP8, concurrency 1, 8192 input / 1024 output tokens, draft block size 8. Target linear-attention verification explicitly uses Triton (--linear-attn-verify-backend triton); draft attention uses FA4.

Three warmups followed by three rounds of two requests, without a profiler. The comparison pairs the same six measured inputs from the baseline and candidate records and aggregates decode throughput over those requests.

Configuration	Average decode tokens/s
Baseline	1091.85
DFlash/KDA optimized	1199.72
Improvement	+9.88%

Timing runs from receipt of the first token to SSE [DONE]. Timed requests had an acceptance length of 8; real-rejection cases were validated separately.

Before optimization, NSYS identified blocking sequence-length transfers and repeated pooled-page-table work; NCU showed low parallelism in KDA recurrence and gated norm.

Memory and KV Capacity

Both runs use the same model, draft, dependencies, and TP8 settings. Target and draft graph buckets are bs=[1] on both sides. Memory was measured separately, after all graphs were captured and before any inference requests.

Metric	Comparison
Target-load NVML sampled peak	Per-GPU difference no greater than 2 MiB
Draft-load NVML sampled peak	Identical on each GPU
Post-graph NVML resident memory	Identical on each GPU, 43.715–44.027 GiB/GPU
Post-graph PyTorch allocated	Candidate increases by 2.5 KiB/GPU
Post-graph PyTorch reserved	41.863281 GiB/GPU on both
Resolved max_total_num_tokens	32768 on both

Loading peaks are total process-memory samples from NVML with a target 50 ms sampling interval. Capacity was confirmed through the scheduler and server interface.

Performance and memory were measured on the recorded snapshots before the subsequent main merge.

Reproduction scripts, launch commands, source revisions, environment, and measurement summaries

---

CI States

Latest PR Test (Base): ❌ Run #37287781676
Latest PR Test (Extra): ⏳ Run #37287781460
Latest PR Test (AMD ROCm 10): ❌ Run #37287781809

[chatgpt-codex-connector]

You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard.
To continue using code reviews, add credits to your account and enable them for code reviews in your settings.

[BBuf]

I would split the optional FP8 MoE layout into a separate PR. It is independent of the DFlash changes and needs its own weight reload/offload validation.

Please include the full baseline/optimized commands, benchmark script and commit SHAs, especially linear-attn-verify-backend, and confirm the results on this head. The Triton KDA change won't run when verify uses FlashInfer.

Please also report peak loading memory, memory after graph capture, and final max_total_num_tokens for both runs with the same settings.

[BBuf]

On a partial weight update, end_weight_update() postprocesses every layer, so weights that weren't reloaded get shuffled a second time. Please undo the packing in restore_weights_before_loading(), or only repack weights that were actually reloaded. The zero-filled view test won't catch this; check that untouched weights stay unchanged.

[leonardHONG]

moved moe to a separate branch. it now restores the original layout before reload; nonzero-weight tests confirm untouched weights stay unchanged.

[BBuf]

Do we need to cache these views? view() doesn't copy data, and CUDA graph replay won't repeat this Python code. Creating the view from the current parameters would remove the cache_views option and offload checks, and avoid having to refresh aliases after reload.

[leonardHONG]

removed the cache and offload checks. views now come directly from the current parameters

[BBuf]

Please move this SM103 rule into _select_recurrent_launch_config instead of overriding its result here. Keep the rounding requirement tied to the selected configuration and explain why this exact shape needs it.

[leonardHONG]

moved into _select_recurrent_launch_config(). the narrow kda tile keeps separate fp32 multiply/add rounding to preserve parity

[BBuf]

Why is BT=8 restricted to exactly (64, 128, bfloat16)? Please document the reason, or use a workload-based rule if the gain also applies to nearby shapes.

[leonardHONG]

tested nearby shapes and expanded bt8 to 8–256 rows for sigmoid rmsnorm on sm103, with bf16 and d=128

[BBuf]

Can we avoid adding another environment variable and enable this layout conversion automatically on SM100 when using FlashInfer FP8 MoE with block_size=[128, 128]?

[leonardHONG]

removed the flag. the separate moe branch enables this automatically for supported sm100-family fp8 configurations