Repository navigation
perf(cake_dense_projection_gemm): 212-SM sm_107a router weight-gradient windows run the head/tail program with the A parts in tensor memory - #8517
Conversation
…'s TMEM head/tail program for sm_107a The head/tail program with the BF16 parts of A in tensor memory (cake_dense_projection_gemm_d384571ea88c28e40e31) serves sm_107a too: its record's arches gain sm_107a, KERNELS["sm_107a"] names its template and build.json lists its sm_107a library. The kernel unit is the current render of the program (with the sm_107a branch of its B shared-memory descriptor mask); every other generated file is unchanged.
…EM head/tail program on 212-SM sm_107a parts On sm_107a parts with exactly 212 SMs the weight gradient's head/tail windows (256 experts, 16172 to 16384 tokens) run the head/tail program with the BF16 parts of A in tensor memory (router_fp32_gemm_ts_r6p6_atomic_nn_t_ht_c2_nreg) instead of router_fp32_gemm_atomic_nn_t_ht_c3_nreg: - ROUTER_HEAD_TAIL_TEMPLATES_BY_SM maps (arch, sm_count) to the head/tail program of the parts of that SM count; router_head_tail_template(arch, sm_count) reads it before ROUTER_HEAD_TAIL_TEMPLATES, and plan_atomic_router passes the device's SM count. 208-SM and 216-SM sm_107a parts keep the three-chunk-buffer program. - The 212-SM table of ROUTER_RULES_BY_SM gains hidden 6144 "H0.5x4" at 16172 to 16384 tokens, a window that holds only with the tensor-memory program. Hidden 4096 keeps the deterministic plan. Tests: the planner, partition and per-SM-count template checks cover the 208, 212 and 216-SM tables; the GPU window test runs the program in the 212-SM windows. README: the 212-SM rows and their measurements.
|
No actionable comments were generated in the recent review. Walkthrough
Merge Risk:
|
|
@flashinfer-bot run tests/experimental/test_cake_dense_projection_gemm.py tests/experimental/test_cake_dense_projection_det_gemm_names.py tests/experimental/test_cake_dense_projection_det_gemm.py |
…amples cover The README's 212-SM clock note said the 200 ms nvidia-smi samples were taken while the head/tail window rows ran. Only 6 of the 12 runs at hidden 5120 and 6 of the 10 at hidden 6144 (those on two of the four parts) were sampled; the 17-47 % share of clock event reason 0x4 is of the non-idle samples. Say so. No code, test or plan change.
|
@flashinfer-bot run tests/experimental/test_cake_dense_projection_gemm.py tests/experimental/test_cake_dense_projection_det_gemm_names.py tests/experimental/test_cake_dense_projection_det_gemm.py |
|
@flashinfer-bot run tests/experimental/test_cake_dense_projection_gemm.py tests/experimental/test_cake_dense_projection_det_gemm_names.py tests/experimental/test_cake_dense_projection_det_gemm.py |
…ad-windows Conflicts with the sm_107a 212-SM router weight-gradient windows of flashinfer-ai#8517, resolved keeping both: - README.md: the router window table footnote keeps the 212-SM clock samples and the GB300 hidden 4096 / 6144 clock samples, and the GB300 bf16x9 ratio x1.500-1.640. - tests/experimental/test_cake_dense_projection_gemm.py: the head/tail program test checks the sm_103a 152-SM windows (hidden 4096, 5120, 6144) and the sm_107a 212-SM windows (hidden 5120, 6144) in one window table; the GPU test of the tensor-memory head/tail program takes the union of both parameter lists.
The
flashinfer.experimental.dense_projection_gemmfamily has three kinds of GEMM. They are BF16-input dense projection GEMMs, BF16-input batched projection GEMMs, and an FP32 router GEMM (256 experts, BF16x3 emulation). Each kind has forward, input-gradient and weight-gradient programs.deterministic=True(the default) callsdense_projection_det_gemm.deterministic=Falseruns this package's split-K / stream-K programs. This PR changes one variant: the router weight gradient (FP32 out, 256 experts) onsm_107aparts with exactly 212 SMs, at 16172 to 16384 tokens.At hidden 5120, this window already exists on FlashInfer main with plan
"H0.5x4"(106 head and 424 tail units). In that window,router_fp32_gemm_ts_r6p6_atomic_nn_t_ht_c2_nregreplacesrouter_fp32_gemm_atomic_nn_t_ht_c3_nreg. The new program is the head/tail program that already ships forsm_103a. It keeps the three BF16 parts ofAin tensor memory. The partition, the products and their accumulation order are unchanged.At hidden 6144, this PR adds a new window with plan
"H0.5x4". That window runs only the tensor-memory program. On FlashInfer main, this row runs the deterministic program at 212 SMs.ROUTER_HEAD_TAIL_TEMPLATES_BY_SMselects the program by SM count. 208-SM and 216-SMsm_107aparts keep their programs and plans.sm_100aandsm_103aalso keep their programs and plans. Hidden 4096 keeps the deterministic plan. The forward and input gradient are unchanged.These results are on R200, against FlashInfer main. Each range below covers all runs. At hidden 5120:
At hidden 6144:
The forward runs the same program in every Cake arm. Its previous/new ratio is 0.979-1.008 at hidden 5120 and 0.994-1.017 at hidden 6144.
An A/A control runs the new program against a second prepared copy of itself. It ran in 6 of the runs at each hidden size. On the weight gradient, it spans 0.994-1.009. On bwd, it spans 1.001-1.030 at hidden 5120 and 1.011-1.037 at hidden 6144. On 2fwd+bwd, it spans 0.994-1.010 at hidden 5120 and 0.941-1.009 at hidden 6144.
The weight-gradient and bwd gains over det Cake are larger than the A/A spread. At hidden 5120 the previous Cake arm is also a separately prepared copy, and the A/A control shows that such a copy can be up to 3 % slower on bwd. Against the second copy of the new program, previous/new bwd is 1.007-1.038 (in the 6 A/A runs). So the hidden-5120 bwd previous/new ratios in the table (1.045-1.051) may be up to about 3 % too high. The hidden-5120 2fwd+bwd gain over det Cake is also larger than the A/A spread. The hidden-5120 2fwd+bwd previous/new ratio (1.010-1.031) is about as large as its A/A spread (0.994-1.010). At hidden 6144, the 2fwd+bwd gain over both Cake baselines (1.015-1.045) is smaller than its A/A spread (0.941-1.009). That row is parity with no regression, not a measured gain.
sm_107apart. Ops are wgrad (router weight gradient, the changed op), bwd (input gradient + weight gradient) and 2fwd+bwd (checkpointed step: 2 x forward + bwd). Modes are eager (calls of the prepared launches) and graph (CUDA graph replays).deterministic=True(dense_projection_det_gemm).deterministic=Falseon FlashInfer main. At hidden 5120, it runsrouter_fp32_gemm_atomic_nn_t_ht_c3_nregat the same plan. At hidden 6144, it runs the deterministic program, which is main's plan for that row.cublasSgemmwithCUBLAS_FP32_EMULATED_BF16X9_MATH.sm_107a, 212 SMs, driver 620.43, CUDA 13.4 (nvcc 13.4.92), PyTorch 2.14.1+cu130, cuBLAS 13.8.0.nvidia-smiwas sampled every 200 ms only in 6 of the 12 hidden-5120 runs and 6 of the 10 hidden-6144 runs (the runs on two of the four GPUs). It reported clock event reason 0x4 in 17-47 % of the non-idle samples.The other variants of the family keep their programs, and the family ships on
sm_100a,sm_103aandsm_107a. These other variants are the projection GEMMs, the router forward and input gradient, and the router weight gradient onsm_100a,sm_103a, 208-SM and 216-SMsm_107aparts, at hidden 4096 and outside these windows.Summary by CodeRabbit
sm_107adevices.sm_107adevices.