[kimseunghyun-kr]

Overview

topkGating (csrc/moe/topk_softmax_kernels.cu) gives each row to one full warp when the row is wide enough (512 experts, 256 experts with 16-bit logits, 128-512 experts with fp32 logits). For each of the k selections it runs a 5-stage shuffle butterfly with 3 dependent shuffles per stage, then lane 0 stores the selection and later re-reads the weights to scale them.

For those rows, without a correction bias, this PR selects each expert with two redux.sync reductions (__reduce_max_sync on the score key, then __reduce_min_sync on the expert among the lanes that hold the max), keeps selection i in the registers of lane i, rescans only in the winning lane, and writes the row once at the end. The output is bitwise equal to main: the same experts, the same order, the same lowest-index tie rule and the same weight bits.

The register-resident output path is from @Gin-Sin (#64731, co-author); the SM80+ redux.sync selection scope is from this PR. #64731 will be closed in favour of this PR.

Claims

• Same output: bitwise equal to main (ids, weights as bits, token-expert indices) in every case run on A100, including tie-heavy, all-equal and NaN/Inf rows.
• Faster on SM80+: E 512 top-10 softmax on A100: 0.59-0.68x of main's time at 1-256 tokens, 0.71-0.76x at 1,024-4,096 tokens. On B300 (measured by @Gin-Sin): 1.7-2.7x faster than main at every full-warp shape.
• End to end: below serving noise. On B300 (Qwen3.8-Flash-Next-FP8, DP4 + EP, MTP 3, with the [Perf][Kernel] Faster full-warp MoE top-k selection with redux.sync #64731 output path) throughput changed by at most about 1 %, inside the main-vs-main spread; GSM8K and greedy outputs are unchanged. Against main, the kernel saves 2.7-4.7 us per router call on A100 (1-256 tokens) and 2.5-6.9 us on B300 (E 512, 1-464 tokens).
• No change elsewhere: biased routing, k > 32, sub-warp rows (for example E 128 with 16-bit logits), ROCm and pre-SM80 GPUs keep the current loop, which is byte-identical to main. No new CPU/GPU sync, no new launch, no change to the op schema.

Validation

Kernel A/B, one A100-SXM4-40GB (sm_80), 31 Oct 2026. Four copies of topk_softmax_kernels.cu are built side by side as separate torch.utils.cpp_extension modules (CUDA 13.0, torch 2.13, -gencode=arch=compute_80,code=sm_80): main 1755058, the first revision of this PR c7f8f96, #64731 7a51d76, and this PR c752616. Each time is per call under CUDA graph replay (100 calls x 20 replays, minimum of 3 interleaved rounds). E 512, top-10, softmax, renormalize.

logits	tokens	main	first revision	#64731	this PR
fp32	1	9.99 us	7.69 us	6.48 us	5.31 us
fp32	8	8.35 us	6.36 us	5.77 us	5.34 us
fp32	64	8.69 us	6.71 us	5.69 us	5.37 us
fp32	256	9.01 us	7.12 us	5.85 us	5.38 us
fp32	1,024	11.84 us	10.90 us	9.17 us	8.71 us
fp32	4,096	26.54 us	25.54 us	21.60 us	20.25 us
bf16	1	7.79 us	5.85 us	5.45 us	5.02 us
bf16	8	8.56 us	6.62 us	5.76 us	5.33 us
bf16	64	8.67 us	6.73 us	5.72 us	5.24 us
bf16	256	9.00 us	7.05 us	5.76 us	5.38 us
bf16	1,024	11.74 us	10.76 us	9.03 us	8.35 us
bf16	4,096	26.35 us	25.37 us	21.47 us	19.98 us

Bitwise checks against main on the same A100: every timed case above, plus 36 tie-heavy cases (fp32 / bf16 / fp16 x tokens 1 / 7 / 64 / 333 x few distinct values / all-equal rows / wide randn): 0 mismatches for all three changed copies.

B300 (sm_100), measured by @Gin-Sin on this head c752616 (comment): CUDA 13.1, FlashInfer CUPTI timing with CUDA graph and cold L2, median of 3 rotated rounds; bitwise equal to main in 336/336 case-runs.

logits	E / top-k	tokens	main	#64731	this PR
bf16	512 / 10	1	6.50 us	4.00 us	3.78 us
bf16	512 / 10	64	8.86 us	3.97 us	3.74 us
bf16	512 / 10	464	11.04 us	4.32 us	4.10 us
bf16	512 / 10	8,192	32.96 us	19.90 us	18.62 us
fp32	256 / 8	64	7.78 us	3.36 us	3.17 us
bf16	128 / 8 (sub-warp, unchanged path)	64	6.24 us	6.24 us	6.24 us

End-to-end serving and GSM8K on B300 with the #64731 output path are in this comment: no difference beyond run-to-run noise at concurrency 16/64, about +1 % at 464 (at the edge of noise); GSM8K 0.967/0.969 (main) vs 0.971/0.969; greedy outputs identical on 200/200 prompts.

Details

Why it is exact. Without a bias the choice score is the output score. After softmax or sigmoid and the existing NaN/Inf clamp every score is >= +0, so the unsigned order of the float bits is the float order. Each lane offers key = score bits + 1 for its best unselected slot (a selected slot has key 0). The first redux.sync finds the maximum key; the second finds the minimum expert among lanes that hold it, which is the lowest expert index, the same tie rule as the butterfly. Only the owner lane clears its slot and rescans. Lane i keeps selection i; lanes 0..k-1 write once at the end, with the same fp32 renormalization sum order and the same scale as main.

Scope. topkGating rows with THREADS_PER_ROW == WARP_SIZE, no bias, k <= 32, __CUDA_ARCH__ >= 800, not ROCm; the launcher picks a separate FULL_WARP_SELECT instantiation, so the existing loop's code generation does not change. Rows beyond num_rows return per warp before the selection, so all 32 lanes reach each redux.sync. TensorRT-LLM uses redux.sync for this selection on SM100+.

Tests. test_topk_without_bias_matches_zero_bias (from #64731, extended) compares a launch without a bias against a launch with an all-zero bias, which keeps the old loop, bit for bit: E 128/352/384/512/576, k 1/10/32, bf16/fp16/fp32, softmax/sigmoid, renormalize on/off, is_padding rows, routed_scaling_factor, tie-heavy rows with NaN/+Inf, and the tie order against a stable sort (720 cases). It has not run in a full vLLM build on our side; CI will run it. The benchmark gains a --dtype option for fp32 router logits.

CI. There is no sm_80 runner. The A100 data come from our cluster.

Related. #64731 (Gin-Sin, same selection; its output path is folded in here). #59737 (PDL for this kernel, SM90+ only) does not overlap. #57823 (padded row strides) touches the same file.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

AI assistance: Claude Code (Claude Opus 5.5) wrote the kernel change, the tests and the benchmark; I reviewed them and ran the A/B on our A100 cluster.

🤖 Generated with Claude Code

[Gin-Sin]

Hi @kimseunghyun-kr, thanks for this PR. I was working on the same topkGating full-warp path and ended up with a very similar redux.sync selection. I just opened #64731 as a draft to share the code and numbers, and I'd like to converge with you rather than compete.

The difference is the output path:

• Each lane keeps its selection in registers, and the row is written once at the end; only the winning lane rescans.
• [Perf][MoE] Select full-warp top-k rows with redux.sync and keep selections in registers (SM80+) #62912 stores to global memory every round and reads the weights back.

On B300 (E=512, top-10, fp32 softmax), all three bitwise equal to main in 11,520 cases:

Tokens	main	#62912	#64731
1	5.33 µs	3.97 µs	3.05 µs
64	8.10 µs	6.96 µs	3.10 µs
128	8.88 µs	7.59 µs	3.10 µs
4096	19.96 µs	19.69 µs	15.71 µs

An ablation (my selection with your per-round store) runs at your PR's speed, so the gap comes from the write path.

Happy to go whichever way you prefer:

1. Fold the output-path change into [Perf][MoE] Select full-warp top-k rows with redux.sync and keep selections in registers (SM80+) #62912 (co-author credit welcome), and I'll close [Perf][Kernel] Faster full-warp MoE top-k selection with redux.sync #64731.
2. Land [Perf][MoE] Select full-warp top-k rows with redux.sync and keep selections in registers (SM80+) #62912 first, and I'll rebase [Perf][Kernel] Faster full-warp MoE top-k selection with redux.sync #64731 on top as a small follow-up.

If you can try it on A100, it'd be great to confirm the gain holds there too.

[kimseunghyun-kr]

Thanks @Gin-Sin, and thanks for the clear ablation. It shows that the write path, not the selection, is where #62912 loses time, so your version is the better kernel. Let's converge.

I prefer option 1: I fold your output path (selection kept in registers, one row write at the end, rescan only in the winning lane) into #62912 and add you as co-author. Then you can close #64731. If you would rather keep your own authorship on that part, option 2 is also fine with me. Tell me which you prefer.

I will run main, #62912 and #64731 on A100-SXM4-40GB (E=512, top-10, the same token counts as your table) and post the numbers here. Our deployment is Qwen3.8-Flash-Next-FP8 on 8 × A100, where decode steps carry 4–32 tokens, so the 1–128 token range matters most for us.

(Drafted with AI assistance; I reviewed it.)

[Gin-Sin]

Thanks @kimseunghyun-kr, option 1 sounds great. Please fold the output path into #62912; co-author credit as Co-authored-by: Gin-Sin <ginxian@foxmail.com> would be appreciated.

The change is in #64731 (single commit 7a51d76027), so you can cherry-pick from there. It also adds:

• test_topk_without_bias_matches_zero_bias in tests/kernels/moe/test_fused_topk.py: bitwise check of the no-bias launch against an all-zero-bias launch, which takes the old loop.
• a --dtype option in benchmarks/kernels/benchmark_fused_topk.py, for fp32 router logits.

Feel free to take or drop either.

I'll follow up here with end-to-end numbers on B300 for Qwen3.8-Flash-Next (DEP4 + MTP3 on the DeepGEMM MegaMoE path from #64315, which calls topk_softmax), with GSM8K, once my last runs finish. Looking forward to your A100 numbers. I'll close #64731 once the change is in #62912.

[kimseunghyun-kr]

@Gin-Sin, here are the A100 numbers. A100-SXM4-40GB, E=512, top-10, renormalize. Each revision of topk_softmax_kernels.cu is built for sm_80 (CUDA 13.0, torch 2.13): main dc4dd5a, #62912 c7f8f96 and #64731 e1a0d3c. Times are per call under CUDA graph replay (100 calls × 20 replays, minimum of 3 interleaved rounds). All outputs (weights, ids, token_expert_indices) are bitwise equal to main in every row.

gating	tokens	main	#62912	#64731
fp32	1	9.88 µs	7.64 µs	9.92 µs
fp32	8	8.42 µs	5.98 µs	8.43 µs
fp32	64	8.74 µs	6.89 µs	8.71 µs
fp32	128	9.15 µs	6.99 µs	9.14 µs
fp32	1024	12.02 µs	11.01 µs	12.02 µs
fp32	4096	26.56 µs	25.59 µs	26.56 µs
bf16	1	7.85 µs	5.89 µs	7.85 µs
bf16	64	8.73 µs	6.72 µs	8.73 µs
bf16	4096	26.44 µs	25.37 µs	26.43 µs

all rows (4–256 tokens)

gating	tokens	main	#62912	#64731
fp32	4	8.17	6.21	8.11
fp32	16	8.45	6.28	8.43
fp32	32	8.56	6.57	8.56
fp32	256	9.18	7.21	9.18
bf16	4	8.38	6.20	8.37
bf16	8	8.39	6.19	8.41
bf16	16	8.51	6.62	8.51
bf16	32	8.57	6.61	8.59
bf16	128	8.90	6.91	8.90
bf16	256	9.16	7.12	9.14

#64731 runs at main's speed here because its full-warp path is compiled only for __CUDA_ARCH__ >= 1000 && < 1100, so on sm_80 it takes the existing loop. redux.sync is available from sm_80, so I think the converged kernel can enable your register-resident output path on SM80+ too. Was there a reason to limit it to SM100 (for example a regression you saw on Hopper)? If not, I will fold your output path into #62912 with the SM80+ gate, add you as co-author, and post A100 and H100 numbers of the combined version here.

(Drafted with AI assistance; I reviewed it.)

[kimseunghyun-kr]

Thanks @Gin-Sin. I am folding 7a51d76027 into this PR now, with Co-authored-by: Gin-Sin <ginxian@foxmail.com>, and I will take the zero-bias test and the --dtype option too. I will gate the full-warp path on SM80+, so A100 and H100 also get your register-resident output path. Before I push, I will check on A100 that the combined kernel is bitwise equal to main and post its numbers here. Looking forward to your B300 end-to-end numbers.

(Drafted with AI assistance; I reviewed it.)

[Gin-Sin]

Thanks for the A100 numbers, @kimseunghyun-kr. e1a0d3c925 was a short-lived revision of #64731 that gated the full-warp path to SM100. I restored the earlier commit shortly after: the current head 7a51d76027 (the commit you are folding in) enables it on SM80+ (__CUDA_ARCH__ >= 800, CUDA only). There was no Hopper regression behind it; I just haven't measured on Hopper or Ampere, so your A100 and H100 numbers for the combined kernel are very welcome. SM80+ sounds right to me.

[Gin-Sin]

End-to-end data on B300 for the folded-in kernel (the #64731 output path), as promised.

Setup.

• Model: Qwen3.8-Flash-Next-FP8 (public checkpoint; 512 experts, top-10, softmax).
• Code: main fac2ce2 plus [Model][Qwen4Exp] Run routed experts on DeepGEMM MegaMoE (MXFP8 or MXFP4) #64315 at 5ed2362 (MegaMoE, MXFP8 routed experts). Its router calls fused_topk → topk_softmax.
• Serving: one B300 node, DP4 + EP on 4 GPUs per engine, MTP3, --max-num-seqs 128.
• Load: vllm bench serve, random 2048 in / 1024 out, --ignore-eos, at concurrency 16, 64 and 464. That is about 16, 64 and 460 routed tokens per rank per verify step, which covers the 4-32 tokens per step range you tested.
• Placement: main and kernel engines side by side on GPUs 0-3 and 4-7, 2 runs each, in both placements, plus a main-vs-main pair for noise.
• The default Blackwell MoE path for this model (FlashInfer TRTLLM) routes inside its own kernel and does not call topk_softmax.

Kernel in-engine (TP1 torch profile, per topkGating call): 8.22 → 4.87 µs at 64 decode tokens, and 12.97 → 11.32 µs at a 4k-token prefill. Standalone (CUPTI, CUDA graph, cold L2), E=512 top-10 bf16:

Tokens	main	#62912	#64731
64	8.80 µs	7.39 µs	4.00 µs
464	10.98 µs	9.54 µs	4.32 µs
8192	32.99 µs	28.13 µs	19.90 µs

Serving (Δ is kernel vs main, mean of 2 runs; acceptance length is the MTP mean per arm):

Conc	Set	main tok/s	kernel tok/s	Δ tok/s	Δ mean TPOT	Accept. length main → kernel
16	main 0-3, kernel 4-7	2,718 / 2,516	2,568 / 2,613	−1.0%	−0.1%	2.40 → 2.41
16	swapped	2,481 / 2,590	2,558 / 2,476	−0.8%	+1.6%	2.39 → 2.33
16	noise: main vs main	2,572 / 2,572	2,651 / 2,591	+1.9%	−0.5%	2.40 → 2.40
64	main 0-3, kernel 4-7	5,938 / 6,023	6,010 / 5,994	+0.4%	−0.5%	2.43 → 2.42
64	swapped	5,846 / 5,942	5,924 / 5,744	−1.0%	+0.8%	2.47 → 2.42
64	noise: main vs main	6,013 / 6,000	5,940 / 6,019	−0.5%	+0.6%	2.45 → 2.40
464	main 0-3, kernel 4-7	17,203 / 16,786	17,103 / 17,331	+1.3%	−1.1%	2.39-2.47 per run, all arms
464	swapped	17,100 / 17,071	17,184 / 17,229	+0.7%	−0.2%
464	repeat	16,728 / 17,319	17,120 / 17,344	+1.2%	−1.0%
464	noise: main vs main	16,868 / 17,221	16,703 / 17,111	−0.8%	−1.5%

Verdict:

• c16 and c64: no measurable e2e difference. Both placements land within the main-vs-main spread (+1.9% at c16, −0.5% at c64). The expected saving is about 1-1.5% of the step (≈0.2 ms of ~13 ms at c16, ≈0.23 ms of ~23 ms at c64), and single runs of the same engine differ by up to 8% at c16.
• c464: about +1%, at the edge of noise. The kernel arm is ahead in all three sets, but two identical engines differ by up to 1.5%. The expected saving is about 0.6%.
• An earlier single-placement c16 run showed −5.5%; it did not reproduce in either placement. Per-run acceptance length ranges 2.28-2.48 at c16 in every arm, main vs main included, and tracks throughput (correlation 0.62). Greedy outputs are identical, so that spread is batch composition, not the kernel.

So the kernel-level gain is real, while e2e on this setup it is at most about 1% and below run-to-run noise at every concurrency tested.

Accuracy (GSM8K, gsm8k_eval.py --use-chat-completions, thinking off, temperature 0, 1,319 questions, 5-shot):

• main 0.967 / 0.969 vs kernel 0.971 / 0.969;
• main vs main 0.970 / 0.965;
• default FlashInfer TRTLLM backend 0.968.

Greedy identity (200 GSM8K chat prompts, 192 tokens): main vs kernel 200/200, and main vs main 200/200.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[kimseunghyun-kr]

@Gin-Sin, the fold is in (c752616, with your Co-authored-by), and the PR is now ready for review.

One correction first: my earlier A100 table tested #64731 at e1a0d3c, which still had the SM100-only gate. Your current head 7a51d76 already gates on SM80+, so my "runs at main's speed" note does not apply to it. Sorry for the confusion.

A100-SXM4-40GB, E 512, top-10, softmax, per call under CUDA graph replay, all four copies built side by side (sm_80):

logits	tokens	main	#62912 (first revision)	#64731 (7a51d76)	#62912 now
fp32	1	9.99 µs	7.69 µs	6.48 µs	5.31 µs
fp32	64	8.69 µs	6.71 µs	5.69 µs	5.37 µs
fp32	4,096	26.54 µs	25.54 µs	21.60 µs	20.25 µs
bf16	1	7.79 µs	5.85 µs	5.45 µs	5.02 µs
bf16	64	8.67 µs	6.73 µs	5.72 µs	5.24 µs
bf16	4,096	26.35 µs	25.37 µs	21.47 µs	19.98 µs

All three changed copies are bitwise equal to main in every timed case and in 36 extra tie-heavy cases (few distinct values, all-equal rows, wide randn; fp32 / bf16 / fp16). The full table is in the PR description.

The combined kernel uses your key/redux structure and your output path; the main differences are a 4-argument selection helper that returns {score, expert, selected_sum} (the kernel does the write) and the existing loop kept byte-identical to main. I took your zero-bias test (extended to E 352/576, k 32, fp16, padding rows and the scaling factor) and the --dtype option. If you can rerun your B300 numbers on c752616, that would complete the picture for reviewers.

(Drafted with AI assistance; I reviewed it.)

[Gin-Sin]

Thanks @kimseunghyun-kr, and thanks for the correction. Here are the B300 numbers for the combined kernel.

B300, topk_softmax alone, renormalize, no bias. FlashInfer CUPTI with CUDA graph and cold L2, median of 3 rounds with the build order rotated each round (largest round-to-round spread 5.4%). All four builds use the same flags (CUDA 13.1, arch 10.0). Outputs (weights as int32 bits, ids, token-expert indices) are bitwise equal to main in 336/336 case-runs.

Logits	E / top-k	Tokens	main	#62912 first revision	#64731 (7a51d76)	#62912 now (c752616)	vs main	vs #64731
bf16	512 / 10	1	6.50 µs	5.18 µs	4.00 µs	3.78 µs	1.72x	1.06x
bf16	512 / 10	64	8.86 µs	7.49 µs	3.97 µs	3.74 µs	2.37x	1.06x
bf16	512 / 10	464	11.04 µs	9.60 µs	4.32 µs	4.10 µs	2.69x	1.05x
bf16	512 / 10	8,192	32.96 µs	28.13 µs	19.90 µs	18.62 µs	1.77x	1.07x
fp32	512 / 10	64	8.42 µs	7.14 µs	4.06 µs	3.87 µs	2.17x	1.05x
fp32	256 / 8	64	7.78 µs	6.11 µs	3.36 µs	3.17 µs	2.45x	1.06x
bf16	128 / 8 (sub-warp)	64	6.24 µs	6.24 µs	6.24 µs	6.24 µs	1.00x	1.00x

All 28 cases

Case	main	#62912 first revision	#64731	#62912 now	vs main	vs #64731
bf16 E512 k10 t1	6.50	5.18	4.00	3.78	1.72x	1.06x
bf16 E512 k10 t4	6.27	5.09	3.87	3.68	1.70x	1.05x
bf16 E512 k10 t16	7.96	6.46	3.90	3.68	2.16x	1.06x
bf16 E512 k10 t64	8.86	7.49	3.97	3.74	2.37x	1.06x
bf16 E512 k10 t128	8.74	7.46	4.06	3.84	2.28x	1.06x
bf16 E512 k10 t256	9.60	8.32	4.13	3.90	2.46x	1.06x
bf16 E512 k10 t464	11.04	9.60	4.32	4.10	2.69x	1.05x
bf16 E512 k10 t1024	11.39	10.37	5.22	5.02	2.27x	1.04x
bf16 E512 k10 t2048	14.43	12.99	7.58	7.17	2.01x	1.06x
bf16 E512 k10 t4096	20.54	17.95	11.42	10.75	1.91x	1.06x
bf16 E512 k10 t8192	32.96	28.13	19.90	18.62	1.77x	1.07x
fp32 E512 k10 t1	5.98	4.67	3.74	3.52	1.70x	1.06x
fp32 E512 k10 t4	6.27	4.99	3.90	3.71	1.69x	1.05x
fp32 E512 k10 t16	7.71	6.43	3.90	3.71	2.08x	1.05x
fp32 E512 k10 t64	8.42	7.14	4.06	3.87	2.17x	1.05x
fp32 E512 k10 t128	9.41	8.13	4.19	4.00	2.35x	1.05x
fp32 E512 k10 t256	9.02	7.71	4.19	4.00	2.25x	1.05x
fp32 E512 k10 t464	10.14	8.80	4.32	4.13	2.46x	1.05x
fp32 E512 k10 t1024	11.62	10.72	5.41	5.19	2.24x	1.04x
fp32 E512 k10 t2048	14.72	13.06	8.03	7.62	1.93x	1.05x
fp32 E512 k10 t4096	21.06	18.72	12.10	11.39	1.85x	1.06x
fp32 E512 k10 t8192	33.25	28.64	20.42	19.07	1.74x	1.07x
fp32 E256 k8 t1	4.77	3.39	3.01	2.82	1.69x	1.07x
fp32 E256 k8 t64	7.78	6.11	3.36	3.17	2.45x	1.06x
fp32 E256 k8 t4096	12.96	11.14	7.20	6.62	1.96x	1.09x
bf16 E128 k8 t1 (sub-warp)	4.58	4.58	4.58	4.58	1.00x	1.00x
bf16 E128 k8 t64 (sub-warp)	6.24	6.24	6.24	6.24	1.00x	1.00x
bf16 E128 k8 t4096 (sub-warp)	8.58	8.61	8.58	8.61	1.00x	1.00x

Times in µs.

So on B300 the combined kernel is another 4-9% faster than #64731 at every full-warp shape (1.7-2.7x vs main), consistent with your A100 numbers. Computing each lane's candidate expert once per rescan, and matching the winner's owner with one compare, takes the per-round index arithmetic out of the loop. Sub-warp rows are unchanged, as expected.

I didn't rerun end to end: at about 0.2 µs per call this is far below the serving noise in my earlier comment, and the outputs are bitwise identical to the kernel measured there. LGTM from my side; I'll close #64731 once this lands.

[kimseunghyun-kr]

Thanks @Gin-Sin, for the B300 run on this head and for the review.

The B300 results match the A100 results: the combined kernel is 4-9 % faster than #64731 at every full-warp shape, 1.7-2.7x faster than main, and bitwise equal in 336/336 case-runs. The sub-warp rows do not change. I added a short B300 table to the PR description, with links to your two comments for the full 28 cases and for the end-to-end and GSM8K data. I also stated in the Claims that the end-to-end gain is below serving noise, as you measured.

I agree that a rerun end to end is not necessary: the change is about 0.2 µs per call against #64731, and the outputs are bitwise identical.

(Drafted with AI assistance; I reviewed it.)

[kimseunghyun-kr]

@Gin-Sin, thanks again for the B300 data. If the combined kernel on the current head (c752616) looks right to you, could you leave an approving review? It would help the maintainers see that both of us checked it.

(Drafted with AI assistance; I reviewed it.)

[Gin-Sin]

Approving c752616. I read the kernel and test changes. On B300 this head is bitwise equal to main in 336/336 case-runs and 1.7-2.7x faster at the full-warp shapes, with sub-warp rows unchanged (numbers in my comment above). The SM80+ gate and the separate FULL_WARP_SELECT instantiation, which keeps the existing loop's code generation untouched, look right to me.

Thanks for folding it in, @kimseunghyun-kr. I'll close #64731 once this lands.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @kimseunghyun-kr.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[kimseunghyun-kr]

@bnellnm, could you take a look when you have time? This PR is rebased on current main (0d29a48). The only conflict was #65059 (PDL launch attribute only on SM90+); both changes are kept, and our two commits are unchanged (git range-diff). The PR body has the A100 numbers, and @Gin-Sin's B300 run on the same head is in the comments above.

(Drafted with AI assistance; I reviewed it.)