[BBuf]

Motivation

Single-token BF16 Qwen3.8-Flash-Next decode launches separate HC combine/norm and MoE shared-gate/finalize/collective work, while the packed GDN value tile leaves substantial register pressure at this shape. On four B200 GPUs, this patch raises the measured 8k-input/1k-output decode rate from fresh stock 242.443 to 270.982–271.759 token/s, preserving the sampled stock accuracy in four native runs.

Modifications

• Overlap the HC gate with block computation, fuse HC apply with grouped Gemma RMSNorm, and pass the prepared norm tensors explicitly across layers. PLE and PP boundaries retain the normal handoff.
• Extend the existing collective with an FP32 shared gate, preserving routed BF16 rounding before the shared-expert FMA. Defer BF16 MoE finalize and overlap the shared branch through a guarded single-token TP4/EP1 path with its own communication buffer.
• Use the smaller GDN value tile for the packed single-token BF16/FP32-state Blackwell shape, preserving the one-warp reduction order. Existing dispatch remains for unsupported shapes.

The MoE specialization requires hidden size 2560, 512 experts/top-10, BF16, SM10, TP4/EP1, attention DP1, with LoRA, speculative decoding and two-batch overlap disabled. No extra command-line switch enables the kernels. The benchmark selects the existing faster custom-all-reduce deployment option for both stock and candidate.

Accuracy Tests

Stock and all four native candidate runs score 126/128 GSM8K, zero invalid answers, zero truncations. These use identical frozen tokenized prompts, normal EOS and the same 2,048-token cap. All 128 per-question correctness and finish-reason verdicts match the fresh stock run.

Historical GPU validation before the standalone test cleanup: 6 HC unit configurations; 12 independent torch-reference GDN configurations; 6 collective configurations on each of 4 ranks, including the original ungated API; 90 trained HC-chain comparisons; 324 changed-input trained collective graph comparisons; 192 seeded GDN output/state bit comparisons. Isolated HC memcheck: zero errors. Full-stack communicator checks have the same handled PyTorch FABRIC capability API errors in stock and candidate, so they are not claimed as clean full-stack sanitizer runs. Pre-commit passes.

Speed Tests and Profiling

Same four B200 GPUs; BF16 weights/KV/indexer, FP32 GDN state, TP4, GPU PLE, context 16384, concurrency one, no speculative decoding or prefix caching. Each speed run has 3 warmups + 32 measured requests. Rate is 1023 / mean(last_token_time - first_token_time); TTFT/E2E are retained separately. Profiling runs are separate.

Run	Decode token/s	TTFT ms	E2E s	GSM8K	Truncated
Fresh stock SGLang, custom AR	242.443	140.04	4.3597	126/128	0
Native candidate, initial validation	271.032	141.60	3.9162	126/128	0
Pair 1, sglang	271.759	141.33	3.9058	126/128	0
Pair 1, vllm	236.285	122.56	4.4522	125/128	1
Pair 2, sglang	271.229	141.02	3.9128	126/128	0
Pair 2, vllm	236.591	123.40	4.4475	126/128	0
Pair 3, sglang	270.982	141.45	3.9167	126/128	0
Pair 3, vllm	236.265	122.46	4.4525	126/128	0

Raw paired speed ratios: 1.150134x, 1.146404x, 1.146941x. Against the fastest observed GPU-PLE vLLM result (240.786 token/s), even the slowest native candidate is 12.54% faster. That conservative denominator includes the quality-failing EP4 result. vLLM's first final TP4 repeat also has one truncation and remains a failed quality observation; this comparison does not establish stable quality parity for vLLM. CPU-offloaded PLE trials are excluded from the matched comparison. GPU-PLE CUTLASS and Triton trials are included in the deployment-search data.

Stock has 1,432 kernels per decode graph. All four candidate traces have 1,241 per graph, including 95 HC apply/norm, 48 gated finalize/all-reduce and 36 packed GDN calls. Kernel sums overlap through PDL and are not critical-path latency. The pinned sources, portable commands, workload hashes, all request timings/verdicts, deployment search, layer guides and NCU evidence are in the reproducer. The results apply to the measured text-decode workload.

SGLang base f03a183719c9e7fb2bd528ed5c7e3b627d9df92d, measured head b66c816cf6475dbaec3e1b5d2d68a3f8c0d77ab1; vLLM b538d807ffe0976af7e5d107ccfe77e676b83b70. Model revision de4b8e4d43b917e7706784d8bb445c9af86a3540.

Checklist

• Format with pre-commit.
• Retain layer-boundary correctness tests.
• Document specialization guards and provide portable reproduction instructions.
• Provide accuracy, speed and profiling evidence.
• Follow SGLang code style.

---

CI States

Latest PR Test (Base): 🚫 Run #37912429343
Latest PR Test (Extra): ⏳ Run #37912428916
Latest PR Test (AMD ROCm 10): ⏳ Run #37912429512

CI follow-ups

The CPU DP fixture now initializes _decode_moe_comm=None; its original 9 tests pass locally and CPU partition 3 passed in run 36958545032. Initialization eligibility now reads the resolved LoRA, speculative-decoding and overlap namespaces. Both config-accessor checks, all 170 runtime-context tests, 8 CPU guard combinations, and pre-commit pass.

The original submission dadeedd3c2449138221ebb961b40d7b61f9c2ae8 changed the CPU fixture and initialization guard/imports after measured head b66c816cf6475dbaec3e1b5d2d68a3f8c0d77ab1. The full-model accuracy and throughput figures above remain attributed to that measured head. The cleanup validation below compares against the original submission; it does not replace the full-model benchmark.

Cleanup validation (2026-10-05)

Head 6f1a16068294cf4f956357184dad491b2a8ae499 removes can_use_qwen4_decode_moe, resolves model eligibility during initialization, and keeps the token-count/dtype/compiler choice at the call site. HC stream initialization is now uniform, redundant shape/dtype guards are removed, and kernel names/types are clearer. CUDA host validation and Triton host checks retain the input contracts. Arithmetic, launch geometry and stream dependencies are unchanged. The three standalone kernel correctness test files were present for this validation; they were subsequently removed from the PR during test cleanup.

Four-B200 validation passed on this exact head: 6 HC configurations, 12 GDN reference configurations, 6 collective configurations across four ranks, and 9 CPU reduction tests. Additional old/new comparisons passed bitwise: 15 HC kernel cases, 20 shared-gate cases, 48 GDN state/output cases and 40 real HC CUDA-graph cases, including FP16 fallback and terminal mixers. Pre-commit passed; 16 initialization and 48 runtime dispatch combinations selected the same paths.

Each timed graph replay first evicts 512 MiB on the measured GPU, outside the timed interval. Each HC measurement times the gate/combine/norm chain without intermediate eviction; GDN state resets happen before eviction. Measurements use 7 blocks of 31 interleaved samples per version. HC and shared-gate median changes were -0.38% to 0%. The initial GDN batch-2 comparison with separate state/output allocations measured +1.86%; a same-buffer repeat on another B200 measured identical old/new medians for batches 1/2/4. GDN and shared-gate PTX and register/spill/shared-memory counts are identical. This is kernel-level validation, not a new end-to-end throughput or GSM8K run.

The main integration and its separate validation are recorded below.

Main integration validation (2026-10-05)

Head 71f1b77e8a1b32de129100890d1aaebd937962d5 merges main through 456d533e61ce9339296cdeb1defe884690528f9b and resolves the Qwen4 stage-boundary migration. Gate partials now travel with the gated residual state; prepared HC norms respect PLE/PP and distributed-layout boundaries. A fused MoE result still completes boundary moves and residual write-back without summing the already-reduced result twice. The specialization is disabled for MoE context-parallel all-gather.

Compared with pre-merge head 6f1a16068294cf4f956357184dad491b2a8ae499 on the same 4x GB300, Torch 2.14.1+cu130:

• GSM8K: 126/128 on both versions, zero invalid answers and zero truncations. All 128 predicted numerical answers, correctness and finish reasons match. Same frozen workload as above, normal EOS and 2,048-token cap.
• 306 layer-boundary tests and 768 subtests passed, plus 6 HC configurations, 12 GDN reference configurations and 6 collective configurations on all four ranks.
• Bitwise comparisons passed for 15 HC kernel cases, 20 shared-gate cases, 48 GDN output/state cases, 40 real HC graph cases and 120 two-layer boundary graph cases. The boundary harness covers FP16/BF16, empty and nonempty rows, PLE contribution and terminal mixers; its pointwise attention/MLP functions isolate boundary integration rather than full distributed-model behavior.
• Local pre-commit, whitespace checks and GitHub lint passed. Kernel arithmetic and launch configurations are unchanged by the merge. GDN/shared-gate PTX and resource counts match the pre-merge version.

Full-model fixed-input decode uses BS1/2/4, 8,192 input tokens and 256 output positions, BF16 weights/KV/indexer, FP32 GDN state, TP4 and GPU PLE. Every rank evicts 512 MiB before each prefill/decode forward, outside timing; there is no eviction between kernels within a forward. Fixed decode inputs prevent generated-token divergence from changing the input workload. All logits were finite. Wall-clock timing includes the benchmark's host preparation and sampling.

Batch	Before merge, median ms	After merge, median ms	Candidate repeat, ms	Change across candidate runs
1	4.8936	4.9809	4.9611	+1.38% to +1.78%
2	5.3884	5.4321	5.4096	+0.39% to +0.81%
4	5.5130	5.5396	5.5130	0% to +0.48%

Fixed-input logits are also not bitwise reproducible within the same candidate: top-1 agreement across two candidate runs was 223/255, 438/510 and 862/1020, versus 224/255, 442/510 and 858/1020 across versions. Common top-5 maximum absolute differences were 4.875/5.3125/4.5 within the candidate and 4.625/5.125/5.375 across versions. The source of this variation has not been localized; these comparisons do not prove full numerical equivalence. The functional evidence is the matching GSM8K answers and direct kernel/boundary checks above.

These are measured small latency increases, not a claim of zero overhead. Three earlier autoregressive synthetic-prompt runs measured +0.53% to +1.99%; they also diverged in generated tokens between identical-candidate runs, so they do not establish token parity or matched routing. The earlier B200 serving-throughput numbers remain attributed to their original measured commit and were not re-measured as B200 numbers here.

Kernel/HC-chain measurements evict max(512 MiB, 4 x L2) = 542,113,792 bytes before every timed replay, excluded by external CUDA events, with no intermediate eviction inside a chain. Seven blocks of 31 interleaved samples retain raw data. Individual HC-chain changes range from -4.30% to +4.49% (about 0.35 us); two-layer boundary changes range from -2.21% to +2.21%; same-buffer GDN changes range from 0% to +3.56% despite identical PTX/resource counts.

Full CI is not claimed green. The inspected AMD rotary failure follows main's dtype revert; the inspected NPU HiCache failure is tensor truth conversion in pool_host/mha.py. Both implicated implementations are unchanged by this PR. Extra workflow gates require run-ci-extra and did not execute tests. No tests or thresholds were weakened to bypass these failures.

Cleanup validation (2026-10-09)

Head 61796310e9c04efbf5356c9454fd82d310f56b0d uses a fixed (residual, normed, optional_partials) HC state, removes tuple-length dispatch and its adapter, and reuses the existing fused gate device kernel for FP32 gate-only output. The five public decode entries now have kernel registry metadata; collective template arguments and residual/communicator documentation are clarified. Existing test files were extended; no test files were added and the three standalone files remain deleted.

Compared with f1733a369f1392d60e82e9737db9591ca6cf18ae on the same 4x GB300, Torch 2.14.1+cu130, Triton 3.8.0:

• Candidate: 377 unit tests and 768 subtests passed, plus 94 gate tests, 18 HC cases, 12 GDN reference cases, and 6 collective configurations on each of four ranks. The collective checks include changed-input graph replay and the original ungated API. Baseline checks also passed. Pre-commit passed.
• 100 FP32 gate comparisons were bitwise equal. Real HC modules matched bitwise for BF16/FP16 and rows 0/1/2/4/32, with five input updates per configuration, prepared-norm reuse and terminal mixers. Empty-row captures contain no GPU work. Gate resource counts are unchanged: 23 registers, zero spills, 64 bytes shared memory. Its cold-L2 graph median was 8.32 us on both versions, using 217 interleaved samples per version and a 512 MiB eviction outside each timed replay. This is a gate-level measurement, not an HC-chain measurement.
• Full-model GSM8K 32/32 on both versions, with all answer texts and finish reasons matching, zero invalid answers and zero truncations. These are the first 32 frozen examples from the earlier 128-question workload, normal EOS and the same 2,048-token cap.
• Autoregressive decode, 8,000 input / 1,024 output tokens, BS1, TP4, BF16 weights/KV/indexer, FP32 GDN state and GPU PLE: 3 warmups + 16 measured requests per version, all successful. Rate was 277.127 -> 276.606 token/s (-0.19%); mean E2E was 3.8232 -> 3.8299 s. This single pair does not establish zero performance overhead. Long synthetic outputs were not token-identical across versions; a fresh same-candidate restart also differed on all four repeated prompts. This reproduces the previously recorded full-model nondeterminism and does not establish full token-level numerical equivalence or identical routing.

Model cache snapshot: f5d08274bafd880402bd16f5e3e6c514136ec06c; configuration and weight-index hashes match the earlier reproducer. Both source snapshots use SGLANG_RUST_BUILD_MODE=auto to build the required Rust extension. The initial unmodified baseline startup failed with the image's inherited never setting because its bundled extension was outside the source snapshot; no production source change was made for this environment issue. All changed source files were SHA-256 checked against the measured snapshot before committing. These results do not replace the historical B200 measurements or claim that full GitHub CI is green.

The CI runs on cleanup head 61796310e9 stopped before code tests at the minimum-base and extra-label gates. On 2026-10-09, merge commit 0df30f3a8f7915ba6e1b629d673ab9bd680e801e merged main 9a1aeacc3bcd3529445a31db8f71394e22e390ca, which includes required base 60878c3fba50f50978c8a38c8ea3346022bf355e. Both run-ci and run-ci-extra are enabled. Pre-commit passed on all 21 PR-changed files after the merge. New base CI and extra CI runs have been triggered; their final results are pending. The GPU results above remain attributed to cleanup head 61796310e9 and were not rerun on this merge commit.

[yuan-luo]

LGTM. Please add registered CUDA regression tests for the new split-gate / apply-norm path.