Skip to content

feat(cake_comm): support SM100 and SM103 TP8 packed-QKV all-gather matmul - #5116

Merged
yyihuang merged 7 commits into
flashinfer-ai:mainfrom
yyihuang:codex/agmm-tp8-qkv-peer-grid4-v32
Sep 14, 2026
Merged

yyihuang merged 7 commits into
flashinfer-ai:mainfrom
yyihuang:codex/agmm-tp8-qkv-peer-grid4-v32

Conversation

@yyihuang

@yyihuang yyihuang commented Sep 11, 2026 •

Copy link
Copy Markdown
Collaborator

Developed by the CAKE team, the kernels in this PR outperform the relevant SOTA baselines at the time of merge on the vast majority of performance benchmark shapes, matching them otherwise.

TP8 packed QKV uses four CTA rows and a serialized fused peer-copy helper on SM103 for BF16 local M=512, K=8192, N=1280. SM100 supports the same packed-QKV shape through the generic peer-copy route with grid y=1. The host allocates 2 + 4 * world_size uint32 mailbox flags, matching the double-buffered synchronization protocol. Architecture-specific manifests and dispatch select the appropriate exported kernel.

Published head: 92b9e398676f7229883a21b59e306794e73dd93d (tree f9a5cf651d35fd49f6c58766766722d71e20f72a). This successor applies the CI-requested line wrapping in four Python files; all four Python ASTs and all generated/native payloads are unchanged from tested commit f037814544ab714be40ffc02fcd3e935b4ca6c6f (tree ff92f6abb73d12096010c63705e884642193f559). Hardware measurements below retain their actual f037 official-wheel identity; they are not new executions at the formatting successor.

Source and export identity

The earlier source performance runs below executed a retained source projection directly. After the SM100 packed-QKV dispatch metadata update, canonical export regenerated both architecture outputs and verified that their generated CUDA bytes were unchanged. The final source integration adds documentation and contract aggregation metadata to that export-proven source. The public candidate contains those exact exports plus its host allocation and architecture-dispatch updates. These source timings are distinct from final-wheel public-API timings.

Both packaged CUDA exports are 121254 bytes, SHA256 ef39c0d6ddfa1e038fc50bdf6137e8df872142faaadfe31a3a1fbe271e11d94d. Manifest SHA256: SM100 ef58c02589f20ced592af02e908bb01d7255b0c77a051a71ea1118559b6a9b55; SM103 b70379dcb14659d8102e00776ff6a126cd5996a448d363489452444d856ab37d. The recorded public correctness and sanitizer payload had exactly tree ff92, installed as a candidate overlay using the preceding package metadata. The subsequent official wheel was built from actual commit f037; no preceding run is relabeled as an official-wheel execution.

Latest prepared-source comparison: eight B200 and eight B300 GPUs

These executions use standalone source tree f7d1b20d76cbb673957fb9fc76a68771261fce83 and canonical timer SHA256 2a7c5f4789acecedca349c6239ed1f07aff2956146738db1bb240265a72dbdca. They are fresh executions of that tree; the retained earlier-source measurements and sanitizer runs below keep their original identities. All 209 inputs to the retained native module match this source tree, so that unchanged module was reused. The source tree label does not identify a new FlashInfer package build.

For source versus cuTile, both architectures use identical BF16 TP8 workloads, K8192/N2048, global M=8×local M, standard-normal inputs, seed 52, the same cuTile baseline commit c656c1196286e99d68656f14118a535dc5248c80 (wheel SHA256 df00e8f880c50d0941d4266869180b7ec2685e09f371c1469e4e098ebec83676), canonical cold-L2 CUPTI, paired ABBA and maximum-rank medians. Each cuTile/source ratio must be strictly greater than 1.0×.

Local M B200 source ms B200 cuTile ms B200 cuTile/source B300 source ms B300 cuTile ms B300 cuTile/source
65536 26.1815210 31.0165200 1.184672× 25.6586900 30.7232445 1.197382×
32768 13.1660490 15.9422045 1.210857× 13.0759460 15.9366535 1.218776×
16384 6.1277045 8.1512165 1.330223× 6.1721185 8.1296240 1.317153×
8192 3.1344920 4.7942985 1.529530× 3.2382530 4.7614855 1.470387×
4096 1.6085275 2.6841110 1.668676× 1.5362035 2.5514250 1.660864×
2048 0.9356660 1.1649305 1.245028× 0.8405770 1.0266560 1.221371×
1024 0.7077295 1.0610580 1.499242× 0.7093420 0.9972320 1.405855×

All seven comparative rows pass on each architecture. B200 geometric mean is 1.370700×, minimum 1.184672×; B300 geometric mean is 1.347343×, minimum 1.197382×. Each architecture also passes correctness-only FP16 M16384, BF16 M19456 and packed-QKV BF16 M512/K8192/N1280, for ten correctness rows and eight rank completions. The seven N2048 rows use grid y=1; their speedups do not establish N1280 performance.

B200 verifies 80 candidate rank correctness checks, 56 baseline preflight checks, complete routes with zero external fallbacks, and 84,576 reportable rank/arm CUPTI samples. Measured-row runtime is 182.087188 s, registered command 217.796813 s, managed runtime 292.065210 s, and submission-to-terminal physical turnaround 292.914876 s. Result SHA256: fc4d6d4cf6db7aa9b54d1fdc14531f07a9972a83c70ad50500408709261d7863; original registered-report SHA256: e014fca416de3607fe456f36b2408d0f5be2f1b8b05453d667f2caba99259635.

B300 verifies 80 candidate rank correctness checks, 56 baseline preflight checks, complete routes with zero external fallbacks, and 83,168 validated rank/arm CUPTI samples. Both architectures use correctness tolerance atol=rtol=0.01. B300 measured-row runtime is 272.771914 s, registered command 311.016003 s, managed runtime 397.258994 s, and submission-to-terminal physical turnaround 398.086380 s. Result SHA256: 44240a844bd47389a393ace7529f76b3622f1070ecb1b1f5773641d34531d76c. Benchmark and bounded receipt audits pass. Original registered-report SHA256: 5843c121e4f53591328fce9d512e8fdd6b64923cbf34de79d2dd350ff8c67bad. All 12 new registered-run artifact hashes reconcile with the saved row and snapshot receipts; all eight rank-completion hashes also match. Full raw reports and CUPTI sidecars remain preserved under their original execution identity.

Latest prepared-source versus official public API

Both comparisons use standalone source tree f7d1b20d76cbb673957fb9fc76a68771261fce83 and the unchanged official f037 wheel identified below. On each architecture, four fresh processes run source/public/public/source for BF16 TP8 local/global M512/4096, K8192/N1280. Each process has three repeats, five untimed warmups and twenty cold-L2 CUPTI samples per repeat. Each row averages the two source maximum-rank medians and the two public maximum-rank medians. Public speedup is source/public; source speedup is the reciprocal.

Architecture Repeat Source ms Official public ms Public speedup Source speedup
SM100 / B200 1 0.76273525 0.70929700 1.075340× 0.929939×
SM100 / B200 2 0.77019950 0.80827375 0.952894× 1.049434×
SM100 / B200 3 0.75148100 0.68580825 1.095760× 0.912609×
SM103 / B300 1 0.40983050 0.30802500 1.330511× 0.751591×
SM103 / B300 2 0.46655850 0.32058425 1.455338× 0.687126×
SM103 / B300 3 0.55792450 0.30716875 1.816345× 0.550556×

B200 public-speedup GM is 1.039366× (reciprocal 0.962125×); its second repeat favors source. B300 public-speedup GM is 1.520758× (reciprocal 0.657567×), with public faster in all three repeats. This source/public comparison is separate from the seven-shape source/cuTile gate.

Each architecture passes 576 correctness checks, 256 source reusable-epoch assertions, 288 source-storage checks, 288 public retained-output checks, 32 clean rank completions, M256/M512 interleaving, matching workload/timing signatures, bit-exact cross-arm outputs, exact CUPTI routes and native isolation. Each has 1,920 rank timing samples. SM100 uses the CE peer-copy route without a fused helper; SM103 uses the fused-P32 route.

Architecture Maximum-rank workload sum s Physical process sum s Batch s Managed s Submission-to-terminal s
SM100 / B200 149.244075 208.648238 232.704296 249.151044 250.150905
SM103 / B300 249.803724 314.744196 358.434339 374.092103 375.209552

B200 result SHA256: 1d3d4c634bb9446e2c6c089cea988927a1f2aef65e1b9322f3e39e3e96bb454c; original full-identity SHA256: c9dc8696abd401d3c7c309b32437083ae342cf0ce537e108bb1d9e90d36c4de8. B300 result SHA256: e478d94e63b83eb90f9e1010716d4b268453177f86adf5699604b7c57f47508e; original full-identity SHA256: b5a0feaa4aebd11f2043f0e49da82a5c8d21b61b19d57d73882c17b93a61383a. All 33 new B300 ABBA result-artifact hashes reconcile with the recorded receipts; the original raw logs remain preserved. These comparisons add no new sanitizer or serving execution claim.

Earlier source projection versus cuTile: SM100 / eight B200s

Identical BF16 TP8 workloads, K8192/N2048, global M=8×local M, against cuTile from FlashInfer c656c1196286e99d68656f14118a535dc5248c80. Timing uses canonical cold-L2 CUPTI, paired ABBA and maximum-rank medians. Every comparative row must exceed 1.0×; geometric mean is a cross-shape summary.

Local M Source ms cuTile ms cuTile/source
65536 25.8811250 30.9193570 1.194668×
32768 12.9458285 15.9008315 1.228259×
16384 6.0791230 8.0983430 1.332156×
8192 3.1621890 4.9279950 1.558413×
4096 1.5663510 2.5899545 1.653496×
2048 0.8223545 0.9905480 1.204527×
1024 0.6745915 1.0002715 1.482781×

All seven rows pass; geometric mean 1.368661×, minimum 1.194668×. Ten correctness rows, complete routes with no external fallbacks, and eight rank completions passed on actual SM100 hardware. Measured-row runtime: 272.267169 s; registered command: 378.448895 s; physical turnaround: 1080.709207 s.

Earlier source projection versus cuTile: SM103 / eight B300s

The same shapes, pinned cuTile baseline, timing method and per-shape acceptance apply.

Local M Source ms cuTile ms cuTile/source
65536 23.8101150 28.9214660 1.214671×
32768 11.9789640 14.9336400 1.246655×
16384 5.6753380 7.7276140 1.361613×
8192 3.0889980 4.5455770 1.471538×
4096 1.5660350 2.5286290 1.614670×
2048 0.7973615 0.9637620 1.208689×
1024 0.6134575 0.9047055 1.474765×

All seven rows pass; geometric mean 1.362857×, minimum 1.208689×. Ten correctness rows, complete routes with no external fallbacks, and eight rank completions passed. Measured-row runtime: 205.790277 s; registered command: 239.132412 s; physical turnaround: 286.879898 s.

Both architectures also passed correctness-only rows for FP16 M16384, BF16 non-power-of-two M19456 and packed QKV BF16 M512/K8192/N1280. The N2048 performance rows use grid y=1; the packed-QKV row does not establish a separate N1280 performance result.

Public candidate and official package

Recorded SM103 public candidate validation passed 144 correctness checks, 144 retained-output checks, 16 reusable epochs per rank, M256/M512 interleaving, exact CUPTI routes and native isolation. Both separate sanitizer runs passed 144 checks and eight clean rank exits. Synccheck reported zero errors; racecheck reported zero hazards, errors and warnings. Sanitizers cover launched kernels except NCCL internal transport kernels and run without CUPTI; routes and isolation come from the separate correctness run.

Original candidate-overlay evidence Measured maximum-rank runtime s Physical turnaround s
Correctness / routes / isolation 69.428493 186.298722
Synccheck 184.205730 248.461311
Racecheck 272.490861 363.788799

The official wheel build and complete installation verification passed for commit f037. The wheel is 36177235 bytes, SHA256 bde34afb1ce3c3fc716ec0660f9a4e2de61077909f9a36d942908494d9841be0, and its build metadata records the actual commit. Official build runtime was 1128.096168 s, total build/install/verification runtime 1248.230369 s, and physical turnaround 1309.323989 s. Target-cluster installation verified 7532 runtime members from this unchanged wheel: successful installation runtime 337.795135 s, attempt turnaround 399.492244 s, and first-resume submission through completion 2430.939085 s. Building or installing the package does not itself add GPU measurements.

Earlier source projection versus official public API: SM100

The official f037 wheel and the export-proven source projection completed a paired comparison on eight actual B200s: BF16 TP8, local/global M512/4096, K8192/N1280. Four fresh processes ran source/public/public/source, each with three repeats, five untimed warmups and twenty cold-L2 CUPTI samples per repeat. Each row averages the two source maximum-rank medians and the two public maximum-rank medians.

Repeat Source ms Official public ms Public speedup (source/public) Source speedup (public/source)
1 0.85142550 0.71838525 1.185193× 0.843744×
2 0.82035275 0.67626625 1.213062× 0.824360×
3 0.92394575 0.76292250 1.211061× 0.825722×

The public API is faster in all three repeats: public-speedup GM 1.203038×, reciprocal source-speedup GM 0.831229×. This source/public comparison is separate from the source/cuTile per-shape gate.

Four verified pass seals and 32 clean rank completions establish 576 correctness checks, 256 source reuse assertions, 288 source storage observations and 288 public retained-output checks. M256/M512 interleaving, reusable epochs, exact CUPTI routes, native isolation and bit-exact outputs pass; workload and timing signatures match across all four processes. SM100 uses the CE peer-copy route with no fused helper or fused counters. The run contains 1920 rank timing samples.

Maximum-rank workload runtime sum: 131.897628 s; physical process sum: 195.762053 s; batch runtime: 221.134788 s; managed runtime: 254.153926 s; actual submission-to-terminal physical turnaround: 258.928194 s. The tested package is official commit f037814544ab714be40ffc02fcd3e935b4ca6c6f, tree ff92f6abb73d12096010c63705e884642193f559, wheel SHA256 bde34afb1ce3c3fc716ec0660f9a4e2de61077909f9a36d942908494d9841be0. Result SHA256: af88b84fa1fa914228ee0787a3eaacc192f5b8b61f786d50e1321f31e9321f26; full original identity SHA256: 5c1b869df4813b5e89f952e2f5560137c5c387ed1be772734a1f8ec206a23310. These measurements retain their f037 execution identity and add no sanitizer claim.

Earlier source projection versus official public API: SM103

The official f037 wheel and the export-proven source projection completed a fresh paired comparison on eight B300s: BF16 TP8, local/global M512/4096, K8192/N1280. Four fresh processes ran source/public/public/source, each with three repeats, five untimed warmups and twenty cold-L2 CUPTI samples per repeat. Each row averages the two source maximum-rank medians and the two public maximum-rank medians.

Repeat Source ms Official public ms Public speedup (source/public)
1 0.69062525 0.63907325 1.080667×
2 0.63772900 0.53622525 1.189293×
3 0.66537700 0.64682550 1.028681×

The official public API is faster in all three repeats, with a public-speedup geometric mean of 1.097540×. The reciprocal source-speedup GM is 0.911128×. This compares source and public implementations separately from the source/cuTile per-shape gate.

Four completed pass seals and 32 clean rank completions establish 576 correctness checks, 256 source reuse assertions, 288 source storage observations, 288 public retained-output checks, and 1504 source zero-counter checks. M256/M512 interleaving, exact routes, native isolation and bit-exact outputs passed; fresh workload and timing signatures matched across all four processes. The run contains 1920 rank timing samples.
Maximum-rank workload runtime sum: 186.744612 s; physical process sum: 249.616150 s; managed runtime: 323.759622 s; actual submission-to-terminal physical turnaround: 324.549698 s. The package binding uses the official f037 wheel SHA256 recorded above. This comparison adds no sanitizer or serving claim.

Retained SM100 source/public sanitizer evidence

Payload / tool Correctness / clean ranks Maximum-rank workload s Physical submission-to-terminal s
Source / synccheck 144 / 8 152.446307 213.633912
Official public API / synccheck 144 / 8 169.487140 257.448675
Source / racecheck 144 / 8 239.891819 299.855235
Official public API / racecheck 144 / 8 248.137925 306.849004

Each pass covers 16 reusable epochs per rank and M256/M512 interleaving. The public runs each include 144 retained-output checks; each source sanitizer run includes 128 reuse assertions and 144 output-storage checks. Each has eight clean rank completions, zero process exit codes, and result/process/sanitizer-log hashes matching its seal. Every synccheck rank reports zero errors; every racecheck rank reports zero hazards, errors and warnings. Scope excludes NCCL internal transport kernels matching ncclDev(Func|Kernel).*; CUPTI is disabled. These sanitizer runs add no CUPTI route/isolation or performance claim; those checks come from the separate comparison above.

The public runs use the unchanged official f037 wheel recorded above. Public synccheck physical process runtime was 188.742848 s, batch 220.751068 s, managed 256.420841 s, and actual submission-to-terminal turnaround 257.448675 s. Result SHA256 c84820fd5916de10d44b9d66ca6a80801919a476cf92f01f66f214c6fa1aade5; full original identity SHA256 dba3dd572b93df8f8e9e83d6266cf8a42de27bf01edb3c4bd5568fef7aaad66c.

Source racecheck physical process was 257.969875 s, batch 281.901288 s, managed 298.919530 s; actual submission-to-terminal physical turnaround 299.855235 s. Result SHA256 c6839f70f91700990fd35d3a6e47c2707003bee5cf0a2806a2276833472f454d; original identity SHA256 a9ae1476aed557f21f966eb841874f862051d31a09a016b5a4b22ad26f0b40d7.

Public racecheck physical process was 266.040934 s, batch 289.472853 s, managed 305.892351 s; actual submission-to-terminal physical turnaround 306.849004 s. Result SHA256 dd6becefc6792612a4c46c72b1c678f435dc5da519ace98cbfb6168b7be7601d; original identity SHA256 fd21ffe643ac37463fe5929f258e167e1c411bbd951d75916cb05fac7f7060e3.

Source synccheck physical process was 169.728984 s, batch 192.660966 s, managed 212.652857 s; actual submission-to-terminal physical turnaround 213.633912 s. All eight sanitizer summaries report zero errors. Result SHA256 eba2f236b298154db1c82d3b64d92bae69ca7e268c055ffda4ec40278442384b; original identity SHA256 f54c736b4cbbff4a058e1c26c76df377da7f3221b3e1184d89f1d4f07a8d15fd. All five staged SM100 cells are complete: source/public ABBA and four separate sanitizer commands. These results retain their original source and f037 package identities.

Integration and published-head CI evidence

Evidence Status and tested identity
Official-wheel SM100 source/public-API comparison PASS: exact package identity, numerical correctness, reusable epochs, routes/isolation and comparative timings are recorded above.
SM100 separate source/public synccheck and racecheck PASS: all four separate commands, each with 144 correctness checks and eight clean rank exits; zero synccheck errors and zero racecheck hazards/errors/warnings.
Earlier real Llama-3.1-70B-Instruct true sequence-parallel TP8 serving Historical PASS on eight B200 and eight B300 GPUs using SGLang 6dc8c9445338c8fbd92a12162dd6a97673f50933 and this official wheel. Both architectures match all 96 candidate/native and 96 native/native token-ID requests; maximum candidate logprob deltas are 0.0003813989460468292 (B200) and 0.0003831770736724138 (B300), native delta 0. Exact eight-rank routes, 80 bindings/rank, 2560 QKV hits/rank/measurement, native isolation and synchronization disabled all pass.
Final published-head CI, review and merge Lint and docs pass at 92b9e398676f7229883a21b59e306794e73dd93d; PR Test attempt 2 passes all 14 jobs, including all 11 configured build/test matrix cells. Required review and merge remain pending.

PR Test run 34748467094, attempt 2 passed all 11 executed matrix cells: four CPU AOT build/import cells (x64/arm64, CUDA 12.9/13.0), five A10G JIT shards, T4 JIT, and H100 JIT. The run is associated with head 92b9e398676f7229883a21b59e306794e73dd93d; the checked-out PR merge was 55887bc291c95c78903cdbe6c404c09e59365e6f (head plus base 488ffbd9fe4455c3b4c3030cf668b5603c6698e4). H100 reported 225,137 passed, 279,817 skipped, and zero failed; its test runtime was 7,546 s, and retry-to-completion turnaround was 7,616 s. These CI lanes do not provide SM100/SM103 TP8 GPU validation. The earlier runner-acquisition failure is superseded by this successful retry. The H100 lane passed 52 affected host tests and 10 dispatch tests; its two affected TP8 GPU E2E tests were skipped. SM100/SM103 correctness, sanitizer and timing evidence retains the separate hardware execution identities above.

Current SGLang integration: true-SP TP8 serving on B200 and B300

SGLang PR36766 now has real-model true sequence-parallel TP8 correctness evidence on eight actual B200 GPUs and eight actual B300 GPUs. Both execute clean SGLang commit 827dbfbf2883c1ab507b2309d0824d94f5ae57e1, tree 7f36509440f9238f918dee650a385213871f4fc6, with the official FlashInfer f037814544ab714be40ffc02fcd3e935b4ca6c6f wheel, tree ff92f6abb73d12096010c63705e884642193f559, wheel SHA256 bde34afb1ce3c3fc716ec0660f9a4e2de61077909f9a36d942908494d9841be0. The version/checksum checks stay enabled; the original runtime receipt verifies 7,532 package files. Published FlashInfer formatting head 92b9e398676f7229883a21b59e306794e73dd93d retains its explicit AST/native-payload equivalence to the tested wheel; these runs are not relabeled as direct executions of 92b9.

The model is unquantized BF16 meta-llama/Llama-3.1-70B-Instruct, revision 1605565b47bb9346c5515c34102e054115b4f98b, with all 30 verified shared shards read-only in place. Each architecture uses one server across native-A/candidate/native-B, each one warmup and three measurements of 32 requests × 4096 input IDs × 32 output tokens, concurrency one and temperature zero. True-SP TP8 uses global/local M4096/512 and packed-QKV K8192/N1280. Synchronization is off, CUDA_LAUNCH_BLOCKING=0, diagnostics off and per-call artifacts disabled. No serving synchronization was added.

Both architectures match 96/96 candidate/native-A and 96/96 native-B/native-A token sequences. Maximum candidate logprob delta is 0.0003813989460468292 on B200 and 0.0003831770736724138 on B300, below 0.05; native stability delta is 0. Every measurement verifies all eight ranks, 80 prepared bindings per rank, 2560 prepared-QKV hits per rank, 32 full-M4096 calls and 992 decode fallbacks, with zero candidate artifacts in native arms. B200 uses CE peer-copy, main grid 20×1×1, no fused helper and start-barrier → main ordering. B300 uses 80 native-bound fused launchers, main grid 20×4×1, helper 32×7×1 and start-barrier → fused-peer-copy → completion-barrier → main ordering. Both have seven peer routes and 34 barrier slots per binding/rank respectively.

Architecture / repeat Native-A tokens/s Candidate tokens/s Native-B tokens/s Interpolated native tokens/s Candidate/native
B200 / 1 23.261132172 22.952848376 22.600646963 22.914018209 1.001694603129×
B200 / 2 23.040452260 23.045842843 22.518273869 22.766422549 1.012273350962×
B200 / 3 23.176225112 22.954360069 22.544051045 22.845061564 1.004784338387×
B300 / 1 22.712963228 22.429038915 22.394396630 22.547879035 0.994729432412×
B300 / 2 23.275056452 22.533593011 22.520824528 22.883326641 0.984716661380×
B300 / 3 23.267995862 22.503065927 23.383714663 23.328426423 0.964619967023×

Geometric-mean candidate/native ratios are 1.006240975155× on B200 and 0.981275224746× on B300. These serving throughput ratios are informational, with no minimum speedup gate. The strictly greater-than-1.0 requirement continues to apply to the seven kernel-comparison shapes above.

Architecture / repeat Native-A measured seconds Candidate measured seconds Native-B measured seconds
B200 / 1 44.021932914 44.613199339 45.308437483
B200 / 2 44.443572047 44.433176385 45.474178259
B200 / 3 44.183209089 44.610261271 45.422182463
B300 / 1 45.084385939 45.655099351 45.725724025
B300 / 2 43.995596836 45.443263286 45.469027954
B300 / 3 44.008947142 45.504910456 43.791160419
Architecture Nine measured serving phases s Three warmups s Server ready / lifetime s Launcher s Managed workload s Workload submission → terminal s
B200 402.510149250 151.937193807 255 / 859 1159.235669136 1181.555971146 1182.519948959
B300 404.678115408 154.850899250 529 / 1134 1213.902644396 1237.637359381 1238.854870796

The B300 workload completed all 12 serving phases and its parity/routes/throughput analyzers. Its original managed result remains HARNESS_FAILED because the finalizer rejected the correct scheduler Unix timestamp encoded as a string. A separate CPU finalizer copy adds only integer-epoch parsing and seals PASS in 3.235293865 s physical, leaving every existing timestamp, source/package, original harness, raw measurement and acceptance check unchanged; no GPU phase was repeated. Workload submission-to-final-seal turnaround is 1637.690051556 s. The full chain from the first B300 SG827 attempt to the final seal is 7202.305768728 s, including prior interrupted attempts and recovery. Final seal time is 2026-09-14 00:36:22.889903 UTC. This distinction preserves the original measured serving durations and the physical recovery cost.

The B300 x86_64 runtime records PyTorch 2.13.0a0+9186a08b2c.nv26.07, CUDA 13.3, sglang-kernel 0.4.6.post1, Transformers 5.8.1 and container digest 356846f927ed8f2fee14ebafd28fe45eb4bd6529a98bf0e0a8719ebd6c6aafbf. It reuses the completed model-input receipt; current mount/access/size verification took 0.093951941 s with zero new weight hashes or copies. All 55 new raw artifacts and referenced phase/rank/workload hashes match, with an independent token/logprob/routes/throughput audit taking 0.200715542 s.

Artifact SHA256
B200 completion seal 5bb10d26e1fd7d5fd1bcf3bca8c19492126734c594a83d39f3d7ff43f3117902
B200 token/logprob and raw phase hashes dc04dd00ce542ab9c037d8deee1afbe83248c9a677aeef223907cf492b8503e9
B200 eight-rank route evidence 0cb827d9601976edeff6677aed72d44149b198b741e3ac4e6213bab96bcdae7a
B200 per-repeat throughput cf4717efaed80b04333d87d90ca374714945ad8de085ed974576ac2855aee2b1
B200 exact runtime 01bd16c8652b86c401c72e13a6bed66a8040ce6b4edecd074aaa758ab8f56bc9
B300 completion seal 8a67dfd151dbddc289b5074136720807bbd60fe65b7c072dfb36b7a84c4b5a3d
B300 token/logprob and raw phase hashes 0553e45b1018c386f4a804728a53ac88ed6c58e547c6bb86e11a10c3eaee1c0e
B300 eight-rank route evidence 94b82de20f6ad835d7761aaf27a2daab07215531436e65e32ced0a1911a912f8
B300 per-repeat throughput 9667b25a9945eda19077c5b6e4bab789c647ea01b809e75af3422134d55e16a4
B300 exact runtime d366c82c6272544b60f64e57c4966bda39bdd4f9c4043aa8627d64b3550fcc13
B200 workload harness manifest b61af7389908a90b24cce139c72a213c8a5469e2ff7fd5715e83ea78ecadedf8
B300 original workload harness manifest 648448c800f443d62f817f9078c270790b551be0a316c9bec565c7f9b044faad
B300 original finalizer in workload harness e244022bfb1eddc43f5dbda2aed77c50725ffb6df2639122219ae71dfea7efb7
B300 separately repaired CPU finalizer df44b102c9903ef20200bebddeb0e62d69d203c995b643999b9040883eb2a750
Shared request input IDs c290140af3c4eaab2c2f4d98d8d5e85e89ad7360165c987c2e0377ecf6388b60
Weight LFS manifest 438a360fe1b9c1f748fdd543757f11eb677c453cc522aa61100c4a5e6dce2c6f

Original per-phase token/logprob records, workload receipts, eight rank-route records and server logs remain under their execution identities. The SGLang PR body contains the complete serving evidence and recovery linkage.

Earlier SGLang execution identities

For those SGLang 6dc8c9445338c8fbd92a12162dd6a97673f50933 executions, measured-phase runtimes are 414.401067 s (B200) and 403.592879 s (B300); submission-command-to-terminal turnaround is 1262.836780 s (B200) and 1264.615915 s (B300). Informational candidate/native throughput ratios are 0.963051/0.965179/1.009118 (B200, GM 0.978888) and 0.991565/0.992482/0.992629 (B300, GM 0.992225). The SGLang PR includes per-repeat native-A/candidate/native-B metrics and reproducibility hashes. These retained executions do not establish validation of a later SGLang integration head.

SGLang end-to-end acceptance is correctness only; throughput is reported without a minimum speedup requirement. Its integration is tracked by SGLang PR36766. Kernel comparative acceptance remains strictly greater than 1.0× for every measured shape.

Earlier source/package measurements retain their original evidence identities. All five earlier SM100 validation cells retain their original completed-run seals; the latest prepared-source measurements above have their own execution identities.

Summary by CodeRabbit

  • Performance

    • Improved packed-QKV all-gather matrix multiplication performance on supported Blackwell hardware.
    • Added optimized fused peer-copy processing for bfloat16 workloads with eight-way parallelism and 512 × 1280 dimensions.
    • Improved workload distribution and synchronization for more reliable execution.
  • Compatibility

    • Expanded prepared packed-QKV support across SM100 and SM103 GPU architectures.
    • Added support for additional supported bfloat16 configurations, including four-way parallelism with 2560 columns.

References

The shared reference corpus for CAKE kernel development includes the following projects, documentation, and existing CAKE work:

These are corpus-level references. PR-specific implementation details, changes, and benchmark references are documented above.

@coderabbitai

coderabbitai Bot commented Sep 11, 2026 •

Copy link
Copy Markdown
Contributor

Review Change StackReview Change Stack

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Advanced

Run ID: e0a927a7-81fb-4876-9090-56ed754ad9ab

📥 Commits

Reviewing files that changed from the base of the PR and between f037814 and 92b9e39.

📒 Files selected for processing (4)
  • flashinfer/comm/all_gather_matmul/all_gather_matmul.py
  • flashinfer/comm/all_gather_matmul/cake_all_gather_matmul.py
  • tests/comm/test_all_gather_matmul_cake.py
  • tests/comm/test_all_gather_matmul_cake_e2e.py
🚧 Files skipped from review as they are similar to previous changes (4)
  • flashinfer/comm/all_gather_matmul/all_gather_matmul.py
  • tests/comm/test_all_gather_matmul_cake_e2e.py
  • tests/comm/test_all_gather_matmul_cake.py
  • flashinfer/comm/all_gather_matmul/cake_all_gather_matmul.py

Included review availability: Your plan provides up to 8 included reviews per hour; 5 remain after this review.


📝 Walkthrough

Walkthrough

The change updates SM100a and SM103a kernels, adds fused peer copying for packed-QKV execution, broadens SM100/SM103 routing, updates launch manifests, and extends validation and end-to-end tests.

Changes

Packed-QKV all-gather matmul

Layer / File(s) Summary
Barrier and kernel contract updates
csrc/cake_all_gather_matmul/sm100a/cake_all_gather_matmul_kernels.cu, csrc/cake_all_gather_matmul/sm103a/cake_all_gather_matmul_kernels.cu
The kernels add FlashInferTensorMap, expanded mbarrier helpers, epoch/mailbox barriers, synchronized TMEM allocation, updated warp/lane calculation, and grid-y peer-pass partitioning.
Fused peer-copy kernel
csrc/cake_all_gather_matmul/sm100a/cake_all_gather_matmul_kernels.cu, csrc/cake_all_gather_matmul/sm103a/cake_all_gather_matmul_kernels.cu
A new kernel copies input data to peer payload buffers and signals completion through counters and readiness slots.
Launch metadata
csrc/cake_all_gather_matmul/sm100a/manifest.json, csrc/cake_all_gather_matmul/sm103a/manifest.json
The manifests register the thirteenth kernel and define fused peer-copy launch ordering, constraints, grid dimensions, and source hashes.
Host fused-path integration
flashinfer/comm/all_gather_matmul/cake_all_gather_matmul.py, flashinfer/comm/all_gather_matmul/all_gather_matmul.py
The backend validates fused buffers, binds prepared launches, replaces per-peer device copies for the SM103 TP8 profile, and accepts the expanded packed-QKV routes.
Validation and end-to-end coverage
tests/comm/test_all_gather_matmul_cake.py, tests/comm/test_all_gather_matmul_cake_e2e.py, tests/comm/test_all_gather_matmul_dispatch.py
Tests update manifest and launch-contract expectations, route validation, helper names, peer-buffer sizes, and SM100/SM103 packed-QKV execution.

Priority: ➖ Normal

Estimated code review effort: 5 (Critical) | ~90 minutes

Change: Feature · Severity of issue fixed: Low

Sequence Diagram(s)

sequenceDiagram
  participant Launcher as _PreparedPackedQkvLauncher
  participant Backend as BindPreparedPackedQkv
  participant Copy as FusedPeerCopyKernel
  participant Main as RunMain
  Launcher->>Backend: invoke prepared packed-QKV launch
  Backend->>Copy: launch fused peer copy on comm stream
  Copy-->>Backend: signal ready_target
  Backend->>Main: run completion barrier and matmul
  Main-->>Launcher: return output
Loading

Merge Risk: ⚪ Minimal · up to 92b9e

No actionable correctness, availability, or integration risk remains from the reviewed changes.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 11.70% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 94 functions across 7 files. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
Title check ✅ Passed The title clearly and concisely identifies the primary change: TP8 packed-QKV all-gather matmul support for SM100 and SM103.
Description check ✅ Passed The description provides detailed scope, implementation behavior, architecture coverage, validation results, performance data, and remaining status. It does not use the template headings or explicitly…
  • Fix all pre-merge checks with AI
✨ Finishing Touches 💡 1
🛠️ Fix failing CI checks 💡
  • Create stacked PR
  • Commit on current branch
🧪 Generate unit tests (beta)
  • Create PR with unit tests

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@yyihuang yyihuang changed the title Parallelize SM103 TP8 packed QKV peer traversal perf(cake_comm): parallelize SM103 TP8 packed QKV peer traversal Sep 11, 2026
@yyihuang
yyihuang marked this pull request as ready for review September 11, 2026 01:16
Both prepared callers already pass the current device stream. Preserve pinned-source lifetime tracking and descriptor-ready events.

Validated by 52 package tests, 128 alternating-stream numerical checks, 16 host-case checks and eight clean rank exits. Timing evidence is limited to CPU diagnostics.
Carry the generated barrier wait hint and current tensor-map declarations into both device bundles. Update source digests while preserving the existing public launch constraints and routes. GPU correctness, sanitizer, and performance validation remains pending.
Fuse the seven peer transfers and publish exact epochs, then complete the existing collective barrier before the unchanged matmul. Preserve mixed-route barrier phases and reusable buffers.

Validated source/export/package tests, TP8 correctness, retained outputs, exact routes, and comparative GPU timing. Sanitizer, regression and final serving gates remain pending.
@yyihuang yyihuang changed the title perf(cake_comm): parallelize SM103 TP8 packed QKV peer traversal Support TP8 packed-QKV all-gather matmul on SM100 and SM103 Sep 13, 2026
@yyihuang yyihuang changed the title Support TP8 packed-QKV all-gather matmul on SM100 and SM103 feat(cake_comm): support SM100 and SM103 TP8 packed-QKV all-gather matmul Sep 14, 2026
@yyihuang

Copy link
Copy Markdown
Collaborator Author

/bot run flashinfer/comm/all_gather_matmul/all_gather_matmul.py tests/comm/test_all_gather_matmul_cake_e2e.py tests/comm/test_all_gather_matmul_dispatch.py

@flashinfer-bot

Copy link
Copy Markdown
Collaborator

Invalid /bot run path(s): flashinfer/comm/all_gather_matmul/all_gather_matmul.py. Expected tests/<dir> or tests/<dir>/<file>.py.

@yyihuang

Copy link
Copy Markdown
Collaborator Author

/bot run tests/comm/test_all_gather_matmul.py tests/comm/test_all_gather_matmul_cake.py tests/comm/test_all_gather_matmul_cake_e2e.py tests/comm/test_all_gather_matmul_dispatch.py

@flashinfer-bot

Copy link
Copy Markdown
Collaborator

GitLab MR !1511 has been created, and the CI pipeline #67715567 is currently running. I'll report back once the pipeline job completes.

@flashinfer-bot

Copy link
Copy Markdown
Collaborator

[SUCCESS] Pipeline #67715567: 18/19 executed test jobs passed

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants