diff --git a/README.md b/README.md index b611234d6..87cb1d41d 100644 --- a/README.md +++ b/README.md @@ -66,6 +66,25 @@ Contributor credits for these OSS CuTe DSL kernels are listed in [Acknowledgemen DSv3 SDPA Benchmark on GB300 (only cuDNN)

+## New OSS Linear Attention Kernels + +#### GDN Forward and Bprop (GB300) +

+ GDN Linear Attention Benchmark on GB300 +

+ +#### KDA Forward and Bprop (GB300) + +

+ KDA Linear Attention Benchmark on GB300 +

+ +#### GDN-2 Forward and Bprop (GB300) + +

+ GDN-2 Linear Attention Benchmark on GB300 +

+ ## Key Features * **Unified Graph API:** Create reusable, persistent `cudnn_frontend::graph::Graph` objects to describe complex subgraphs. diff --git a/benchmark/linear_attention/README.md b/benchmark/linear_attention/README.md index 7c2dac8c2..a37977a93 100644 --- a/benchmark/linear_attention/README.md +++ b/benchmark/linear_attention/README.md @@ -8,6 +8,8 @@ This directory contains benchmarking tools for linear attention operations (GDN/ - `Dockerfile` - Docker container setup for running benchmarks - `benchmark_single_linear_attention.py` - Single linear attention benchmark script +- `plot_results.py` - Renders the charts under `results/` from a sweep CSV +- `results///` - Dated sweep CSVs and the charts rendered from them ## Quick Start @@ -92,4 +94,24 @@ The `kda` and `gdn2` variants fuse q/k L2 normalization in-kernel on every backe | `flash_qla` | FlashQLA (TileLang fused GDN kernels, `gdn` variant only) | | `flash_kda` | FlashKDA (`kda` forward variant only) | -The cuDNN backend routes through the pygraph engines: FROST (Cutlass DSL) on SM100-class devices, the cuTile engines elsewhere. \ No newline at end of file +The cuDNN backend routes through the pygraph engines: FROST (Cutlass DSL) on SM100-class devices, the cuTile engines elsewhere. + +## Results + +Forward and backward TFLOPS at batch 4 over the sequence-length sweep, one +chart per variant, rendered by `plot_results.py` from the dated CSVs under +`results///`. The `cudnn (state on)` bars dump the per-chunk +state-checkpoint series in the forward pass and reuse it in the backward +pass. Runs were captured on GB200 and GB300 (GB300 results shown below). + +### GB300 - GDN +![GDN on GB300](results/gdn/gb300/gdn_fixed_batch_flops.png) +- `batch=4; num_q_heads=64; num_kv_heads=64; head_dim=128; seqlen 2048-32768; bf16` + +### GB300 - KDA +![KDA on GB300](results/kda/gb300/kda_fixed_batch_flops.png) +- `batch=4; num_q_heads=64; num_kv_heads=64; head_dim=128; seqlen 2048-32768; bf16` + +### GB300 - GDN-2 +![GDN-2 on GB300](results/gdn2/gb300/gdn2_fixed_batch_flops.png) +- `batch=4; num_q_heads=64; num_kv_heads=64; head_dim=128; seqlen 2048-32768; bf16` \ No newline at end of file diff --git a/benchmark/linear_attention/results/gdn/gb200/gdn_20260814.csv b/benchmark/linear_attention/results/gdn/gb200/gdn_20260814.csv deleted file mode 100644 index b95b666f2..000000000 --- a/benchmark/linear_attention/results/gdn/gb200/gdn_20260814.csv +++ /dev/null @@ -1,41 +0,0 @@ -case_tag,backend,variant,batch_size,seqlen,num_q_heads,num_kv_heads,head_dim,fwd_ms,bwd_ms,fwd_tflops,bwd_tflops,max_diff,num_iters,fwd_bw,bwd_bw -gdn_h64,cudnn,gdn,4,2048,64,64,128,0.176,0.661,440,351,0.000000,20,3.08,2.66 -gdn_hon,cudnn_state_on,gdn,4,2048,64,64,128,0.227,0.522,340,444,0.000000,20,2.38,3.37 -gdn_h64,fla,gdn,4,2048,64,64,128,0.756,2.323,102,100,0.000000,20,0.72,0.76 -gdn_h64,flash_qla,gdn,4,2048,64,64,128,0.272,0.726,284,320,0.000000,20,1.99,2.42 -gdn_h64,cudnn,gdn,4,4096,64,64,128,0.305,1.271,507,365,0.000000,20,3.55,2.77 -gdn_hon,cudnn_state_on,gdn,4,4096,64,64,128,0.409,0.998,378,465,0.000000,20,2.65,3.52 -gdn_h64,fla,gdn,4,4096,64,64,128,1.478,4.596,105,101,0.000000,20,0.73,0.76 -gdn_h64,flash_qla,gdn,4,4096,64,64,128,0.513,1.409,301,329,0.000000,20,2.11,2.49 -gdn_h64,cudnn,gdn,4,8192,64,64,128,0.567,2.490,545,373,0.000000,20,3.82,2.82 -gdn_hon,cudnn_state_on,gdn,4,8192,64,64,128,0.789,1.953,392,475,0.000000,20,2.74,3.60 -gdn_h64,fla,gdn,4,8192,64,64,128,2.927,9.198,106,101,0.000000,20,0.74,0.76 -gdn_h64,flash_qla,gdn,4,8192,64,64,128,0.993,2.757,312,337,0.000000,20,2.18,2.55 -gdn_h64,cudnn,gdn,4,16384,64,64,128,1.062,4.930,582,376,0.000000,20,4.07,2.85 -gdn_hon,cudnn_state_on,gdn,4,16384,64,64,128,1.575,3.866,393,480,0.000000,20,2.75,3.64 -gdn_h64,fla,gdn,4,16384,64,64,128,5.842,18.483,106,100,0.000000,20,0.74,0.76 -gdn_h64,flash_qla,gdn,4,16384,64,64,128,1.961,5.502,315,337,0.000000,20,2.21,2.56 -gdn_h64,cudnn,gdn,4,32768,64,64,128,2.086,9.814,593,378,0.000000,20,4.15,2.87 -gdn_hon,cudnn_state_on,gdn,4,32768,64,64,128,3.256,7.689,380,483,0.000000,20,2.66,3.66 -gdn_h64,fla,gdn,4,32768,64,64,128,11.715,37.419,106,99,0.000000,20,0.74,0.75 -gdn_h64,flash_qla,gdn,4,32768,64,64,128,3.888,10.967,318,338,0.000000,20,2.23,2.56 -gdn_h64,cudnn,gdn,1,8192,64,64,128,0.178,0.667,434,347,0.000000,20,3.04,2.63 -gdn_hon,cudnn_state_on,gdn,1,8192,64,64,128,0.226,0.525,342,442,0.000000,20,2.39,3.35 -gdn_h64,fla,gdn,1,8192,64,64,128,0.821,2.482,94,93,0.000000,20,0.66,0.71 -gdn_h64,flash_qla,gdn,1,8192,64,64,128,0.342,1.313,226,177,0.000000,20,1.58,1.34 -gdn_h64,cudnn,gdn,2,8192,64,64,128,0.284,1.244,544,373,0.000000,20,3.81,2.82 -gdn_hon,cudnn_state_on,gdn,2,8192,64,64,128,0.387,0.973,400,477,0.000000,20,2.80,3.61 -gdn_h64,fla,gdn,2,8192,64,64,128,1.471,4.840,105,96,0.000000,20,0.74,0.73 -gdn_h64,flash_qla,gdn,2,8192,64,64,128,0.503,1.434,307,324,0.000000,20,2.15,2.45 -gdn_h64,cudnn,gdn,4,8192,64,64,128,0.565,2.489,547,373,0.000000,20,3.83,2.82 -gdn_hon,cudnn_state_on,gdn,4,8192,64,64,128,0.786,1.951,393,475,0.000000,20,2.75,3.60 -gdn_h64,fla,gdn,4,8192,64,64,128,2.927,9.200,106,101,0.000000,20,0.74,0.76 -gdn_h64,flash_qla,gdn,4,8192,64,64,128,0.992,2.769,312,335,0.000000,20,2.18,2.54 -gdn_h64,cudnn,gdn,8,8192,64,64,128,1.057,4.978,585,373,0.000000,20,4.10,2.82 -gdn_hon,cudnn_state_on,gdn,8,8192,64,64,128,1.555,3.876,398,479,0.000000,20,2.78,3.63 -gdn_h64,fla,gdn,8,8192,64,64,128,5.601,17.705,110,105,0.000000,20,0.77,0.79 -gdn_h64,flash_qla,gdn,8,8192,64,64,128,1.963,5.390,315,344,0.000000,20,2.20,2.61 -gdn_h64,cudnn,gdn,16,8192,64,64,128,1.891,9.164,654,405,0.000000,20,4.58,3.07 -gdn_hon,cudnn_state_on,gdn,16,8192,64,64,128,2.905,7.242,426,512,0.000000,20,2.98,3.88 -gdn_h64,fla,gdn,16,8192,64,64,128,11.050,35.154,112,106,0.000000,20,0.78,0.80 -gdn_h64,flash_qla,gdn,16,8192,64,64,128,3.622,9.670,342,384,0.000000,20,2.39,2.91 diff --git a/benchmark/linear_attention/results/gdn/gb200/gdn_20260824.csv b/benchmark/linear_attention/results/gdn/gb200/gdn_20260824.csv new file mode 100644 index 000000000..3a5047029 --- /dev/null +++ b/benchmark/linear_attention/results/gdn/gb200/gdn_20260824.csv @@ -0,0 +1,41 @@ +case_tag,backend,variant,batch_size,seqlen,num_q_heads,num_kv_heads,head_dim,fwd_ms,bwd_ms,fwd_tflops,bwd_tflops,max_diff,num_iters,fwd_bw,bwd_bw +gdn_h64,cudnn,gdn,4,2048,64,64,128,0.166,0.596,467,389,0.000000,20,3.27,2.95 +gdn_h64,fla,gdn,4,2048,64,64,128,0.755,2.341,102,99,0.000000,20,0.72,0.75 +gdn_h64,flash_qla,gdn,4,2048,64,64,128,0.271,0.732,285,317,0.000000,20,1.99,2.40 +gdn_h64,cudnn,gdn,4,4096,64,64,128,0.288,1.130,537,410,0.000000,20,3.76,3.11 +gdn_h64,fla,gdn,4,4096,64,64,128,1.477,4.637,105,100,0.000000,20,0.73,0.76 +gdn_h64,flash_qla,gdn,4,4096,64,64,128,0.511,1.406,303,330,0.000000,20,2.12,2.50 +gdn_h64,cudnn,gdn,4,8192,64,64,128,0.532,2.186,581,424,0.000000,20,4.06,3.22 +gdn_h64,fla,gdn,4,8192,64,64,128,2.930,9.304,106,100,0.000000,20,0.74,0.76 +gdn_h64,flash_qla,gdn,4,8192,64,64,128,0.995,2.805,311,331,0.000000,20,2.18,2.51 +gdn_h64,cudnn,gdn,4,16384,64,64,128,1.014,4.295,610,432,0.000000,20,4.27,3.27 +gdn_h64,fla,gdn,4,16384,64,64,128,5.844,18.691,106,99,0.000000,20,0.74,0.75 +gdn_h64,flash_qla,gdn,4,16384,64,64,128,1.967,5.603,314,331,0.000000,20,2.20,2.51 +gdn_h64,cudnn,gdn,4,32768,64,64,128,1.966,8.499,629,437,0.000000,20,4.40,3.31 +gdn_h64,fla,gdn,4,32768,64,64,128,11.720,37.833,106,98,0.000000,20,0.74,0.74 +gdn_h64,flash_qla,gdn,4,32768,64,64,128,3.893,11.070,318,335,0.000000,20,2.22,2.54 +gdn_h64,cudnn,gdn,1,8192,64,64,128,0.169,0.557,457,417,0.000000,20,3.20,3.16 +gdn_h64,fla,gdn,1,8192,64,64,128,0.822,2.508,94,92,0.000000,20,0.66,0.70 +gdn_h64,flash_qla,gdn,1,8192,64,64,128,0.349,1.335,222,174,0.000000,20,1.55,1.32 +gdn_h64,cudnn,gdn,2,8192,64,64,128,0.273,1.024,566,453,0.000000,20,3.96,3.43 +gdn_h64,fla,gdn,2,8192,64,64,128,1.481,4.881,104,95,0.000000,20,0.73,0.72 +gdn_h64,flash_qla,gdn,2,8192,64,64,128,0.502,1.433,308,324,0.000000,20,2.16,2.45 +gdn_h64,cudnn,gdn,4,8192,64,64,128,0.533,2.179,580,426,0.000000,20,4.06,3.23 +gdn_h64,fla,gdn,4,8192,64,64,128,2.928,9.301,106,100,0.000000,20,0.74,0.76 +gdn_h64,flash_qla,gdn,4,8192,64,64,128,0.994,2.786,311,333,0.000000,20,2.18,2.52 +gdn_h64,cudnn,gdn,8,8192,64,64,128,1.010,4.493,612,413,0.000000,20,4.29,3.13 +gdn_h64,fla,gdn,8,8192,64,64,128,5.660,17.921,109,104,0.000000,20,0.76,0.78 +gdn_h64,flash_qla,gdn,8,8192,64,64,128,1.971,5.437,314,341,0.000000,20,2.20,2.59 +gdn_h64,cudnn,gdn,16,8192,64,64,128,1.755,8.112,705,457,0.000000,20,4.93,3.47 +gdn_h64,fla,gdn,16,8192,64,64,128,11.150,35.579,111,104,0.000000,20,0.78,0.79 +gdn_h64,flash_qla,gdn,16,8192,64,64,128,3.627,9.704,341,382,0.000000,20,2.39,2.90 +gdn_hon,cudnn_state_on,gdn,4,2048,64,64,128,0.222,0.473,349,490,0.000000,20,2.44,3.71 +gdn_hon,cudnn_state_on,gdn,4,4096,64,64,128,0.401,0.887,386,523,0.000000,20,2.70,3.96 +gdn_hon,cudnn_state_on,gdn,4,8192,64,64,128,0.767,1.697,403,547,0.000000,20,2.82,4.14 +gdn_hon,cudnn_state_on,gdn,4,16384,64,64,128,1.500,3.285,412,565,0.000000,20,2.89,4.28 +gdn_hon,cudnn_state_on,gdn,4,32768,64,64,128,2.970,6.472,416,573,0.000000,20,2.91,4.34 +gdn_hon,cudnn_state_on,gdn,1,8192,64,64,128,0.223,0.425,347,546,0.000000,20,2.43,4.14 +gdn_hon,cudnn_state_on,gdn,2,8192,64,64,128,0.384,0.771,403,602,0.000000,20,2.82,4.56 +gdn_hon,cudnn_state_on,gdn,4,8192,64,64,128,0.768,1.693,403,548,0.000000,20,2.82,4.15 +gdn_hon,cudnn_state_on,gdn,8,8192,64,64,128,1.527,3.475,405,534,0.000000,20,2.83,4.05 +gdn_hon,cudnn_state_on,gdn,16,8192,64,64,128,2.745,6.285,451,590,0.000000,20,3.15,4.47 diff --git a/benchmark/linear_attention/results/gdn/gb200/gdn_fixed_batch_bw.png b/benchmark/linear_attention/results/gdn/gb200/gdn_fixed_batch_bw.png index bd44211e5..42391cc4a 100644 Binary files a/benchmark/linear_attention/results/gdn/gb200/gdn_fixed_batch_bw.png and b/benchmark/linear_attention/results/gdn/gb200/gdn_fixed_batch_bw.png differ diff --git a/benchmark/linear_attention/results/gdn/gb200/gdn_fixed_batch_flops.png b/benchmark/linear_attention/results/gdn/gb200/gdn_fixed_batch_flops.png index 5f4a442c7..7160b2ce7 100644 Binary files a/benchmark/linear_attention/results/gdn/gb200/gdn_fixed_batch_flops.png and b/benchmark/linear_attention/results/gdn/gb200/gdn_fixed_batch_flops.png differ diff --git a/benchmark/linear_attention/results/gdn/gb200/gdn_fixed_seq_bw.png b/benchmark/linear_attention/results/gdn/gb200/gdn_fixed_seq_bw.png index 4b834a1b5..7ba9d117b 100644 Binary files a/benchmark/linear_attention/results/gdn/gb200/gdn_fixed_seq_bw.png and b/benchmark/linear_attention/results/gdn/gb200/gdn_fixed_seq_bw.png differ diff --git a/benchmark/linear_attention/results/gdn/gb200/gdn_fixed_seq_flops.png b/benchmark/linear_attention/results/gdn/gb200/gdn_fixed_seq_flops.png index 4368b9288..ee0e7cdef 100644 Binary files a/benchmark/linear_attention/results/gdn/gb200/gdn_fixed_seq_flops.png and b/benchmark/linear_attention/results/gdn/gb200/gdn_fixed_seq_flops.png differ diff --git a/benchmark/linear_attention/results/gdn/gb300/gdn_20260814.csv b/benchmark/linear_attention/results/gdn/gb300/gdn_20260814.csv deleted file mode 100644 index a784e4c45..000000000 --- a/benchmark/linear_attention/results/gdn/gb300/gdn_20260814.csv +++ /dev/null @@ -1,41 +0,0 @@ -case_tag,backend,variant,batch_size,seqlen,num_q_heads,num_kv_heads,head_dim,fwd_ms,bwd_ms,fwd_tflops,bwd_tflops,max_diff,num_iters,fwd_bw,bwd_bw -gdn_h64,cudnn,gdn,4,2048,64,64,128,0.165,0.660,469,352,0.000000,20,3.28,2.66 -gdn_hon,cudnn_state_on,gdn,4,2048,64,64,128,0.217,0.501,356,463,0.000000,20,2.49,3.51 -gdn_h64,fla,gdn,4,2048,64,64,128,0.740,2.284,104,102,0.000000,20,0.73,0.77 -gdn_h64,flash_qla,gdn,4,2048,64,64,128,0.273,0.729,283,318,0.000000,20,1.98,2.41 -gdn_h64,cudnn,gdn,4,4096,64,64,128,0.289,1.219,534,381,0.000000,20,3.74,2.88 -gdn_hon,cudnn_state_on,gdn,4,4096,64,64,128,0.396,0.954,390,486,0.000000,20,2.73,3.68 -gdn_h64,fla,gdn,4,4096,64,64,128,1.447,4.530,107,102,0.000000,20,0.75,0.78 -gdn_h64,flash_qla,gdn,4,4096,64,64,128,0.516,1.406,299,330,0.000000,20,2.10,2.50 -gdn_h64,cudnn,gdn,4,8192,64,64,128,0.546,2.389,566,388,0.000000,20,3.96,2.94 -gdn_hon,cudnn_state_on,gdn,4,8192,64,64,128,0.776,1.864,398,498,0.000000,20,2.79,3.77 -gdn_h64,fla,gdn,4,8192,64,64,128,2.865,9.067,108,102,0.000000,20,0.76,0.78 -gdn_h64,flash_qla,gdn,4,8192,64,64,128,1.000,2.766,309,335,0.000000,20,2.16,2.54 -gdn_h64,cudnn,gdn,4,16384,64,64,128,1.049,4.715,590,394,0.000000,20,4.13,2.98 -gdn_hon,cudnn_state_on,gdn,4,16384,64,64,128,1.557,3.675,397,505,0.000000,20,2.78,3.83 -gdn_h64,fla,gdn,4,16384,64,64,128,5.712,18.208,108,102,0.000000,20,0.76,0.77 -gdn_h64,flash_qla,gdn,4,16384,64,64,128,1.978,5.503,313,337,0.000000,20,2.19,2.55 -gdn_h64,cudnn,gdn,4,32768,64,64,128,2.073,9.381,597,396,0.000000,20,4.18,3.00 -gdn_hon,cudnn_state_on,gdn,4,32768,64,64,128,3.187,7.294,388,509,0.000000,20,2.72,3.85 -gdn_h64,fla,gdn,4,32768,64,64,128,11.451,36.798,108,101,0.000000,20,0.76,0.76 -gdn_h64,flash_qla,gdn,4,32768,64,64,128,3.918,10.903,316,340,0.000000,20,2.21,2.58 -gdn_h64,cudnn,gdn,1,8192,64,64,128,0.171,0.620,453,374,0.000000,20,3.17,2.83 -gdn_hon,cudnn_state_on,gdn,1,8192,64,64,128,0.221,0.480,350,483,0.000000,20,2.45,3.66 -gdn_h64,fla,gdn,1,8192,64,64,128,0.798,2.443,97,95,0.000000,20,0.68,0.72 -gdn_h64,flash_qla,gdn,1,8192,64,64,128,0.370,1.321,209,176,0.000000,20,1.46,1.33 -gdn_h64,cudnn,gdn,2,8192,64,64,128,0.278,1.146,556,405,0.000000,20,3.89,3.07 -gdn_hon,cudnn_state_on,gdn,2,8192,64,64,128,0.379,0.880,408,527,0.000000,20,2.85,4.00 -gdn_h64,fla,gdn,2,8192,64,64,128,1.447,4.768,107,97,0.000000,20,0.75,0.74 -gdn_h64,flash_qla,gdn,2,8192,64,64,128,0.505,1.418,306,327,0.000000,20,2.14,2.48 -gdn_h64,cudnn,gdn,4,8192,64,64,128,0.546,2.431,566,382,0.000000,20,3.96,2.89 -gdn_hon,cudnn_state_on,gdn,4,8192,64,64,128,0.778,1.865,398,497,0.000000,20,2.78,3.77 -gdn_h64,fla,gdn,4,8192,64,64,128,2.864,9.062,108,102,0.000000,20,0.76,0.78 -gdn_h64,flash_qla,gdn,4,8192,64,64,128,1.003,2.760,308,336,0.000000,20,2.16,2.55 -gdn_h64,cudnn,gdn,8,8192,64,64,128,1.042,4.935,593,376,0.000000,20,4.15,2.85 -gdn_hon,cudnn_state_on,gdn,8,8192,64,64,128,1.533,3.861,403,481,0.000000,20,2.82,3.64 -gdn_h64,fla,gdn,8,8192,64,64,128,5.505,17.444,112,106,0.000000,20,0.79,0.81 -gdn_h64,flash_qla,gdn,8,8192,64,64,128,1.977,5.358,313,346,0.000000,20,2.19,2.62 -gdn_h64,cudnn,gdn,16,8192,64,64,128,1.910,9.012,648,412,0.000000,20,4.53,3.12 -gdn_hon,cudnn_state_on,gdn,16,8192,64,64,128,2.866,7.139,432,520,0.000000,20,3.02,3.94 -gdn_h64,fla,gdn,16,8192,64,64,128,10.881,34.661,114,107,0.000000,20,0.80,0.81 -gdn_h64,flash_qla,gdn,16,8192,64,64,128,3.648,9.575,339,388,0.000000,20,2.37,2.94 diff --git a/benchmark/linear_attention/results/gdn/gb300/gdn_20260824.csv b/benchmark/linear_attention/results/gdn/gb300/gdn_20260824.csv new file mode 100644 index 000000000..3ff81373c --- /dev/null +++ b/benchmark/linear_attention/results/gdn/gb300/gdn_20260824.csv @@ -0,0 +1,41 @@ +case_tag,backend,variant,batch_size,seqlen,num_q_heads,num_kv_heads,head_dim,fwd_ms,bwd_ms,fwd_tflops,bwd_tflops,max_diff,num_iters,fwd_bw,bwd_bw +gdn_h64,cudnn,gdn,4,2048,64,64,128,0.157,0.582,493,399,0.000000,20,3.45,3.02 +gdn_h64,fla,gdn,4,2048,64,64,128,0.739,2.282,105,102,0.000000,20,0.73,0.77 +gdn_h64,flash_qla,gdn,4,2048,64,64,128,0.272,0.726,284,319,0.000000,20,1.99,2.42 +gdn_h64,cudnn,gdn,4,4096,64,64,128,0.275,1.120,563,414,0.000000,20,3.94,3.14 +gdn_h64,fla,gdn,4,4096,64,64,128,1.444,4.537,107,102,0.000000,20,0.75,0.77 +gdn_h64,flash_qla,gdn,4,4096,64,64,128,0.514,1.398,301,332,0.000000,20,2.11,2.51 +gdn_h64,cudnn,gdn,4,8192,64,64,128,0.516,2.158,600,430,0.000000,20,4.20,3.26 +gdn_h64,fla,gdn,4,8192,64,64,128,2.861,9.065,108,102,0.000000,20,0.76,0.78 +gdn_h64,flash_qla,gdn,4,8192,64,64,128,0.997,2.767,310,335,0.000000,20,2.17,2.54 +gdn_h64,cudnn,gdn,4,16384,64,64,128,0.992,4.224,624,439,0.000000,20,4.37,3.33 +gdn_h64,fla,gdn,4,16384,64,64,128,5.707,18.210,108,102,0.000000,20,0.76,0.77 +gdn_h64,flash_qla,gdn,4,16384,64,64,128,1.972,5.452,314,340,0.000000,20,2.20,2.58 +gdn_h64,cudnn,gdn,4,32768,64,64,128,1.939,8.394,638,442,0.000000,20,4.47,3.35 +gdn_h64,fla,gdn,4,32768,64,64,128,11.440,36.799,108,101,0.000000,20,0.76,0.76 +gdn_h64,flash_qla,gdn,4,32768,64,64,128,3.902,10.938,317,339,0.000000,20,2.22,2.57 +gdn_h64,cudnn,gdn,1,8192,64,64,128,0.160,0.552,482,420,0.000000,20,3.37,3.19 +gdn_h64,fla,gdn,1,8192,64,64,128,0.803,2.440,96,95,0.000000,20,0.67,0.72 +gdn_h64,flash_qla,gdn,1,8192,64,64,128,0.366,1.330,211,174,0.000000,20,1.48,1.32 +gdn_h64,cudnn,gdn,2,8192,64,64,128,0.264,1.014,586,457,0.000000,20,4.10,3.47 +gdn_h64,fla,gdn,2,8192,64,64,128,1.448,4.763,107,97,0.000000,20,0.75,0.74 +gdn_h64,flash_qla,gdn,2,8192,64,64,128,0.503,1.428,308,325,0.000000,20,2.15,2.46 +gdn_h64,cudnn,gdn,4,8192,64,64,128,0.514,2.147,602,432,0.000000,20,4.21,3.27 +gdn_h64,fla,gdn,4,8192,64,64,128,2.861,9.067,108,102,0.000000,20,0.76,0.78 +gdn_h64,flash_qla,gdn,4,8192,64,64,128,0.999,2.760,309,336,0.000000,20,2.17,2.55 +gdn_h64,cudnn,gdn,8,8192,64,64,128,0.988,4.425,626,419,0.000000,20,4.38,3.18 +gdn_h64,fla,gdn,8,8192,64,64,128,5.500,17.423,112,106,0.000000,20,0.79,0.81 +gdn_h64,flash_qla,gdn,8,8192,64,64,128,1.964,5.374,315,345,0.000000,20,2.20,2.62 +gdn_h64,cudnn,gdn,16,8192,64,64,128,1.729,7.956,715,466,0.000000,20,5.01,3.53 +gdn_h64,fla,gdn,16,8192,64,64,128,10.868,34.628,114,107,0.000000,20,0.80,0.81 +gdn_h64,flash_qla,gdn,16,8192,64,64,128,3.630,9.615,341,386,0.000000,20,2.38,2.92 +gdn_hon,cudnn_state_on,gdn,4,2048,64,64,128,0.218,0.468,355,495,0.000000,20,2.48,3.75 +gdn_hon,cudnn_state_on,gdn,4,4096,64,64,128,0.393,0.876,393,530,0.000000,20,2.75,4.01 +gdn_hon,cudnn_state_on,gdn,4,8192,64,64,128,0.753,1.680,411,552,0.000000,20,2.87,4.18 +gdn_hon,cudnn_state_on,gdn,4,16384,64,64,128,1.478,3.277,418,566,0.000000,20,2.93,4.29 +gdn_hon,cudnn_state_on,gdn,4,32768,64,64,128,2.924,6.467,423,574,0.000000,20,2.96,4.35 +gdn_hon,cudnn_state_on,gdn,1,8192,64,64,128,0.219,0.424,354,547,0.000000,20,2.48,4.15 +gdn_hon,cudnn_state_on,gdn,2,8192,64,64,128,0.375,0.768,413,604,0.000000,20,2.89,4.58 +gdn_hon,cudnn_state_on,gdn,4,8192,64,64,128,0.761,1.678,407,553,0.000000,20,2.85,4.19 +gdn_hon,cudnn_state_on,gdn,8,8192,64,64,128,1.489,3.446,415,538,0.000000,20,2.91,4.08 +gdn_hon,cudnn_state_on,gdn,16,8192,64,64,128,2.675,6.195,462,599,0.000000,20,3.24,4.54 diff --git a/benchmark/linear_attention/results/gdn/gb300/gdn_fixed_batch_bw.png b/benchmark/linear_attention/results/gdn/gb300/gdn_fixed_batch_bw.png index 1f427d9c4..79c1e7313 100644 Binary files a/benchmark/linear_attention/results/gdn/gb300/gdn_fixed_batch_bw.png and b/benchmark/linear_attention/results/gdn/gb300/gdn_fixed_batch_bw.png differ diff --git a/benchmark/linear_attention/results/gdn/gb300/gdn_fixed_batch_flops.png b/benchmark/linear_attention/results/gdn/gb300/gdn_fixed_batch_flops.png index eca28eefd..389f41fad 100644 Binary files a/benchmark/linear_attention/results/gdn/gb300/gdn_fixed_batch_flops.png and b/benchmark/linear_attention/results/gdn/gb300/gdn_fixed_batch_flops.png differ diff --git a/benchmark/linear_attention/results/gdn/gb300/gdn_fixed_seq_bw.png b/benchmark/linear_attention/results/gdn/gb300/gdn_fixed_seq_bw.png index 113ca830d..3280117c0 100644 Binary files a/benchmark/linear_attention/results/gdn/gb300/gdn_fixed_seq_bw.png and b/benchmark/linear_attention/results/gdn/gb300/gdn_fixed_seq_bw.png differ diff --git a/benchmark/linear_attention/results/gdn/gb300/gdn_fixed_seq_flops.png b/benchmark/linear_attention/results/gdn/gb300/gdn_fixed_seq_flops.png index 2afe662f9..2e594fd64 100644 Binary files a/benchmark/linear_attention/results/gdn/gb300/gdn_fixed_seq_flops.png and b/benchmark/linear_attention/results/gdn/gb300/gdn_fixed_seq_flops.png differ diff --git a/benchmark/linear_attention/results/gdn2/gb200/gdn2_20260814.csv b/benchmark/linear_attention/results/gdn2/gb200/gdn2_20260814.csv deleted file mode 100644 index e3a323218..000000000 --- a/benchmark/linear_attention/results/gdn2/gb200/gdn2_20260814.csv +++ /dev/null @@ -1,31 +0,0 @@ -case_tag,backend,variant,batch_size,seqlen,num_q_heads,num_kv_heads,head_dim,fwd_ms,bwd_ms,fwd_tflops,bwd_tflops,max_diff,num_iters,fwd_bw,bwd_bw -gdn2_h64,cudnn,gdn2,4,2048,64,64,128,0.257,2.050,176,66,0.000000,20,4.18,2.42 -gdn2_h64,fla,gdn2,4,2048,64,64,128,1.536,5.829,29,23,0.000000,20,0.70,0.85 -gdn2_h64,cudnn,gdn2,4,4096,64,64,128,0.474,4.075,190,66,0.000000,20,4.53,2.44 -gdn2_h64,fla,gdn2,4,4096,64,64,128,3.039,11.577,30,23,0.000000,20,0.71,0.86 -gdn2_h64,cudnn,gdn2,4,8192,64,64,128,0.908,8.136,199,67,0.000000,20,4.73,2.44 -gdn2_h64,fla,gdn2,4,8192,64,64,128,6.008,23.062,30,23,0.000000,20,0.71,0.86 -gdn2_h64,cudnn,gdn2,4,16384,64,64,128,1.770,16.278,204,66,0.000000,20,4.85,2.44 -gdn2_h64,fla,gdn2,4,16384,64,64,128,12.007,46.079,30,23,0.000000,20,0.72,0.86 -gdn2_h64,cudnn,gdn2,4,32768,64,64,128,3.508,32.541,206,67,0.000000,20,4.90,2.44 -gdn2_h64,fla,gdn2,4,32768,64,64,128,23.826,88.467,30,24,0.000000,20,0.72,0.90 -gdn2_h64,cudnn,gdn2,1,8192,64,64,128,0.476,3.672,95,37,0.000000,20,2.26,1.35 -gdn2_h64,fla,gdn2,1,8192,64,64,128,1.688,6.099,27,22,0.000000,20,0.64,0.81 -gdn2_h64,cudnn,gdn2,2,8192,64,64,128,0.461,4.031,196,67,0.000000,20,4.66,2.46 -gdn2_h64,fla,gdn2,2,8192,64,64,128,3.060,11.576,29,23,0.000000,20,0.70,0.86 -gdn2_h64,cudnn,gdn2,4,8192,64,64,128,0.906,8.130,199,67,0.000000,20,4.74,2.44 -gdn2_h64,fla,gdn2,4,8192,64,64,128,6.004,23.060,30,23,0.000000,20,0.72,0.86 -gdn2_h64,cudnn,gdn2,8,8192,64,64,128,1.718,16.113,210,67,0.000000,20,5.00,2.47 -gdn2_h64,fla,gdn2,8,8192,64,64,128,11.849,45.914,30,24,0.000000,20,0.72,0.87 -gdn2_h64,cudnn,gdn2,16,8192,64,64,128,3.001,28.796,240,75,0.000000,20,5.72,2.76 -gdn2_h64,fla,gdn2,16,8192,64,64,128,23.265,87.351,31,25,0.000000,20,0.74,0.91 -gdn2_hon,cudnn_state_on,gdn2,4,2048,64,64,128,0.384,1.763,117,77,0.000000,20,2.80,2.82 -gdn2_hon,cudnn_state_on,gdn2,4,4096,64,64,128,0.744,3.494,121,77,0.000000,20,2.89,2.84 -gdn2_hon,cudnn_state_on,gdn2,4,8192,64,64,128,1.468,6.990,123,77,0.000000,20,2.93,2.84 -gdn2_hon,cudnn_state_on,gdn2,4,16384,64,64,128,2.925,13.954,123,78,0.000000,20,2.94,2.85 -gdn2_hon,cudnn_state_on,gdn2,4,32768,64,64,128,5.832,27.875,124,78,0.000000,20,2.95,2.85 -gdn2_hon,cudnn_state_on,gdn2,1,8192,64,64,128,0.691,3.115,65,43,0.000000,20,1.55,1.59 -gdn2_hon,cudnn_state_on,gdn2,2,8192,64,64,128,0.718,3.457,126,78,0.000000,20,2.99,2.87 -gdn2_hon,cudnn_state_on,gdn2,4,8192,64,64,128,1.467,6.990,123,77,0.000000,20,2.93,2.84 -gdn2_hon,cudnn_state_on,gdn2,8,8192,64,64,128,2.919,12.448,124,87,0.000000,20,2.94,3.19 -gdn2_hon,cudnn_state_on,gdn2,16,8192,64,64,128,5.677,24.345,127,89,0.000000,20,3.03,3.26 diff --git a/benchmark/linear_attention/results/gdn2/gb200/gdn2_20260824.csv b/benchmark/linear_attention/results/gdn2/gb200/gdn2_20260824.csv new file mode 100644 index 000000000..1eb96c7d0 --- /dev/null +++ b/benchmark/linear_attention/results/gdn2/gb200/gdn2_20260824.csv @@ -0,0 +1,31 @@ +case_tag,backend,variant,batch_size,seqlen,num_q_heads,num_kv_heads,head_dim,fwd_ms,bwd_ms,fwd_tflops,bwd_tflops,max_diff,num_iters,fwd_bw,bwd_bw +gdn2_h64,cudnn,gdn2,4,2048,64,64,128,0.252,1.734,179,78,0.000000,20,4.27,2.86 +gdn2_h64,fla,gdn2,4,2048,64,64,128,1.545,5.850,29,23,0.000000,20,0.69,0.85 +gdn2_h64,cudnn,gdn2,4,4096,64,64,128,0.471,3.426,191,79,0.000000,20,4.56,2.90 +gdn2_h64,fla,gdn2,4,4096,64,64,128,3.055,11.660,30,23,0.000000,20,0.70,0.85 +gdn2_h64,cudnn,gdn2,4,8192,64,64,128,0.913,6.840,197,79,0.000000,20,4.70,2.90 +gdn2_h64,fla,gdn2,4,8192,64,64,128,6.061,23.219,30,23,0.000000,20,0.71,0.86 +gdn2_h64,cudnn,gdn2,4,16384,64,64,128,1.796,13.660,201,79,0.000000,20,4.78,2.91 +gdn2_h64,fla,gdn2,4,16384,64,64,128,12.070,46.353,30,23,0.000000,20,0.71,0.86 +gdn2_h64,cudnn,gdn2,4,32768,64,64,128,3.561,27.325,203,79,0.000000,20,4.82,2.91 +gdn2_h64,fla,gdn2,4,32768,64,64,128,24.022,89.034,30,24,0.000000,20,0.72,0.89 +gdn2_h64,cudnn,gdn2,1,8192,64,64,128,0.474,3.176,95,43,0.000000,20,2.27,1.56 +gdn2_h64,fla,gdn2,1,8192,64,64,128,1.703,6.136,26,22,0.000000,20,0.63,0.81 +gdn2_h64,cudnn,gdn2,2,8192,64,64,128,0.466,3.841,193,70,0.000000,20,4.61,2.59 +gdn2_h64,fla,gdn2,2,8192,64,64,128,3.076,11.649,29,23,0.000000,20,0.70,0.85 +gdn2_h64,cudnn,gdn2,4,8192,64,64,128,0.913,6.839,198,79,0.000000,20,4.70,2.90 +gdn2_h64,fla,gdn2,4,8192,64,64,128,6.042,23.218,30,23,0.000000,20,0.71,0.86 +gdn2_h64,cudnn,gdn2,8,8192,64,64,128,1.741,14.251,207,76,0.000000,20,4.94,2.79 +gdn2_h64,fla,gdn2,8,8192,64,64,128,11.914,46.208,30,23,0.000000,20,0.72,0.86 +gdn2_h64,cudnn,gdn2,16,8192,64,64,128,3.047,26.265,237,82,0.000000,20,5.64,3.03 +gdn2_h64,fla,gdn2,16,8192,64,64,128,23.412,87.895,31,25,0.000000,20,0.73,0.90 +gdn2_hon,cudnn_state_on,gdn2,4,2048,64,64,128,0.387,1.492,117,91,0.000000,20,2.77,3.33 +gdn2_hon,cudnn_state_on,gdn2,4,4096,64,64,128,0.745,2.835,121,95,0.000000,20,2.88,3.50 +gdn2_hon,cudnn_state_on,gdn2,4,8192,64,64,128,1.469,5.626,123,96,0.000000,20,2.92,3.53 +gdn2_hon,cudnn_state_on,gdn2,4,16384,64,64,128,2.928,11.214,123,97,0.000000,20,2.93,3.54 +gdn2_hon,cudnn_state_on,gdn2,4,32768,64,64,128,5.854,22.455,123,96,0.000000,20,2.93,3.54 +gdn2_hon,cudnn_state_on,gdn2,1,8192,64,64,128,0.687,2.677,66,51,0.000000,20,1.56,1.86 +gdn2_hon,cudnn_state_on,gdn2,2,8192,64,64,128,0.727,3.252,124,83,0.000000,20,2.95,3.05 +gdn2_hon,cudnn_state_on,gdn2,4,8192,64,64,128,1.468,5.624,123,96,0.000000,20,2.93,3.53 +gdn2_hon,cudnn_state_on,gdn2,8,8192,64,64,128,2.929,11.722,123,92,0.000000,20,2.93,3.39 +gdn2_hon,cudnn_state_on,gdn2,16,8192,64,64,128,5.413,21.685,133,100,0.000000,20,3.17,3.66 diff --git a/benchmark/linear_attention/results/gdn2/gb200/gdn2_fixed_batch_bw.png b/benchmark/linear_attention/results/gdn2/gb200/gdn2_fixed_batch_bw.png index 4adfd1387..cd01fe2fe 100644 Binary files a/benchmark/linear_attention/results/gdn2/gb200/gdn2_fixed_batch_bw.png and b/benchmark/linear_attention/results/gdn2/gb200/gdn2_fixed_batch_bw.png differ diff --git a/benchmark/linear_attention/results/gdn2/gb200/gdn2_fixed_batch_flops.png b/benchmark/linear_attention/results/gdn2/gb200/gdn2_fixed_batch_flops.png index 3f13b8393..fb531f6fe 100644 Binary files a/benchmark/linear_attention/results/gdn2/gb200/gdn2_fixed_batch_flops.png and b/benchmark/linear_attention/results/gdn2/gb200/gdn2_fixed_batch_flops.png differ diff --git a/benchmark/linear_attention/results/gdn2/gb200/gdn2_fixed_seq_bw.png b/benchmark/linear_attention/results/gdn2/gb200/gdn2_fixed_seq_bw.png index 1bec90ae8..cf58f61da 100644 Binary files a/benchmark/linear_attention/results/gdn2/gb200/gdn2_fixed_seq_bw.png and b/benchmark/linear_attention/results/gdn2/gb200/gdn2_fixed_seq_bw.png differ diff --git a/benchmark/linear_attention/results/gdn2/gb200/gdn2_fixed_seq_flops.png b/benchmark/linear_attention/results/gdn2/gb200/gdn2_fixed_seq_flops.png index 86d159c44..c6dec9ac3 100644 Binary files a/benchmark/linear_attention/results/gdn2/gb200/gdn2_fixed_seq_flops.png and b/benchmark/linear_attention/results/gdn2/gb200/gdn2_fixed_seq_flops.png differ diff --git a/benchmark/linear_attention/results/gdn2/gb300/gdn2_20260814.csv b/benchmark/linear_attention/results/gdn2/gb300/gdn2_20260814.csv deleted file mode 100644 index 11ed31700..000000000 --- a/benchmark/linear_attention/results/gdn2/gb300/gdn2_20260814.csv +++ /dev/null @@ -1,31 +0,0 @@ -case_tag,backend,variant,batch_size,seqlen,num_q_heads,num_kv_heads,head_dim,fwd_ms,bwd_ms,fwd_tflops,bwd_tflops,max_diff,num_iters,fwd_bw,bwd_bw -gdn2_h64,cudnn,gdn2,4,2048,64,64,128,0.249,2.034,181,67,0.000000,20,4.32,2.44 -gdn2_h64,fla,gdn2,4,2048,64,64,128,1.479,5.669,30,24,0.000000,20,0.73,0.88 -gdn2_h64,cudnn,gdn2,4,4096,64,64,128,0.473,4.032,191,67,0.000000,20,4.54,2.46 -gdn2_h64,fla,gdn2,4,4096,64,64,128,2.916,11.276,31,24,0.000000,20,0.74,0.88 -gdn2_h64,cudnn,gdn2,4,8192,64,64,128,0.918,8.036,197,67,0.000000,20,4.68,2.47 -gdn2_h64,fla,gdn2,4,8192,64,64,128,5.772,22.466,31,24,0.000000,20,0.74,0.88 -gdn2_h64,cudnn,gdn2,4,16384,64,64,128,1.809,16.039,199,67,0.000000,20,4.75,2.48 -gdn2_h64,fla,gdn2,4,16384,64,64,128,11.469,44.851,31,24,0.000000,20,0.75,0.89 -gdn2_h64,cudnn,gdn2,4,32768,64,64,128,3.590,32.059,201,68,0.000000,20,4.79,2.48 -gdn2_h64,fla,gdn2,4,32768,64,64,128,22.873,86.102,32,25,0.000000,20,0.75,0.92 -gdn2_h64,cudnn,gdn2,1,8192,64,64,128,0.477,3.610,94,37,0.000000,20,2.25,1.38 -gdn2_h64,fla,gdn2,1,8192,64,64,128,1.622,5.937,28,23,0.000000,20,0.66,0.84 -gdn2_h64,cudnn,gdn2,2,8192,64,64,128,0.471,4.033,192,67,0.000000,20,4.56,2.46 -gdn2_h64,fla,gdn2,2,8192,64,64,128,2.931,11.275,31,24,0.000000,20,0.73,0.88 -gdn2_h64,cudnn,gdn2,4,8192,64,64,128,0.918,8.032,196,67,0.000000,20,4.68,2.47 -gdn2_h64,fla,gdn2,4,8192,64,64,128,5.764,22.456,31,24,0.000000,20,0.75,0.88 -gdn2_h64,cudnn,gdn2,8,8192,64,64,128,1.752,14.667,206,74,0.000000,20,4.90,2.71 -gdn2_h64,fla,gdn2,8,8192,64,64,128,11.361,44.740,32,24,0.000000,20,0.76,0.89 -gdn2_h64,cudnn,gdn2,16,8192,64,64,128,3.066,28.539,235,76,0.000000,20,5.60,2.78 -gdn2_h64,fla,gdn2,16,8192,64,64,128,22.279,84.983,32,25,0.000000,20,0.77,0.93 -gdn2_hon,cudnn_state_on,gdn2,4,2048,64,64,128,0.375,1.762,120,77,0.000000,20,2.86,2.82 -gdn2_hon,cudnn_state_on,gdn2,4,4096,64,64,128,0.724,3.483,125,78,0.000000,20,2.97,2.85 -gdn2_hon,cudnn_state_on,gdn2,4,8192,64,64,128,1.426,6.937,126,78,0.000000,20,3.01,2.86 -gdn2_hon,cudnn_state_on,gdn2,4,16384,64,64,128,2.833,13.838,127,78,0.000000,20,3.03,2.87 -gdn2_hon,cudnn_state_on,gdn2,4,32768,64,64,128,5.643,27.648,128,78,0.000000,20,3.04,2.87 -gdn2_hon,cudnn_state_on,gdn2,1,8192,64,64,128,0.684,3.115,66,43,0.000000,20,1.57,1.59 -gdn2_hon,cudnn_state_on,gdn2,2,8192,64,64,128,0.708,3.452,127,78,0.000000,20,3.03,2.88 -gdn2_hon,cudnn_state_on,gdn2,4,8192,64,64,128,1.428,6.937,126,78,0.000000,20,3.01,2.86 -gdn2_hon,cudnn_state_on,gdn2,8,8192,64,64,128,2.850,12.342,127,88,0.000000,20,3.01,3.22 -gdn2_hon,cudnn_state_on,gdn2,16,8192,64,64,128,5.265,24.217,137,89,0.000000,20,3.26,3.28 diff --git a/benchmark/linear_attention/results/gdn2/gb300/gdn2_20260824.csv b/benchmark/linear_attention/results/gdn2/gb300/gdn2_20260824.csv new file mode 100644 index 000000000..b0543ccde --- /dev/null +++ b/benchmark/linear_attention/results/gdn2/gb300/gdn2_20260824.csv @@ -0,0 +1,31 @@ +case_tag,backend,variant,batch_size,seqlen,num_q_heads,num_kv_heads,head_dim,fwd_ms,bwd_ms,fwd_tflops,bwd_tflops,max_diff,num_iters,fwd_bw,bwd_bw +gdn2_h64,cudnn,gdn2,4,2048,64,64,128,0.246,1.720,183,79,0.000000,20,4.36,2.89 +gdn2_h64,fla,gdn2,4,2048,64,64,128,1.477,5.666,31,24,0.000000,20,0.73,0.88 +gdn2_h64,cudnn,gdn2,4,4096,64,64,128,0.467,3.397,193,80,0.000000,20,4.59,2.92 +gdn2_h64,fla,gdn2,4,4096,64,64,128,2.911,11.255,31,24,0.000000,20,0.74,0.88 +gdn2_h64,cudnn,gdn2,4,8192,64,64,128,0.908,6.769,199,80,0.000000,20,4.73,2.93 +gdn2_h64,fla,gdn2,4,8192,64,64,128,5.766,22.445,31,24,0.000000,20,0.74,0.89 +gdn2_h64,cudnn,gdn2,4,16384,64,64,128,1.791,13.519,201,80,0.000000,20,4.80,2.94 +gdn2_h64,fla,gdn2,4,16384,64,64,128,11.481,44.833,31,24,0.000000,20,0.75,0.89 +gdn2_h64,cudnn,gdn2,4,32768,64,64,128,3.555,27.019,203,80,0.000000,20,4.83,2.94 +gdn2_h64,fla,gdn2,4,32768,64,64,128,22.823,85.995,32,25,0.000000,20,0.75,0.92 +gdn2_h64,cudnn,gdn2,1,8192,64,64,128,0.471,3.085,96,44,0.000000,20,2.28,1.61 +gdn2_h64,fla,gdn2,1,8192,64,64,128,1.621,5.931,28,23,0.000000,20,0.66,0.84 +gdn2_h64,cudnn,gdn2,2,8192,64,64,128,0.464,3.809,195,71,0.000000,20,4.63,2.61 +gdn2_h64,fla,gdn2,2,8192,64,64,128,2.929,11.261,31,24,0.000000,20,0.73,0.88 +gdn2_h64,cudnn,gdn2,4,8192,64,64,128,0.908,6.768,199,80,0.000000,20,4.73,2.94 +gdn2_h64,fla,gdn2,4,8192,64,64,128,5.755,22.438,31,24,0.000000,20,0.75,0.89 +gdn2_h64,cudnn,gdn2,8,8192,64,64,128,1.736,14.128,208,77,0.000000,20,4.95,2.81 +gdn2_h64,fla,gdn2,8,8192,64,64,128,11.343,44.689,32,24,0.000000,20,0.76,0.89 +gdn2_h64,cudnn,gdn2,16,8192,64,64,128,3.040,25.933,237,83,0.000000,20,5.65,3.06 +gdn2_h64,fla,gdn2,16,8192,64,64,128,22.264,84.903,32,25,0.000000,20,0.77,0.94 +gdn2_hon,cudnn_state_on,gdn2,4,2048,64,64,128,0.379,1.533,119,88,0.000000,20,2.83,3.24 +gdn2_hon,cudnn_state_on,gdn2,4,4096,64,64,128,0.729,2.807,124,96,0.000000,20,2.95,3.54 +gdn2_hon,cudnn_state_on,gdn2,4,8192,64,64,128,1.435,5.581,126,97,0.000000,20,2.99,3.56 +gdn2_hon,cudnn_state_on,gdn2,4,16384,64,64,128,2.851,11.159,127,97,0.000000,20,3.01,3.56 +gdn2_hon,cudnn_state_on,gdn2,4,32768,64,64,128,5.682,22.293,127,97,0.000000,20,3.02,3.56 +gdn2_hon,cudnn_state_on,gdn2,1,8192,64,64,128,0.682,2.618,66,52,0.000000,20,1.57,1.90 +gdn2_hon,cudnn_state_on,gdn2,2,8192,64,64,128,0.717,3.239,126,84,0.000000,20,3.00,3.07 +gdn2_hon,cudnn_state_on,gdn2,4,8192,64,64,128,1.433,5.585,126,97,0.000000,20,3.00,3.56 +gdn2_hon,cudnn_state_on,gdn2,8,8192,64,64,128,2.846,11.635,127,93,0.000000,20,3.02,3.41 +gdn2_hon,cudnn_state_on,gdn2,16,8192,64,64,128,5.249,21.566,137,100,0.000000,20,3.27,3.68 diff --git a/benchmark/linear_attention/results/gdn2/gb300/gdn2_fixed_batch_bw.png b/benchmark/linear_attention/results/gdn2/gb300/gdn2_fixed_batch_bw.png index 952aa3897..a29ca5d9f 100644 Binary files a/benchmark/linear_attention/results/gdn2/gb300/gdn2_fixed_batch_bw.png and b/benchmark/linear_attention/results/gdn2/gb300/gdn2_fixed_batch_bw.png differ diff --git a/benchmark/linear_attention/results/gdn2/gb300/gdn2_fixed_batch_flops.png b/benchmark/linear_attention/results/gdn2/gb300/gdn2_fixed_batch_flops.png index 0d37fe54e..d109a4b1e 100644 Binary files a/benchmark/linear_attention/results/gdn2/gb300/gdn2_fixed_batch_flops.png and b/benchmark/linear_attention/results/gdn2/gb300/gdn2_fixed_batch_flops.png differ diff --git a/benchmark/linear_attention/results/gdn2/gb300/gdn2_fixed_seq_bw.png b/benchmark/linear_attention/results/gdn2/gb300/gdn2_fixed_seq_bw.png index 892285ac6..37c8b043d 100644 Binary files a/benchmark/linear_attention/results/gdn2/gb300/gdn2_fixed_seq_bw.png and b/benchmark/linear_attention/results/gdn2/gb300/gdn2_fixed_seq_bw.png differ diff --git a/benchmark/linear_attention/results/gdn2/gb300/gdn2_fixed_seq_flops.png b/benchmark/linear_attention/results/gdn2/gb300/gdn2_fixed_seq_flops.png index 054ff4f92..90f519a63 100644 Binary files a/benchmark/linear_attention/results/gdn2/gb300/gdn2_fixed_seq_flops.png and b/benchmark/linear_attention/results/gdn2/gb300/gdn2_fixed_seq_flops.png differ diff --git a/benchmark/linear_attention/results/kda/gb200/kda_20260814.csv b/benchmark/linear_attention/results/kda/gb200/kda_20260814.csv deleted file mode 100644 index 242c8cae6..000000000 --- a/benchmark/linear_attention/results/kda/gb200/kda_20260814.csv +++ /dev/null @@ -1,41 +0,0 @@ -case_tag,backend,variant,batch_size,seqlen,num_q_heads,num_kv_heads,head_dim,fwd_ms,bwd_ms,fwd_tflops,bwd_tflops,max_diff,num_iters,fwd_bw,bwd_bw -kda_h64,cudnn,kda,4,2048,64,64,128,0.265,1.875,170,72,0.000000,20,3.05,2.22 -kda_hon,cudnn_state_on,kda,4,2048,64,64,128,0.354,1.566,127,86,0.000000,20,2.28,2.66 -kda_h64,fla,kda,4,2048,64,64,128,1.462,5.450,31,25,0.000000,20,0.55,0.76 -kda_h64,flash_kda,kda,4,2048,64,64,128,0.436,0.000,103,0,0.000000,20,1.85,0.00 -kda_h64,cudnn,kda,4,4096,64,64,128,0.506,3.727,178,73,0.000000,20,3.19,2.24 -kda_hon,cudnn_state_on,kda,4,4096,64,64,128,0.679,3.104,133,87,0.000000,20,2.38,2.68 -kda_h64,fla,kda,4,4096,64,64,128,2.892,10.862,31,25,0.000000,20,0.56,0.77 -kda_h64,flash_kda,kda,4,4096,64,64,128,0.851,0.000,106,0,0.000000,20,1.90,0.00 -kda_h64,cudnn,kda,4,8192,64,64,128,0.975,7.420,185,73,0.000000,20,3.31,2.25 -kda_hon,cudnn_state_on,kda,4,8192,64,64,128,1.337,6.181,135,88,0.000000,20,2.41,2.70 -kda_h64,fla,kda,4,8192,64,64,128,5.699,21.766,32,25,0.000000,20,0.57,0.77 -kda_h64,flash_kda,kda,4,8192,64,64,128,1.692,0.000,107,0,0.000000,20,1.91,0.00 -kda_h64,cudnn,kda,4,16384,64,64,128,1.923,14.847,188,73,0.000000,20,3.36,2.25 -kda_hon,cudnn_state_on,kda,4,16384,64,64,128,2.674,12.330,135,88,0.000000,20,2.42,2.70 -kda_h64,fla,kda,4,16384,64,64,128,11.386,43.630,32,25,0.000000,20,0.57,0.76 -kda_h64,flash_kda,kda,4,16384,64,64,128,3.368,0.000,107,0,0.000000,20,1.92,0.00 -kda_h64,cudnn,kda,4,32768,64,64,128,3.896,29.622,185,73,0.000000,20,3.32,2.25 -kda_hon,cudnn_state_on,kda,4,32768,64,64,128,5.256,24.661,137,88,0.000000,20,2.46,2.70 -kda_h64,fla,kda,4,32768,64,64,128,22.736,87.647,32,25,0.000000,20,0.57,0.76 -kda_h64,flash_kda,kda,4,32768,64,64,128,6.717,0.000,107,0,0.000000,20,1.92,0.00 -kda_h64,cudnn,kda,1,8192,64,64,128,0.485,2.861,93,47,0.000000,20,1.67,1.46 -kda_hon,cudnn_state_on,kda,1,8192,64,64,128,0.662,2.357,68,57,0.000000,20,1.22,1.77 -kda_h64,fla,kda,1,8192,64,64,128,1.609,5.757,28,24,0.000000,20,0.50,0.72 -kda_h64,flash_kda,kda,1,8192,64,64,128,0.902,0.000,50,0,0.000000,20,0.90,0.00 -kda_h64,cudnn,kda,2,8192,64,64,128,0.500,3.692,180,73,0.000000,20,3.23,2.26 -kda_hon,cudnn_state_on,kda,2,8192,64,64,128,0.668,3.127,135,87,0.000000,20,2.42,2.66 -kda_h64,fla,kda,2,8192,64,64,128,2.892,10.908,31,25,0.000000,20,0.56,0.76 -kda_h64,flash_kda,kda,2,8192,64,64,128,1.074,0.000,84,0,0.000000,20,1.50,0.00 -kda_h64,cudnn,kda,4,8192,64,64,128,0.984,7.425,183,73,0.000000,20,3.28,2.24 -kda_hon,cudnn_state_on,kda,4,8192,64,64,128,1.330,6.186,136,87,0.000000,20,2.43,2.69 -kda_h64,fla,kda,4,8192,64,64,128,5.710,21.769,32,25,0.000000,20,0.57,0.77 -kda_h64,flash_kda,kda,4,8192,64,64,128,1.692,0.000,107,0,0.000000,20,1.91,0.00 -kda_h64,cudnn,kda,8,8192,64,64,128,1.898,14.872,190,73,0.000000,20,3.40,2.24 -kda_hon,cudnn_state_on,kda,8,8192,64,64,128,2.706,12.279,133,88,0.000000,20,2.39,2.71 -kda_h64,fla,kda,8,8192,64,64,128,11.241,43.337,32,25,0.000000,20,0.57,0.77 -kda_h64,flash_kda,kda,8,8192,64,64,128,3.363,0.000,107,0,0.000000,20,1.92,0.00 -kda_h64,cudnn,kda,16,8192,64,64,128,3.292,26.193,219,83,0.000000,20,3.92,2.55 -kda_hon,cudnn_state_on,kda,16,8192,64,64,128,4.926,21.509,146,101,0.000000,20,2.62,3.10 -kda_h64,fla,kda,16,8192,64,64,128,22.024,85.810,33,25,0.000000,20,0.59,0.78 -kda_h64,flash_kda,kda,16,8192,64,64,128,6.369,0.000,113,0,0.000000,20,2.03,0.00 diff --git a/benchmark/linear_attention/results/kda/gb200/kda_20260824.csv b/benchmark/linear_attention/results/kda/gb200/kda_20260824.csv new file mode 100644 index 000000000..51a9b5d3d --- /dev/null +++ b/benchmark/linear_attention/results/kda/gb200/kda_20260824.csv @@ -0,0 +1,41 @@ +case_tag,backend,variant,batch_size,seqlen,num_q_heads,num_kv_heads,head_dim,fwd_ms,bwd_ms,fwd_tflops,bwd_tflops,max_diff,num_iters,fwd_bw,bwd_bw +kda_h64,cudnn,kda,4,2048,64,64,128,0.268,1.647,168,82,0.000000,20,3.01,2.53 +kda_h64,fla,kda,4,2048,64,64,128,1.465,5.480,31,25,0.000000,20,0.55,0.76 +kda_h64,flash_kda,kda,4,2048,64,64,128,0.435,0.000,104,0,0.000000,20,1.86,0.00 +kda_h64,cudnn,kda,4,4096,64,64,128,0.507,3.260,178,83,0.000000,20,3.18,2.56 +kda_h64,fla,kda,4,4096,64,64,128,2.903,10.943,31,25,0.000000,20,0.56,0.76 +kda_h64,flash_kda,kda,4,4096,64,64,128,0.852,0.000,106,0,0.000000,20,1.90,0.00 +kda_h64,cudnn,kda,4,8192,64,64,128,0.989,6.515,182,83,0.000000,20,3.26,2.56 +kda_h64,fla,kda,4,8192,64,64,128,5.733,21.943,31,25,0.000000,20,0.56,0.76 +kda_h64,flash_kda,kda,4,8192,64,64,128,1.698,0.000,106,0,0.000000,20,1.90,0.00 +kda_h64,cudnn,kda,4,16384,64,64,128,1.962,13.009,184,83,0.000000,20,3.29,2.56 +kda_h64,fla,kda,4,16384,64,64,128,11.463,43.991,31,25,0.000000,20,0.56,0.76 +kda_h64,flash_kda,kda,4,16384,64,64,128,3.383,0.000,107,0,0.000000,20,1.91,0.00 +kda_h64,cudnn,kda,4,32768,64,64,128,3.867,25.983,187,83,0.000000,20,3.34,2.57 +kda_h64,fla,kda,4,32768,64,64,128,22.875,88.295,32,25,0.000000,20,0.56,0.76 +kda_h64,flash_kda,kda,4,32768,64,64,128,6.750,0.000,107,0,0.000000,20,1.91,0.00 +kda_h64,cudnn,kda,1,8192,64,64,128,0.486,3.142,93,43,0.000000,20,1.66,1.33 +kda_h64,fla,kda,1,8192,64,64,128,1.621,5.804,28,23,0.000000,20,0.50,0.72 +kda_h64,flash_kda,kda,1,8192,64,64,128,0.904,0.000,50,0,0.000000,20,0.89,0.00 +kda_h64,cudnn,kda,2,8192,64,64,128,0.498,3.687,181,73,0.000000,20,3.24,2.26 +kda_h64,fla,kda,2,8192,64,64,128,2.910,10.997,31,25,0.000000,20,0.55,0.76 +kda_h64,flash_kda,kda,2,8192,64,64,128,1.075,0.000,84,0,0.000000,20,1.50,0.00 +kda_h64,cudnn,kda,4,8192,64,64,128,0.991,6.512,182,83,0.000000,20,3.26,2.56 +kda_h64,fla,kda,4,8192,64,64,128,5.750,21.933,31,25,0.000000,20,0.56,0.76 +kda_h64,flash_kda,kda,4,8192,64,64,128,1.700,0.000,106,0,0.000000,20,1.90,0.00 +kda_h64,cudnn,kda,8,8192,64,64,128,1.880,13.706,192,79,0.000000,20,3.44,2.43 +kda_h64,fla,kda,8,8192,64,64,128,11.308,43.677,32,25,0.000000,20,0.57,0.76 +kda_h64,flash_kda,kda,8,8192,64,64,128,3.374,0.000,107,0,0.000000,20,1.91,0.00 +kda_h64,cudnn,kda,16,8192,64,64,128,3.270,24.911,221,87,0.000000,20,3.95,2.68 +kda_h64,fla,kda,16,8192,64,64,128,22.164,86.483,33,25,0.000000,20,0.58,0.77 +kda_h64,flash_kda,kda,16,8192,64,64,128,6.399,0.000,113,0,0.000000,20,2.02,0.00 +kda_hon,cudnn_state_on,kda,4,2048,64,64,128,0.355,1.355,127,100,0.000000,20,2.27,3.08 +kda_hon,cudnn_state_on,kda,4,4096,64,64,128,0.682,2.685,132,101,0.000000,20,2.37,3.10 +kda_hon,cudnn_state_on,kda,4,8192,64,64,128,1.339,5.381,135,101,0.000000,20,2.41,3.10 +kda_hon,cudnn_state_on,kda,4,16384,64,64,128,2.650,10.727,136,101,0.000000,20,2.44,3.11 +kda_hon,cudnn_state_on,kda,4,32768,64,64,128,5.284,21.439,137,101,0.000000,20,2.44,3.11 +kda_hon,cudnn_state_on,kda,1,8192,64,64,128,0.661,2.621,68,52,0.000000,20,1.22,1.59 +kda_hon,cudnn_state_on,kda,2,8192,64,64,128,0.675,3.097,134,87,0.000000,20,2.39,2.69 +kda_hon,cudnn_state_on,kda,4,8192,64,64,128,1.339,5.380,135,101,0.000000,20,2.41,3.10 +kda_hon,cudnn_state_on,kda,8,8192,64,64,128,2.727,11.280,132,96,0.000000,20,2.37,2.96 +kda_hon,cudnn_state_on,kda,16,8192,64,64,128,4.944,20.580,146,105,0.000000,20,2.61,3.24 diff --git a/benchmark/linear_attention/results/kda/gb200/kda_fixed_batch_bw.png b/benchmark/linear_attention/results/kda/gb200/kda_fixed_batch_bw.png index 37db34ffe..e73d14fc8 100644 Binary files a/benchmark/linear_attention/results/kda/gb200/kda_fixed_batch_bw.png and b/benchmark/linear_attention/results/kda/gb200/kda_fixed_batch_bw.png differ diff --git a/benchmark/linear_attention/results/kda/gb200/kda_fixed_batch_flops.png b/benchmark/linear_attention/results/kda/gb200/kda_fixed_batch_flops.png index 90bc35b48..23b97616e 100644 Binary files a/benchmark/linear_attention/results/kda/gb200/kda_fixed_batch_flops.png and b/benchmark/linear_attention/results/kda/gb200/kda_fixed_batch_flops.png differ diff --git a/benchmark/linear_attention/results/kda/gb200/kda_fixed_seq_bw.png b/benchmark/linear_attention/results/kda/gb200/kda_fixed_seq_bw.png index 165b17bca..0a0f0e4b1 100644 Binary files a/benchmark/linear_attention/results/kda/gb200/kda_fixed_seq_bw.png and b/benchmark/linear_attention/results/kda/gb200/kda_fixed_seq_bw.png differ diff --git a/benchmark/linear_attention/results/kda/gb200/kda_fixed_seq_flops.png b/benchmark/linear_attention/results/kda/gb200/kda_fixed_seq_flops.png index b95798e4d..2db7d97df 100644 Binary files a/benchmark/linear_attention/results/kda/gb200/kda_fixed_seq_flops.png and b/benchmark/linear_attention/results/kda/gb200/kda_fixed_seq_flops.png differ diff --git a/benchmark/linear_attention/results/kda/gb300/kda_20260814.csv b/benchmark/linear_attention/results/kda/gb300/kda_20260814.csv deleted file mode 100644 index 5b9d21bf2..000000000 --- a/benchmark/linear_attention/results/kda/gb300/kda_20260814.csv +++ /dev/null @@ -1,41 +0,0 @@ -case_tag,backend,variant,batch_size,seqlen,num_q_heads,num_kv_heads,head_dim,fwd_ms,bwd_ms,fwd_tflops,bwd_tflops,max_diff,num_iters,fwd_bw,bwd_bw -kda_h64,cudnn,kda,4,2048,64,64,128,0.261,1.867,172,72,0.000000,20,3.09,2.23 -kda_hon,cudnn_state_on,kda,4,2048,64,64,128,0.349,1.564,129,86,0.000000,20,2.31,2.66 -kda_h64,fla,kda,4,2048,64,64,128,1.421,5.352,32,25,0.000000,20,0.57,0.78 -kda_h64,flash_kda,kda,4,2048,64,64,128,0.433,0.000,104,0,0.000000,20,1.86,0.00 -kda_h64,cudnn,kda,4,4096,64,64,128,0.496,3.709,182,73,0.000000,20,3.26,2.25 -kda_hon,cudnn_state_on,kda,4,4096,64,64,128,0.673,3.097,134,87,0.000000,20,2.40,2.69 -kda_h64,fla,kda,4,4096,64,64,128,2.803,10.697,32,25,0.000000,20,0.58,0.78 -kda_h64,flash_kda,kda,4,4096,64,64,128,0.849,0.000,106,0,0.000000,20,1.90,0.00 -kda_h64,cudnn,kda,4,8192,64,64,128,0.965,7.391,187,73,0.000000,20,3.35,2.26 -kda_hon,cudnn_state_on,kda,4,8192,64,64,128,1.319,6.176,137,88,0.000000,20,2.45,2.70 -kda_h64,fla,kda,4,8192,64,64,128,5.533,21.361,33,25,0.000000,20,0.58,0.78 -kda_h64,flash_kda,kda,4,8192,64,64,128,1.687,0.000,107,0,0.000000,20,1.91,0.00 -kda_h64,cudnn,kda,4,16384,64,64,128,1.899,14.751,190,73,0.000000,20,3.40,2.26 -kda_hon,cudnn_state_on,kda,4,16384,64,64,128,2.607,12.312,138,88,0.000000,20,2.48,2.71 -kda_h64,fla,kda,4,16384,64,64,128,11.072,42.801,33,25,0.000000,20,0.58,0.78 -kda_h64,flash_kda,kda,4,16384,64,64,128,3.357,0.000,107,0,0.000000,20,1.92,0.00 -kda_h64,cudnn,kda,4,32768,64,64,128,3.844,29.469,188,73,0.000000,20,3.36,2.26 -kda_hon,cudnn_state_on,kda,4,32768,64,64,128,5.200,24.597,139,88,0.000000,20,2.48,2.71 -kda_h64,fla,kda,4,32768,64,64,128,22.044,85.871,33,25,0.000000,20,0.59,0.78 -kda_h64,flash_kda,kda,4,32768,64,64,128,6.700,0.000,108,0,0.000000,20,1.93,0.00 -kda_h64,cudnn,kda,1,8192,64,64,128,0.481,2.860,94,47,0.000000,20,1.68,1.46 -kda_hon,cudnn_state_on,kda,1,8192,64,64,128,0.655,2.524,69,54,0.000000,20,1.23,1.65 -kda_h64,fla,kda,1,8192,64,64,128,1.566,5.651,29,24,0.000000,20,0.52,0.74 -kda_h64,flash_kda,kda,1,8192,64,64,128,0.897,0.000,50,0,0.000000,20,0.90,0.00 -kda_h64,cudnn,kda,2,8192,64,64,128,0.499,3.676,181,74,0.000000,20,3.24,2.27 -kda_hon,cudnn_state_on,kda,2,8192,64,64,128,0.654,3.121,138,87,0.000000,20,2.47,2.67 -kda_h64,fla,kda,2,8192,64,64,128,2.814,10.715,32,25,0.000000,20,0.57,0.78 -kda_h64,flash_kda,kda,2,8192,64,64,128,1.070,0.000,84,0,0.000000,20,1.51,0.00 -kda_h64,cudnn,kda,4,8192,64,64,128,0.976,7.393,185,73,0.000000,20,3.31,2.25 -kda_hon,cudnn_state_on,kda,4,8192,64,64,128,1.319,6.175,137,88,0.000000,20,2.45,2.70 -kda_h64,fla,kda,4,8192,64,64,128,5.550,21.362,33,25,0.000000,20,0.58,0.78 -kda_h64,flash_kda,kda,4,8192,64,64,128,1.686,0.000,107,0,0.000000,20,1.92,0.00 -kda_h64,cudnn,kda,8,8192,64,64,128,1.855,13.464,194,80,0.000000,20,3.48,2.48 -kda_hon,cudnn_state_on,kda,8,8192,64,64,128,2.669,11.021,135,98,0.000000,20,2.42,3.02 -kda_h64,fla,kda,8,8192,64,64,128,10.939,42.524,33,25,0.000000,20,0.59,0.78 -kda_h64,flash_kda,kda,8,8192,64,64,128,3.351,0.000,108,0,0.000000,20,1.93,0.00 -kda_h64,cudnn,kda,16,8192,64,64,128,3.224,26.006,224,83,0.000000,20,4.01,2.56 -kda_hon,cudnn_state_on,kda,16,8192,64,64,128,4.767,21.475,151,101,0.000000,20,2.71,3.10 -kda_h64,fla,kda,16,8192,64,64,128,21.432,84.215,34,26,0.000000,20,0.60,0.79 -kda_h64,flash_kda,kda,16,8192,64,64,128,6.346,0.000,114,0,0.000000,20,2.04,0.00 diff --git a/benchmark/linear_attention/results/kda/gb300/kda_20260824.csv b/benchmark/linear_attention/results/kda/gb300/kda_20260824.csv new file mode 100644 index 000000000..c0b885e7d --- /dev/null +++ b/benchmark/linear_attention/results/kda/gb300/kda_20260824.csv @@ -0,0 +1,41 @@ +case_tag,backend,variant,batch_size,seqlen,num_q_heads,num_kv_heads,head_dim,fwd_ms,bwd_ms,fwd_tflops,bwd_tflops,max_diff,num_iters,fwd_bw,bwd_bw +kda_h64,cudnn,kda,4,2048,64,64,128,0.255,1.632,177,83,0.000000,20,3.16,2.55 +kda_h64,fla,kda,4,2048,64,64,128,1.421,5.348,32,25,0.000000,20,0.57,0.78 +kda_h64,cudnn,kda,4,4096,64,64,128,0.490,3.231,184,84,0.000000,20,3.29,2.58 +kda_h64,fla,kda,4,4096,64,64,128,2.795,10.683,32,25,0.000000,20,0.58,0.78 +kda_h64,cudnn,kda,4,8192,64,64,128,0.936,6.429,193,84,0.000000,20,3.45,2.59 +kda_h64,fla,kda,4,8192,64,64,128,5.531,21.341,33,25,0.000000,20,0.58,0.78 +kda_h64,cudnn,kda,4,16384,64,64,128,1.845,12.829,196,84,0.000000,20,3.50,2.60 +kda_h64,fla,kda,4,16384,64,64,128,11.021,42.768,33,25,0.000000,20,0.59,0.78 +kda_h64,cudnn,kda,4,32768,64,64,128,3.734,25.634,193,84,0.000000,20,3.46,2.60 +kda_h64,fla,kda,4,32768,64,64,128,22.078,85.822,33,25,0.000000,20,0.59,0.78 +kda_h64,cudnn,kda,1,8192,64,64,128,0.478,3.033,94,45,0.000000,20,1.69,1.37 +kda_h64,fla,kda,1,8192,64,64,128,1.564,5.644,29,24,0.000000,20,0.52,0.74 +kda_h64,cudnn,kda,2,8192,64,64,128,0.489,3.652,184,74,0.000000,20,3.30,2.28 +kda_h64,fla,kda,2,8192,64,64,128,2.812,10.699,32,25,0.000000,20,0.57,0.78 +kda_h64,cudnn,kda,4,8192,64,64,128,0.937,6.428,193,84,0.000000,20,3.45,2.59 +kda_h64,fla,kda,4,8192,64,64,128,5.554,21.337,32,25,0.000000,20,0.58,0.78 +kda_h64,cudnn,kda,8,8192,64,64,128,1.801,13.542,200,80,0.000000,20,3.59,2.46 +kda_h64,fla,kda,8,8192,64,64,128,10.925,42.514,33,25,0.000000,20,0.59,0.78 +kda_h64,cudnn,kda,16,8192,64,64,128,3.155,24.592,229,88,0.000000,20,4.09,2.71 +kda_h64,fla,kda,16,8192,64,64,128,21.419,84.140,34,26,0.000000,20,0.60,0.79 +kda_hon,cudnn_state_on,kda,4,2048,64,64,128,0.351,1.359,128,100,0.000000,20,2.30,3.07 +kda_hon,cudnn_state_on,kda,4,4096,64,64,128,0.673,2.682,134,101,0.000000,20,2.40,3.11 +kda_hon,cudnn_state_on,kda,4,8192,64,64,128,1.318,5.356,137,101,0.000000,20,2.45,3.11 +kda_hon,cudnn_state_on,kda,4,16384,64,64,128,2.601,10.668,139,101,0.000000,20,2.48,3.13 +kda_hon,cudnn_state_on,kda,4,32768,64,64,128,5.199,21.323,139,102,0.000000,20,2.48,3.13 +kda_hon,cudnn_state_on,kda,1,8192,64,64,128,0.653,2.522,69,54,0.000000,20,1.24,1.65 +kda_hon,cudnn_state_on,kda,2,8192,64,64,128,0.655,3.086,138,88,0.000000,20,2.46,2.70 +kda_hon,cudnn_state_on,kda,4,8192,64,64,128,1.319,5.356,137,101,0.000000,20,2.45,3.11 +kda_hon,cudnn_state_on,kda,8,8192,64,64,128,2.678,11.181,135,97,0.000000,20,2.41,2.98 +kda_hon,cudnn_state_on,kda,16,8192,64,64,128,4.802,20.399,150,106,0.000000,20,2.69,3.27 +kda_h64,flash_kda,kda,4,2048,64,64,128,0.432,0.000,104,0,0.000000,20,1.87,0.00 +kda_h64,flash_kda,kda,4,4096,64,64,128,0.848,0.000,106,0,0.000000,20,1.90,0.00 +kda_h64,flash_kda,kda,4,8192,64,64,128,1.688,0.000,107,0,0.000000,20,1.91,0.00 +kda_h64,flash_kda,kda,4,16384,64,64,128,3.353,0.000,108,0,0.000000,20,1.93,0.00 +kda_h64,flash_kda,kda,4,32768,64,64,128,6.690,0.000,108,0,0.000000,20,1.93,0.00 +kda_h64,flash_kda,kda,1,8192,64,64,128,0.898,0.000,50,0,0.000000,20,0.90,0.00 +kda_h64,flash_kda,kda,2,8192,64,64,128,1.070,0.000,84,0,0.000000,20,1.51,0.00 +kda_h64,flash_kda,kda,4,8192,64,64,128,1.686,0.000,107,0,0.000000,20,1.92,0.00 +kda_h64,flash_kda,kda,8,8192,64,64,128,3.349,0.000,108,0,0.000000,20,1.93,0.00 +kda_h64,flash_kda,kda,16,8192,64,64,128,6.342,0.000,114,0,0.000000,20,2.04,0.00 diff --git a/benchmark/linear_attention/results/kda/gb300/kda_fixed_batch_bw.png b/benchmark/linear_attention/results/kda/gb300/kda_fixed_batch_bw.png index 0f39146ce..ac8f9f107 100644 Binary files a/benchmark/linear_attention/results/kda/gb300/kda_fixed_batch_bw.png and b/benchmark/linear_attention/results/kda/gb300/kda_fixed_batch_bw.png differ diff --git a/benchmark/linear_attention/results/kda/gb300/kda_fixed_batch_flops.png b/benchmark/linear_attention/results/kda/gb300/kda_fixed_batch_flops.png index 331d8c900..610a93cca 100644 Binary files a/benchmark/linear_attention/results/kda/gb300/kda_fixed_batch_flops.png and b/benchmark/linear_attention/results/kda/gb300/kda_fixed_batch_flops.png differ diff --git a/benchmark/linear_attention/results/kda/gb300/kda_fixed_seq_bw.png b/benchmark/linear_attention/results/kda/gb300/kda_fixed_seq_bw.png index bf4badc36..3cacec1e4 100644 Binary files a/benchmark/linear_attention/results/kda/gb300/kda_fixed_seq_bw.png and b/benchmark/linear_attention/results/kda/gb300/kda_fixed_seq_bw.png differ diff --git a/benchmark/linear_attention/results/kda/gb300/kda_fixed_seq_flops.png b/benchmark/linear_attention/results/kda/gb300/kda_fixed_seq_flops.png index 91a081c06..344c074d4 100644 Binary files a/benchmark/linear_attention/results/kda/gb300/kda_fixed_seq_flops.png and b/benchmark/linear_attention/results/kda/gb300/kda_fixed_seq_flops.png differ diff --git a/python/cudnn/frost/tile_dsl/pointwise.py b/python/cudnn/frost/tile_dsl/pointwise.py index 8cb8273f0..81c89ad1c 100644 --- a/python/cudnn/frost/tile_dsl/pointwise.py +++ b/python/cudnn/frost/tile_dsl/pointwise.py @@ -309,12 +309,72 @@ def mul_f16x2(lhs: cutlass.Int32, rhs: cutlass.Int32, input_dtype: cutlass.Const @cute.jit def sigmoid_f16x2(logit_pair: cutlass.Int32, input_dtype: cutlass.Constexpr): """Sigmoid of a packed 16-bit logit pair, returned as two fp32 values.""" - logit_vec_f32 = cutlass.Vector.from_elements((logit_pair,), cutlass.Int32).bitcast(input_dtype).to(cutlass.Float32) + return sigmoid2(logit_vec_f32[0], logit_vec_f32[1]) + + +L2_NORM_EPS = 1.0e-12 + + +@cute.jit +def lane_group_sum(value: cutlass.Float32, lanes: cutlass.Constexpr[int]) -> cutlass.Float32: + """Sum ``value`` across a power-of-two group of consecutive lanes via + butterfly shuffles (every lane ends up holding the group total).""" + offset = lanes // 2 + while offset >= 1: + value = value + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, value, offset, 31, kind=nvvm.Shfl.BFLY)) + offset = offset // 2 + return value + + +@cute.jit +def l2norm_inv(sum_sq: cutlass.Float32) -> cutlass.Float32: + """Inverse L2 norm with the shared epsilon floor: rows at or below the + floor normalize by ``1 / L2_NORM_EPS`` instead of dividing by zero.""" + norm_floor_sq = cutlass.Float32(L2_NORM_EPS * L2_NORM_EPS) + return cute.math.rsqrt(cute.math.max(sum_sq, norm_floor_sq), fastmath=True) + + +@cute.jit +def sigmoid(x: cutlass.Float32) -> cutlass.Float32: + """sigmoid(x) via the tanh identity (single MUFU on Blackwell).""" half = cutlass.Float32(0.5) - value0 = cute.math.tanh(logit_vec_f32[0] * half, approx=True) * half + half - value1 = cute.math.tanh(logit_vec_f32[1] * half, approx=True) * half + half - return value0, value1 + return cute.math.tanh(x * half, approx=True) * half + half + + +@cute.jit +def sigmoid2(x_lo, x_hi): + """``(sigmoid(x_lo), sigmoid(x_hi))`` via the tanh identity, with the + halving and the scale-bias folded into one FMUL2 and one FFMA2.""" + half = opaque_f32_zero() + cutlass.Float32(0.5) + scaled_lo, scaled_hi = fmul2(x_lo, x_hi, half, half) + tanh_lo = cute.math.tanh(scaled_lo, approx=True) + tanh_hi = cute.math.tanh(scaled_hi, approx=True) + return ffma2(tanh_lo, tanh_hi, half, half, half, half) + + +@cute.jit +def softplus(x: cutlass.Float32) -> cutlass.Float32: + """log(1 + exp(x)) with the linear tail (x > 20 returns x: exp saturates + fp32 there and log1p(exp(x)) == x to fp32 precision).""" + result = x + if x < cutlass.Float32(20.0): + result = cute.math.log(cutlass.Float32(1.0) + cute.math.exp(x, fastmath=True), fastmath=True) + return result + + +@cute.jit +def softplus2(x_lo, x_hi): + """``(softplus(x_lo), softplus(x_hi))`` with the ``1 + exp`` step packed + into one FADD2 and the linear tail applied as a select.""" + one = cutlass.Float32(1.0) + tail = cutlass.Float32(20.0) + exp_lo = cute.math.exp(x_lo, fastmath=True) + exp_hi = cute.math.exp(x_hi, fastmath=True) + sum_lo, sum_hi = fadd2(exp_lo, exp_hi, one, one) + log_lo = cute.math.log(sum_lo, fastmath=True) + log_hi = cute.math.log(sum_hi, fastmath=True) + return (log_lo if x_lo < tail else x_lo), (log_hi if x_hi < tail else x_hi) @cute.jit diff --git a/python/cudnn/frost/tile_dsl/swizzle.py b/python/cudnn/frost/tile_dsl/swizzle.py index b660b95f5..cfedb5ae1 100644 --- a/python/cudnn/frost/tile_dsl/swizzle.py +++ b/python/cudnn/frost/tile_dsl/swizzle.py @@ -60,9 +60,3 @@ def swizzle_lin_128b(lin, *, row_stride_log2: cutlass.Constexpr[int], elem_bytes shift = cutlass.const_expr(row_stride_log2 - chunk_log2) mask = cutlass.const_expr(0x7 << chunk_log2) return lin ^ ((lin >> shift) & mask) - - -@cute.jit -def swizzle_lin_S(lin, *, bbits: cutlass.Constexpr[int], mbase: cutlass.Constexpr[int], sshift: cutlass.Constexpr[int]): - yyy = (lin >> cutlass.const_expr(mbase + sshift)) & cutlass.const_expr((1 << bbits) - 1) - return lin ^ (yyy << cutlass.const_expr(mbase)) diff --git a/python/cudnn/linear_attention/cutile/gdn_engine.py b/python/cudnn/linear_attention/cutile/gdn_engine.py index 9980868ed..d1681a510 100644 --- a/python/cudnn/linear_attention/cutile/gdn_engine.py +++ b/python/cudnn/linear_attention/cutile/gdn_engine.py @@ -94,7 +94,7 @@ def __init__(self, graph): ("wy_dg", layout.add(total * HV * 4), f32, (total, HV)), ] - self.ws_bytes = layout.size + self.workspace_size = layout.size self.carve_names = [name for name, _off, _dtype, _shape in regions] self.carve = carve_plan(self.plan_name, [(off, dtype, shape) for _name, off, dtype, shape in regions]) self.expect = expect_table(node) @@ -120,7 +120,7 @@ def __init__(self, graph): self.indices = None def get_workspace_size(self) -> int: - return self.ws_bytes + return self.workspace_size def execute(self, graph, variant_pack, ctx) -> None: if self.ports is None: @@ -132,8 +132,8 @@ def execute(self, graph, variant_pack, ctx) -> None: check_layouts_compact(self.plan_name, self.expect, self.names, views) nb = dict(zip(self.names, views)) stream = ctx.stream if ctx.stream is not None else 0 - ws = Workspace.over(variant_pack, self.ws_bytes, self.plan_name) - region = dict(zip(self.carve_names, ws.carve(self.carve))) + workspace = Workspace.over(variant_pack, self.workspace_size, self.plan_name) + region = dict(zip(self.carve_names, workspace.carve(self.carve))) self.common.build_chunk_table( region["chunk_table"], region["chunk_count"], diff --git a/python/cudnn/linear_attention/cutile/kda_engine.py b/python/cudnn/linear_attention/cutile/kda_engine.py index 9b8f60d3e..f6db5f999 100644 --- a/python/cudnn/linear_attention/cutile/kda_engine.py +++ b/python/cudnn/linear_attention/cutile/kda_engine.py @@ -108,7 +108,7 @@ def __init__(self, graph): if to_buffer_dtype(g.get_data_type()) != f32: regions.append(("dg_cum", layout.add(total * HV * K * 4), f32, (total, HV, K))) - self.ws_bytes = layout.size + self.workspace_size = layout.size self.carve_names = [name for name, _off, _dtype, _shape in regions] self.carve = carve_plan(self.plan_name, [(off, dtype, shape) for _name, off, dtype, shape in regions]) self.expect = expect_table(node) @@ -146,7 +146,7 @@ def __init__(self, graph): self.indices = None def get_workspace_size(self) -> int: - return self.ws_bytes + return self.workspace_size def execute(self, graph, variant_pack, ctx) -> None: if self.ports is None: @@ -158,8 +158,8 @@ def execute(self, graph, variant_pack, ctx) -> None: check_layouts_compact(self.plan_name, self.expect, self.names, views) nb = dict(zip(self.names, views)) stream = ctx.stream if ctx.stream is not None else 0 - ws = Workspace.over(variant_pack, self.ws_bytes, self.plan_name) - region = dict(zip(self.carve_names, ws.carve(self.carve))) + workspace = Workspace.over(variant_pack, self.workspace_size, self.plan_name) + region = dict(zip(self.carve_names, workspace.carve(self.carve))) self.common.build_chunk_table( region["chunk_table"], region["chunk_count"], diff --git a/python/cudnn/linear_attention/frost/common/elementwise.py b/python/cudnn/linear_attention/frost/common/elementwise.py deleted file mode 100644 index ad1b5c04f..000000000 --- a/python/cudnn/linear_attention/frost/common/elementwise.py +++ /dev/null @@ -1,48 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Element-wise device helpers shared by the FROST LA kernels: a lane-group -butterfly reduction, the inverse L2 norm with its epsilon floor, and the -sigmoid / softplus activations behind the safe gate and beta.""" - -import cutlass -import cutlass.cute as cute -import cutlass.experimental.primitives as nvvm - -L2_NORM_EPS = 1.0e-12 - - -@cute.jit -def lane_group_sum(value: cutlass.Float32, lanes: cutlass.Constexpr[int]) -> cutlass.Float32: - """Sum ``value`` across a power-of-two group of consecutive lanes via - butterfly shuffles (every lane ends up holding the group total).""" - offset = lanes // 2 - while offset >= 1: - value = value + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, value, offset, 31, kind=nvvm.Shfl.BFLY)) - offset = offset // 2 - return value - - -@cute.jit -def l2norm_inv(sum_sq: cutlass.Float32) -> cutlass.Float32: - """Inverse L2 norm with the shared epsilon floor: rows at or below the - floor normalize by ``1 / L2_NORM_EPS`` instead of dividing by zero.""" - norm_floor_sq = cutlass.Float32(L2_NORM_EPS * L2_NORM_EPS) - return cute.math.rsqrt(cute.math.max(sum_sq, norm_floor_sq), fastmath=True) - - -@cute.jit -def sigmoid(x: cutlass.Float32) -> cutlass.Float32: - """sigmoid(x) via the tanh identity (single MUFU on Blackwell).""" - half = cutlass.Float32(0.5) - return cute.math.tanh(x * half, approx=True) * half + half - - -@cute.jit -def softplus(x: cutlass.Float32) -> cutlass.Float32: - """log(1 + exp(x)) with the linear tail (x > 20 returns x: exp saturates - fp32 there and log1p(exp(x)) == x to fp32 precision).""" - result = x - if x < cutlass.Float32(20.0): - result = cute.math.log(cutlass.Float32(1.0) + cute.math.exp(x, fastmath=True), fastmath=True) - return result diff --git a/python/cudnn/linear_attention/frost/common/gate_bwd.py b/python/cudnn/linear_attention/frost/common/gate_bwd.py index e86550385..da44ffdd8 100644 --- a/python/cudnn/linear_attention/frost/common/gate_bwd.py +++ b/python/cudnn/linear_attention/frost/common/gate_bwd.py @@ -32,9 +32,7 @@ from cutlass.cute.runtime import from_dlpack from cudnn.frost.buffers import data_ptr -from cudnn.frost.tile_dsl.pointwise import fadd2, ffma2, fmul2 - -from .elementwise import lane_group_sum, sigmoid, softplus +from cudnn.frost.tile_dsl.pointwise import fadd2, ffma2, fmul2, lane_group_sum, sigmoid, sigmoid2, softplus GATE_BWD_BLOCKS = 128 # channel-gate token stripes (partials carve = 128 * HO * 128 fp32) SCALAR_BLOCK_CAP = 8192 # scalar-gate stripe ceiling @@ -108,14 +106,14 @@ def scalar_gate_bwd_finish_kernel( pairwise and one butterfly tree across the 32 lanes — a fixed-shape bracketing regardless of the stripe count.""" bid = cute.arch.block_idx() - lane = cutlass.Int32(cute.arch.thread_idx()[0]) + lane_idx = cutlass.Int32(cute.arch.thread_idx()[0]) h = cutlass.Int32(bid[0]) a8 = cutlass.Array(cutlass.Float32, 8) d8 = cutlass.Array(cutlass.Float32, 8) for j in cutlass.range_constexpr(8): a8[j] = cutlass.Float32(0.0) d8[j] = cutlass.Float32(0.0) - s = lane + s = lane_idx while s + cutlass.Int32(224) < n_blocks: for j in cutlass.range_constexpr(8): idx = (s + cutlass.Int32(32 * j)) * h_o + h @@ -135,7 +133,7 @@ def scalar_gate_bwd_finish_kernel( acc_a, acc_dt = fadd2(q0a, q0d, q1a, q1d) acc_a = lane_group_sum(acc_a, 32) acc_dt = lane_group_sum(acc_dt, 32) - if lane == 0: + if lane_idx == 0: mDA[h] = acc_a mDDt[h] = acc_dt @@ -166,8 +164,8 @@ def channel_gate_bwd_partial_kernel( g_blk = cutlass.Int32(bid[0]) h = cutlass.Int32(bid[1]) wrp = tidx // cutlass.Int32(32) - lane = tidx % cutlass.Int32(32) - d0 = lane * cutlass.Int32(4) + lane_idx = tidx % cutlass.Int32(32) + d0 = lane_idx * cutlass.Int32(4) exp_a = cute.math.exp(mALog[h], fastmath=True) bias = cutlass.Array(cutlass.Float32, 4) for q in cutlass.range_constexpr(4): @@ -209,8 +207,7 @@ def channel_gate_bwd_partial_kernel( one = cutlass.Float32(1.0) y_lo, y_hi = fadd2(gvv[i], gvv[j], bias[i], bias[j]) z_lo, z_hi = fmul2(exp_a, exp_a, y_lo, y_hi) - sig_lo = sigmoid(z_lo) - sig_hi = sigmoid(z_hi) + sig_lo, sig_hi = sigmoid2(z_lo, z_hi) c_lo, c_hi = fadd2(one, one, -sig_lo, -sig_hi) k_lo, k_hi = fmul2(sig_lo, sig_hi, c_lo, c_hi) b_lo, b_hi = fmul2(dgv[i], dgv[j], lower_bound, lower_bound) @@ -260,7 +257,7 @@ def channel_gate_bwd_finish_kernel( h = cutlass.Int32(bid[0]) d = tidx wrp = tidx // cutlass.Int32(32) - lane = tidx % cutlass.Int32(32) + lane_idx = tidx % cutlass.Int32(32) a8 = cutlass.Array(cutlass.Float32, 8) d8 = cutlass.Array(cutlass.Float32, 8) for j in cutlass.range_constexpr(8): @@ -278,7 +275,7 @@ def channel_gate_bwd_finish_kernel( mDDt[h, d] = col_dt sWa = cutlass.Array(cutlass.Float32, 4, space=cutlass.AddressSpace.smem, alignment=16) va = lane_group_sum(col_a, 32) - if lane == 0: + if lane_idx == 0: sWa[wrp] = va nvvm.barrier_cta_sync() if tidx == 0: diff --git a/python/cudnn/linear_attention/frost/common/l2norm.py b/python/cudnn/linear_attention/frost/common/l2norm.py index 795662ed0..1dec6b83d 100644 --- a/python/cudnn/linear_attention/frost/common/l2norm.py +++ b/python/cudnn/linear_attention/frost/common/l2norm.py @@ -19,12 +19,11 @@ from cutlass.cute.runtime import from_dlpack from cudnn.frost.buffers import data_ptr -from cudnn.frost.tile_dsl.pointwise import f16x2_to_f32, ffma2, fmul2, fp32_to_fp16 -from .elementwise import l2norm_inv, lane_group_sum +from cudnn.frost.tile_dsl.pointwise import f16x2_to_f32, ffma2, fmul2, fp32_to_fp16, l2norm_inv, lane_group_sum THREADS_PER_ROW = 16 ROWS_PER_CTA = 8 -FWD_LANES = 4 # fwd lanes per row: 32 elems/lane (4 x 128-bit loads), 2-step butterfly +FWD_LANES = 4 # fwd lanes per row: 32 elements/lane (4 x 128-bit loads), 2-step butterfly FWD_ROWS_PER_GROUP = 2 # fwd rows batched per lane group: 8 independent loads in flight per thread @@ -53,11 +52,11 @@ def l2norm_qk_kernel( bid = cute.arch.block_idx() tidx = cutlass.Int32(cute.arch.thread_idx()[0]) grp = tidx // cutlass.Int32(FWD_LANES) - lane = tidx % cutlass.Int32(FWD_LANES) + lane_idx = tidx % cutlass.Int32(FWD_LANES) row0 = (cutlass.Int32(bid[0]) * cutlass.Int32(128 // FWD_LANES) + grp) * cutlass.Int32(FWD_ROWS_PER_GROUP) - v0 = lane * cutlass.Int32(32) + v0 = lane_idx * cutlass.Int32(32) rows = [] - ws_addrs = [] + workspace_addrs = [] nrm_addrs = [] vals = [] for r in cutlass.range_constexpr(FWD_ROWS_PER_GROUP): @@ -67,22 +66,22 @@ def l2norm_qk_kernel( # [T, H, 128] strides while workspace rows sit at row * 128. Both # branches are traced, so the addresses have to exist beforehand. src_addr = cutlass.Int64(0) - ws_addr = cutlass.Int64(0) + workspace_addr = cutlass.Int64(0) nrm_addr = cutlass.Int64(0) if row_r < n_q_rows: t = row_r // h_q h = row_r % h_q - src_elems = t * cutlass.Int64(mQ.stride[0]) + h * cutlass.Int64(mQ.stride[1]) + cutlass.Int64(v0) - src_addr = mQ.iterator.toint() + src_elems * cutlass.Int64(2) - ws_addr = mQn.iterator.toint() + (cutlass.Int64(row_r) * cutlass.Int64(128) + cutlass.Int64(v0)) * cutlass.Int64(2) + src_elements = t * cutlass.Int64(mQ.stride[0]) + h * cutlass.Int64(mQ.stride[1]) + cutlass.Int64(v0) + src_addr = mQ.iterator.toint() + src_elements * cutlass.Int64(2) + workspace_addr = mQn.iterator.toint() + (cutlass.Int64(row_r) * cutlass.Int64(128) + cutlass.Int64(v0)) * cutlass.Int64(2) nrm_addr = mInvQ.iterator.toint() + cutlass.Int64(row_r) * cutlass.Int64(4) else: k_row = row_r - n_q_rows t = k_row // h_k h = k_row % h_k - src_elems = t * cutlass.Int64(mK.stride[0]) + h * cutlass.Int64(mK.stride[1]) + cutlass.Int64(v0) - src_addr = mK.iterator.toint() + src_elems * cutlass.Int64(2) - ws_addr = mKn.iterator.toint() + (cutlass.Int64(k_row) * cutlass.Int64(128) + cutlass.Int64(v0)) * cutlass.Int64(2) + src_elements = t * cutlass.Int64(mK.stride[0]) + h * cutlass.Int64(mK.stride[1]) + cutlass.Int64(v0) + src_addr = mK.iterator.toint() + src_elements * cutlass.Int64(2) + workspace_addr = mKn.iterator.toint() + (cutlass.Int64(k_row) * cutlass.Int64(128) + cutlass.Int64(v0)) * cutlass.Int64(2) nrm_addr = mInvK.iterator.toint() + cutlass.Int64(k_row) * cutlass.Int64(4) chunks = [] for c in cutlass.range_constexpr(4): @@ -97,7 +96,7 @@ def l2norm_qk_kernel( f6, f7 = f16x2_to_f32(w3, dtype=mQ.element_type) chunks.append((f0, f1, f2, f3, f4, f5, f6, f7)) rows.append(row) - ws_addrs.append(ws_addr) + workspace_addrs.append(workspace_addr) nrm_addrs.append(nrm_addr) vals.append(chunks) for r in cutlass.range_constexpr(FWD_ROWS_PER_GROUP): @@ -119,9 +118,9 @@ def l2norm_qk_kernel( w3 = fp32_to_fp16(s6, s7, dtype=mQ.element_type) inline_ptx( "st.global.v4.b32 [$0], {$1, $2, $3, $4};", - read_only_args=[ws_addrs[r] + cutlass.Int64(16 * c), w0, w1, w2, w3], + read_only_args=[workspace_addrs[r] + cutlass.Int64(16 * c), w0, w1, w2, w3], ) - if lane == cutlass.Int32(0): + if lane_idx == cutlass.Int32(0): inline_ptx("st.global.f32 [$0], $1;", read_only_args=[nrm_addrs[r], inv]) @@ -147,26 +146,26 @@ def l2norm_qk_bwd_kernel( bid = cute.arch.block_idx() tidx = cutlass.Int32(cute.arch.thread_idx()[0]) row = cutlass.Int32(bid[0]) * cutlass.Int32(ROWS_PER_CTA) + tidx // cutlass.Int32(THREADS_PER_ROW) - lane = tidx % cutlass.Int32(THREADS_PER_ROW) + lane_idx = tidx % cutlass.Int32(THREADS_PER_ROW) if row < n_rows: - v0 = lane * cutlass.Int32(8) + v0 = lane_idx * cutlass.Int32(8) grad_addr = cutlass.Int64(0) - ws_addr = cutlass.Int64(0) + workspace_addr = cutlass.Int64(0) nrm_addr = cutlass.Int64(0) if row < n_q_rows: t = row // h_q h = row % h_q - grad_elems = t * cutlass.Int64(mDq.stride[0]) + h * cutlass.Int64(mDq.stride[1]) + cutlass.Int64(v0) - grad_addr = mDq.iterator.toint() + grad_elems * cutlass.Int64(2) - ws_addr = mQn.iterator.toint() + (cutlass.Int64(row) * cutlass.Int64(128) + cutlass.Int64(v0)) * cutlass.Int64(2) + grad_elements = t * cutlass.Int64(mDq.stride[0]) + h * cutlass.Int64(mDq.stride[1]) + cutlass.Int64(v0) + grad_addr = mDq.iterator.toint() + grad_elements * cutlass.Int64(2) + workspace_addr = mQn.iterator.toint() + (cutlass.Int64(row) * cutlass.Int64(128) + cutlass.Int64(v0)) * cutlass.Int64(2) nrm_addr = mInvQ.iterator.toint() + cutlass.Int64(row) * cutlass.Int64(4) else: k_row = row - n_q_rows t = k_row // h_k h = k_row % h_k - grad_elems = t * cutlass.Int64(mDk.stride[0]) + h * cutlass.Int64(mDk.stride[1]) + cutlass.Int64(v0) - grad_addr = mDk.iterator.toint() + grad_elems * cutlass.Int64(2) - ws_addr = mKn.iterator.toint() + (cutlass.Int64(k_row) * cutlass.Int64(128) + cutlass.Int64(v0)) * cutlass.Int64(2) + grad_elements = t * cutlass.Int64(mDk.stride[0]) + h * cutlass.Int64(mDk.stride[1]) + cutlass.Int64(v0) + grad_addr = mDk.iterator.toint() + grad_elements * cutlass.Int64(2) + workspace_addr = mKn.iterator.toint() + (cutlass.Int64(k_row) * cutlass.Int64(128) + cutlass.Int64(v0)) * cutlass.Int64(2) nrm_addr = mInvK.iterator.toint() + cutlass.Int64(k_row) * cutlass.Int64(4) gw0, gw1, gw2, gw3 = inline_ptx( "ld.global.v4.b32 {$0, $1, $2, $3}, [$4];", @@ -180,7 +179,7 @@ def l2norm_qk_bwd_kernel( nw0, nw1, nw2, nw3 = inline_ptx( "ld.global.v4.b32 {$0, $1, $2, $3}, [$4];", write_only_types=[cutlass.Int32, cutlass.Int32, cutlass.Int32, cutlass.Int32], - read_only_args=[ws_addr], + read_only_args=[workspace_addr], ) n0, n1 = f16x2_to_f32(nw0, dtype=mQn.element_type) n2, n3 = f16x2_to_f32(nw1, dtype=mQn.element_type) diff --git a/python/cudnn/linear_attention/frost/common/split_k.py b/python/cudnn/linear_attention/frost/common/split_k.py index 2127b4499..c350b3fa5 100644 --- a/python/cudnn/linear_attention/frost/common/split_k.py +++ b/python/cudnn/linear_attention/frost/common/split_k.py @@ -16,22 +16,23 @@ Work-item table row (``WORK_ITEM_FIELDS`` x int32, chunk units):: - [batch_idx, head_idx, wstart, wend, cstart, cend, batch_start, batch_end] + [batch_idx, head_idx, write_start, write_end, compute_start, compute_end, batch_start, batch_end] ``batch_start``/``batch_end`` are the token bounds ``cu_seqlens[b]`` / ``cu_seqlens[b+1]`` denormalized into the row: decode reads one 32-byte vectorizable row instead of chasing a dependent ``cu_seqlens`` load pair. -The item OWNS (writes outputs for) chunks ``[wstart, wend)``. The forward -kernel COMPUTES chunks ``[cstart, wend)`` — ``[cstart, wstart)`` is the -left warmup that rebuilds the incoming state from zero (accurate to -``2^log2_threshold`` because the gate decay over the window saturates). -The backward kernel computes ``[wstart, cend)`` — ``[wend, cend)`` is the -right warmup for the reverse dstate recurrence (the forward states come exactly -from the per-chunk state checkpoints). ``cstart == 0`` items seed the true -initial state; ``cend == num_chunks`` items seed the true ``d_final_state`` -— so the un-cut degenerate item ``(0, nc, 0, nc)`` reproduces the serial -kernel exactly. +The item OWNS (writes outputs for) chunks ``[write_start, write_end)``. +The forward kernel COMPUTES chunks ``[compute_start, write_end)`` — +``[compute_start, write_start)`` is the left warmup that rebuilds the +incoming state from zero (accurate to ``2^log2_threshold`` because the +gate decay over the window saturates). The backward kernel computes +``[write_start, compute_end)`` — ``[write_end, compute_end)`` is the +right warmup for the reverse dstate recurrence (the forward states come +exactly from the per-chunk state checkpoints). ``compute_start == 0`` +items seed the true initial state; ``compute_end == num_chunks`` items +seed the true ``d_final_state`` — so the un-cut degenerate item +``(0, nc, 0, nc)`` reproduces the serial kernel exactly. Piece choice per (batch, head): spans never exceed ``ideal_chunks`` (total work / SM count), so outlier-long sequences are always cut down to the @@ -57,11 +58,11 @@ ``item_scratch``. 3. order (:func:`order_body`, hosted by each kernel module's prologue kernel alongside its TMA-descriptor build — one launch for - both): bitonic-sort the items into ``work_items``, longest ``[cstart, - cend)`` first, so the ticket scheduler consumes them in LPT order — - the makespan tail is set by whatever starts last, so the big items - must go first. This is what keeps ragged varlen batches balanced - without cutting them. + both): bitonic-sort the items into ``work_items``, longest + ``[compute_start, compute_end)`` first, so the ticket scheduler + consumes them in LPT order — the makespan tail is set by whatever + starts last, so the big items must go first. This is what keeps + ragged varlen batches balanced without cutting them. The order body also zeroes the main kernels' scheduler ticket rings (dirty on exit), and with ``split=False`` it replaces the whole pipeline: @@ -84,7 +85,7 @@ from cudnn.frost.buffers import data_ptr -from .elementwise import softplus +from cudnn.frost.tile_dsl.pointwise import sigmoid, softplus WORK_ITEM_FIELDS = 8 WARMUP_CAP_CHUNKS = 32 # hard warmup cap: a cut must saturate within one warp of chunks per side @@ -101,8 +102,8 @@ P_BELOW = 8 # how far below the ideal-cap floor the fill-regime search may go ORDER_THREADS = 1024 -ORDER_ELEMS = 4 -ORDER_CAPACITY = ORDER_THREADS * ORDER_ELEMS # sort capacity (32 KB SMEM); past this the device-side branch copies through unsorted +ORDER_ELEMENTS = 4 +ORDER_CAPACITY = ORDER_THREADS * ORDER_ELEMENTS # sort capacity (32 KB SMEM); past this the device-side branch copies through unsorted DEFAULT_LOG2_THRESHOLD = -10.0 / math.log(2.0) # e^-10, in log2 units RCP_LN2 = 1.4426950408889634 # 1/ln(2): natural-log gates -> the scan's log2 domain @@ -139,37 +140,38 @@ def max_work_items(total_tokens: int, batch_size: int, n_heads_out: int, ideal_c def decode_work_item(cfg, tile_idx, mWorkItems): """Tile decode shared by every warp body of the main kernels: read the work-item row (an uncut table row IS the whole sequence). Returns - ``(batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, - wstart, wend, cstart, cend)`` with chunk-unit bounds.""" + ``(batch_idx, head_idx, batch_start, batch_end, batch_seqlen, + batch_num_chunks, write_start, write_end, compute_start, + compute_end)`` with chunk-unit bounds.""" batch_idx = mWorkItems[tile_idx, 0] head_idx = mWorkItems[tile_idx, 1] - wstart = mWorkItems[tile_idx, 2] - wend = mWorkItems[tile_idx, 3] - cstart = mWorkItems[tile_idx, 4] - cend = mWorkItems[tile_idx, 5] + write_start = mWorkItems[tile_idx, 2] + write_end = mWorkItems[tile_idx, 3] + compute_start = mWorkItems[tile_idx, 4] + compute_end = mWorkItems[tile_idx, 5] batch_start = mWorkItems[tile_idx, 6] batch_end = mWorkItems[tile_idx, 7] - seqlen_b = batch_end - batch_start - num_chunks_b = cute.ceil_div(seqlen_b, cfg.b_t) - return batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend + batch_seqlen = batch_end - batch_start + batch_num_chunks = cute.ceil_div(batch_seqlen, cfg.b_t) + return batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end @cute.jit -def emit_item(mWorkItems, mCount, batch_idx, head_idx, wstart, wend, cstart, cend, batch_start, batch_end): +def emit_item(mWorkItems, mCount, batch_idx, head_idx, write_start, write_end, compute_start, compute_end, batch_start, batch_end): count_addr = mCount.iterator.toint() - wi = inline_ptx( + slot = inline_ptx( "atom.global.add.s32 {$w0}, [{$r0}], 1;", write_only_types=[cutlass.Int32], read_only_args=[count_addr], ) - mWorkItems[wi, 0] = batch_idx - mWorkItems[wi, 1] = head_idx - mWorkItems[wi, 2] = wstart - mWorkItems[wi, 3] = wend - mWorkItems[wi, 4] = cstart - mWorkItems[wi, 5] = cend - mWorkItems[wi, 6] = batch_start - mWorkItems[wi, 7] = batch_end + mWorkItems[slot, 0] = batch_idx + mWorkItems[slot, 1] = head_idx + mWorkItems[slot, 2] = write_start + mWorkItems[slot, 3] = write_end + mWorkItems[slot, 4] = compute_start + mWorkItems[slot, 5] = compute_end + mWorkItems[slot, 6] = batch_start + mWorkItems[slot, 7] = batch_end @cute.jit @@ -184,12 +186,12 @@ def clamped_log2(log_gate: cutlass.Constexpr[bool], gate_val: cutlass.Float32) - @cute.jit -def piece_choice(overhead_chunks: cutlass.Constexpr[int], num_chunks_b, n_tiles, num_sms, ideal_chunks): +def piece_choice(overhead_chunks: cutlass.Constexpr[int], batch_num_chunks, n_tiles, num_sms, ideal_chunks): """Per-tile piece choice. Returns ``(span, num_blocks)``.""" # even spread: spans never exceed ideal_chunks (total work / SM count) - p_hi = num_chunks_b if num_chunks_b < cutlass.Int32(MAX_BLOCKS) else cutlass.Int32(MAX_BLOCKS) + p_hi = batch_num_chunks if batch_num_chunks < cutlass.Int32(MAX_BLOCKS) else cutlass.Int32(MAX_BLOCKS) p_hi = p_hi if p_hi > 0 else cutlass.Int32(1) - p = (num_chunks_b + ideal_chunks - cutlass.Int32(1)) // ideal_chunks + p = (batch_num_chunks + ideal_chunks - cutlass.Int32(1)) // ideal_chunks p = p if p > 0 else cutlass.Int32(1) p = p if p < p_hi else p_hi if n_tiles < cutlass.Int32(2) * num_sms: @@ -199,40 +201,40 @@ def piece_choice(overhead_chunks: cutlass.Constexpr[int], num_chunks_b, n_tiles, p_start = p_start if p_start > 0 else cutlass.Int32(1) best = cutlass.Int32(2147483647) for dp in cutlass.range_constexpr(P_WINDOW): - cand = p_start + cutlass.Int32(dp) - cand = cand if cand < p_hi else p_hi - span_c = (num_chunks_b + cand - cutlass.Int32(1)) // cand - waves = (n_tiles * cand + num_sms - cutlass.Int32(1)) // num_sms - est = waves * (span_c + cutlass.Int32(overhead_chunks)) - hit = est < best - best = est if hit else best - p = cand if hit else p + candidate = p_start + cutlass.Int32(dp) + candidate = candidate if candidate < p_hi else p_hi + span_c = (batch_num_chunks + candidate - cutlass.Int32(1)) // candidate + waves = (n_tiles * candidate + num_sms - cutlass.Int32(1)) // num_sms + estimate = waves * (span_c + cutlass.Int32(overhead_chunks)) + hit = estimate < best + best = estimate if hit else best + p = candidate if hit else p # the estimate flatters marginal cuts; cut only on a clear margin over uncut - est1 = ((n_tiles + num_sms - cutlass.Int32(1)) // num_sms) * (num_chunks_b + cutlass.Int32(overhead_chunks)) - if cutlass.Int32(4) * best > cutlass.Int32(3) * est1: + uncut_estimate = ((n_tiles + num_sms - cutlass.Int32(1)) // num_sms) * (batch_num_chunks + cutlass.Int32(overhead_chunks)) + if cutlass.Int32(4) * best > cutlass.Int32(3) * uncut_estimate: p = cutlass.Int32(1) span = cutlass.Int32(0) num_blocks = cutlass.Int32(0) - if num_chunks_b > 0: - span = (num_chunks_b + p - cutlass.Int32(1)) // p - num_blocks = (num_chunks_b + span - cutlass.Int32(1)) // span + if batch_num_chunks > 0: + span = (batch_num_chunks + p - cutlass.Int32(1)) // p + num_blocks = (batch_num_chunks + span - cutlass.Int32(1)) // span return span, num_blocks @cute.jit def tile_spans(b_t: cutlass.Constexpr[int], overhead_chunks: cutlass.Constexpr[int], n_heads_out, n_tiles, num_sms, ideal_chunks, mCuSeqlens, tile): """Per-tile decode + piece choice. Returns ``(batch_idx, head_idx, - batch_start, batch_end, num_chunks_b, cv_base, span, num_blocks)``; - ``cv_base`` is the tile's row base in the GMEM chunk scratch.""" + batch_start, batch_end, batch_num_chunks, chunk_value_base, span, num_blocks)``; + ``chunk_value_base`` is the tile's row base in the GMEM chunk scratch.""" batch_idx = tile // n_heads_out head_idx = tile % n_heads_out batch_start = cutlass.Int32(mCuSeqlens[batch_idx]) batch_end = cutlass.Int32(mCuSeqlens[batch_idx + 1]) - seqlen_b = batch_end - batch_start - num_chunks_b = cute.ceil_div(seqlen_b, b_t) - cv_base = batch_start // cutlass.Int32(b_t) + batch_idx - span, num_blocks = piece_choice(overhead_chunks, num_chunks_b, n_tiles, num_sms, ideal_chunks) - return batch_idx, head_idx, batch_start, batch_end, num_chunks_b, cv_base, span, num_blocks + batch_seqlen = batch_end - batch_start + batch_num_chunks = cute.ceil_div(batch_seqlen, b_t) + chunk_value_base = batch_start // cutlass.Int32(b_t) + batch_idx + span, num_blocks = piece_choice(overhead_chunks, batch_num_chunks, n_tiles, num_sms, ideal_chunks) + return batch_idx, head_idx, batch_start, batch_end, batch_num_chunks, chunk_value_base, span, num_blocks @cute.jit @@ -280,7 +282,7 @@ def scan_kernel( head_idx = cutlass.Int32(bidx[1]) if cutlass.Int32(bidx[0]) == 0 and head_idx == 0 and tidx == 0: mCount[0] = cutlass.Int32(0) - lidx = tidx % cutlass.Int32(WARP_SIZE) + lane_idx = tidx % cutlass.Int32(WARP_SIZE) widx = tidx // cutlass.Int32(WARP_SIZE) row0 = (cutlass.Int32(bidx[0]) * cutlass.Int32(SCAN_WARPS) + widx) * cutlass.Int32(SCAN_ROWS_PER_WARP) @@ -295,10 +297,10 @@ def scan_kernel( batch_idx = lo batch_start = cutlass.Int32(mCuSeqlens[batch_idx]) batch_end = cutlass.Int32(mCuSeqlens[batch_idx + 1]) - num_chunks_b = cute.ceil_div(batch_end - batch_start, b_t) - cv_base = batch_start // cutlass.Int32(b_t) + batch_idx + batch_num_chunks = cute.ceil_div(batch_end - batch_start, b_t) + chunk_value_base = batch_start // cutlass.Int32(b_t) + batch_idx # the piece choice is computed per batch, NOT per row - span, num_blocks = piece_choice(overhead_chunks, num_chunks_b, n_tiles, num_sms, ideal_chunks) + span, num_blocks = piece_choice(overhead_chunks, batch_num_chunks, n_tiles, num_sms, ideal_chunks) for rr in cutlass.range_constexpr(SCAN_ROWS_PER_WARP): row = row0 + cutlass.Int32(rr) while (batch_idx + cutlass.Int32(1) < batch_size) and ( @@ -307,17 +309,17 @@ def scan_kernel( batch_idx = batch_idx + cutlass.Int32(1) batch_start = cutlass.Int32(mCuSeqlens[batch_idx]) batch_end = cutlass.Int32(mCuSeqlens[batch_idx + 1]) - num_chunks_b = cute.ceil_div(batch_end - batch_start, b_t) - cv_base = batch_start // cutlass.Int32(b_t) + batch_idx - span, num_blocks = piece_choice(overhead_chunks, num_chunks_b, n_tiles, num_sms, ideal_chunks) - c = row - cv_base - if (c >= 0) and (c < num_chunks_b) and (num_blocks > 1): + batch_num_chunks = cute.ceil_div(batch_end - batch_start, b_t) + chunk_value_base = batch_start // cutlass.Int32(b_t) + batch_idx + span, num_blocks = piece_choice(overhead_chunks, batch_num_chunks, n_tiles, num_sms, ideal_chunks) + c = row - chunk_value_base + if (c >= 0) and (c < batch_num_chunks) and (num_blocks > 1): if near_boundary(c, span if span > 0 else cutlass.Int32(1), num_blocks): # chunk value = max over channels of the per-channel # clamped-log2 sums (each lane owns a contiguous channel run) cpl = gate_channels // WARP_SIZE - row_elems = cutlass.Int32(mGate.stride[0]) - lane_base = cutlass.Int64(head_idx * cutlass.Int32(mGate.stride[1]) + lidx * cutlass.Int32(cpl)) + row_elements = cutlass.Int32(mGate.stride[0]) + lane_base = cutlass.Int64(head_idx * cutlass.Int32(mGate.stride[1]) + lane_idx * cutlass.Int32(cpl)) gate_addr = mGate.iterator.toint() + lane_base * cutlass.Int64(4) gate_ptr = mGate.iterator + lane_base a_exp = cutlass.Float32(1.0) @@ -326,7 +328,7 @@ def scan_kernel( # per-head rate + per-lane channel biases, fixed for the whole chunk a_exp = cute.math.exp2(mALog[head_idx].to(cutlass.Float32) * cutlass.Float32(RCP_LN2), fastmath=True) for q in cutlass.range_constexpr(cpl): - dt_vals[q] = mDtBias[head_idx, lidx * cutlass.Int32(cpl) + cutlass.Int32(q)].to(cutlass.Float32) + dt_vals[q] = mDtBias[head_idx, lane_idx * cutlass.Int32(cpl) + cutlass.Int32(q)].to(cutlass.Float32) ch_acc = cutlass.Array(cutlass.Float32, cpl, alignment=16) for q in cutlass.range_constexpr(cpl): ch_acc[q] = cutlass.Float32(0.0) @@ -335,7 +337,7 @@ def scan_kernel( pos = batch_start + c * cutlass.Int32(b_t) + cutlass.Int32(tt) inb = pos < batch_end pos_r = pos if inb else batch_start - grow = cutlass.Int64(pos_r) * cutlass.Int64(row_elems) + grow = cutlass.Int64(pos_r) * cutlass.Int64(row_elements) if cutlass.const_expr(cpl % 4 == 0): for q4 in cutlass.range_constexpr(cpl // 4): addr = gate_addr + (grow + cutlass.Int64(4 * q4)) * cutlass.Int64(4) @@ -347,9 +349,7 @@ def scan_kernel( for qq, gvq in enumerate((g0, g1, g2, g3)): q = 4 * q4 + qq if cutlass.const_expr(safe_gate): - # the main kernel's transform, in log2 domain: scale * sigmoid(exp(a_log) * (g + dt_bias)) - half = cutlass.Float32(0.5) - sig = cute.math.tanh(a_exp * (gvq + dt_vals[q]) * half, approx=True) * half + half + sig = sigmoid(a_exp * (gvq + dt_vals[q])) contrib = gate_scale_log2 * sig contrib = contrib if inb else cutlass.Float32(0.0) else: @@ -360,8 +360,7 @@ def scan_kernel( for q in cutlass.range_constexpr(cpl): gv = (gate_ptr + grow + cutlass.Int32(q)).load() if cutlass.const_expr(safe_gate): - half = cutlass.Float32(0.5) - sig = cute.math.tanh(a_exp * (gv + dt_vals[q]) * half, approx=True) * half + half + sig = sigmoid(a_exp * (gv + dt_vals[q])) contrib = gate_scale_log2 * sig contrib = contrib if inb else cutlass.Float32(0.0) else: @@ -374,8 +373,8 @@ def scan_kernel( for off in [1, 2, 4, 8, 16]: other = cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, m, off, 31, kind=nvvm.Shfl.BFLY)) m = m if m > other else other - if lidx == 0: - mChunkVals[cv_base + c, head_idx] = m + if lane_idx == 0: + mChunkVals[chunk_value_base + c, head_idx] = m @cute.kernel @@ -409,16 +408,15 @@ def scan_scalar_kernel( tidx = cutlass.Int32(tidx) if cutlass.Int32(bidx[0]) == 0 and cutlass.Int32(bidx[1]) == 0 and tidx == 0: mCount[0] = cutlass.Int32(0) - lidx = tidx % cutlass.Int32(WARP_SIZE) + lane_idx = tidx % cutlass.Int32(WARP_SIZE) widx = tidx // cutlass.Int32(WARP_SIZE) - h = cutlass.Int32(bidx[1]) * cutlass.Int32(WARP_SIZE) + lidx + h = cutlass.Int32(bidx[1]) * cutlass.Int32(WARP_SIZE) + lane_idx h_ok = h < n_heads_out h_r = h if h_ok else n_heads_out - cutlass.Int32(1) - a_l2 = cutlass.Float32(0.0) + a = cutlass.Float32(0.0) bias = cutlass.Float32(0.0) if cutlass.const_expr(safe_gate): - # per-head transform constants, fixed for the lane's whole sweep - a_l2 = -cute.math.exp2(mALog[h_r].to(cutlass.Float32) * cutlass.Float32(RCP_LN2), fastmath=True) * cutlass.Float32(RCP_LN2) + a = -cute.math.exp2(mALog[h_r].to(cutlass.Float32) * cutlass.Float32(RCP_LN2), fastmath=True) * cutlass.Float32(RCP_LN2) bias = mDtBias[h_r].to(cutlass.Float32) row0 = (cutlass.Int32(bidx[0]) * cutlass.Int32(SCAN_WARPS) + widx) * cutlass.Int32(SCAN_ROWS_PER_WARP) @@ -433,9 +431,9 @@ def scan_scalar_kernel( batch_idx = lo batch_start = cutlass.Int32(mCuSeqlens[batch_idx]) batch_end = cutlass.Int32(mCuSeqlens[batch_idx + 1]) - num_chunks_b = cute.ceil_div(batch_end - batch_start, b_t) - cv_base = batch_start // cutlass.Int32(b_t) + batch_idx - span, num_blocks = piece_choice(overhead_chunks, num_chunks_b, n_tiles, num_sms, ideal_chunks) + batch_num_chunks = cute.ceil_div(batch_end - batch_start, b_t) + chunk_value_base = batch_start // cutlass.Int32(b_t) + batch_idx + span, num_blocks = piece_choice(overhead_chunks, batch_num_chunks, n_tiles, num_sms, ideal_chunks) for rr in cutlass.range_constexpr(SCAN_ROWS_PER_WARP): row = row0 + cutlass.Int32(rr) while (batch_idx + cutlass.Int32(1) < batch_size) and ( @@ -444,11 +442,11 @@ def scan_scalar_kernel( batch_idx = batch_idx + cutlass.Int32(1) batch_start = cutlass.Int32(mCuSeqlens[batch_idx]) batch_end = cutlass.Int32(mCuSeqlens[batch_idx + 1]) - num_chunks_b = cute.ceil_div(batch_end - batch_start, b_t) - cv_base = batch_start // cutlass.Int32(b_t) + batch_idx - span, num_blocks = piece_choice(overhead_chunks, num_chunks_b, n_tiles, num_sms, ideal_chunks) - c = row - cv_base - if (c >= 0) and (c < num_chunks_b) and (num_blocks > 1): + batch_num_chunks = cute.ceil_div(batch_end - batch_start, b_t) + chunk_value_base = batch_start // cutlass.Int32(b_t) + batch_idx + span, num_blocks = piece_choice(overhead_chunks, batch_num_chunks, n_tiles, num_sms, ideal_chunks) + c = row - chunk_value_base + if (c >= 0) and (c < batch_num_chunks) and (num_blocks > 1): if near_boundary(c, span if span > 0 else cutlass.Int32(1), num_blocks): oob = cutlass.Float32(0.0) if cutlass.const_expr(log_gate) else cutlass.Float32(1.0) acc = cutlass.Float32(0.0) @@ -458,13 +456,13 @@ def scan_scalar_kernel( pos_r = pos if inb else batch_start gv = (mGate.iterator + cutlass.Int64(pos_r) * cutlass.Int64(mGate.stride[0]) + h_r).load() if cutlass.const_expr(safe_gate): - contrib = a_l2 * softplus(gv + bias) + contrib = a * softplus(gv + bias) acc = acc + (contrib if inb else cutlass.Float32(0.0)) else: gv = gv if inb else oob acc = acc + clamped_log2(log_gate, gv) if h_ok: - mChunkVals[cv_base + c, h] = acc + mChunkVals[chunk_value_base + c, h] = acc @cute.kernel @@ -489,73 +487,73 @@ def walk_kernel( bidx = cute.arch.block_idx()[0] tidx = cutlass.Int32(tidx) - batch_idx, head_idx, batch_start, batch_end, num_chunks_b, cv_base, span, num_blocks = tile_spans( + batch_idx, head_idx, batch_start, batch_end, batch_num_chunks, chunk_value_base, span, num_blocks = tile_spans( b_t, overhead_chunks, n_heads_out, n_tiles, num_sms, ideal_chunks, mCuSeqlens, cutlass.Int32(bidx) ) if num_blocks <= 1: # single piece: no cuts, nothing scanned if tidx == 0: - emit_item(mStaging, mCount, batch_idx, head_idx, cutlass.Int32(0), num_chunks_b, cutlass.Int32(0), num_chunks_b, batch_start, batch_end) + emit_item(mStaging, mCount, batch_idx, head_idx, cutlass.Int32(0), batch_num_chunks, cutlass.Int32(0), batch_num_chunks, batch_start, batch_end) else: - # packed per-boundary probe results: warm_b | warm_f << 8, 0 = no cut - sWarm = cutlass.Array(cutlass.Int32, MAX_BLOCKS, space=cutlass.AddressSpace.smem, alignment=16) - lidx = tidx % cutlass.Int32(WARP_SIZE) + # packed per-boundary probe results: warmup_before | warmup_after << 8, 0 = no cut + sWarmup = cutlass.Array(cutlass.Int32, MAX_BLOCKS, space=cutlass.AddressSpace.smem, alignment=16) + lane_idx = tidx % cutlass.Int32(WARP_SIZE) widx = tidx // cutlass.Int32(WARP_SIZE) big = cutlass.Int32(2 * WARMUP_CAP_CHUNKS) j = widx + cutlass.Int32(1) while j < num_blocks: - wend = j * span - # fwd warmup: smallest chunk suffix of [0, wend) that saturates - idx = wend - cutlass.Int32(1) - lidx + write_end = j * span + # fwd warmup: smallest chunk suffix of [0, write_end) that saturates + idx = write_end - cutlass.Int32(1) - lane_idx ok_b = idx >= 0 - v = mChunkVals[cv_base + (idx if ok_b else cutlass.Int32(0)), head_idx] + v = mChunkVals[chunk_value_base + (idx if ok_b else cutlass.Int32(0)), head_idx] acc = v if ok_b else cutlass.Float32(0.0) for off in [1, 2, 4, 8, 16]: o = cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, acc, off, 0, kind=nvvm.Shfl.UP)) - acc = acc + (o if lidx >= cutlass.Int32(off) else cutlass.Float32(0.0)) - cand = lidx + cutlass.Int32(1) if (ok_b and acc <= log2_thresh) else big + acc = acc + (o if lane_idx >= cutlass.Int32(off) else cutlass.Float32(0.0)) + candidate = lane_idx + cutlass.Int32(1) if (ok_b and acc <= log2_thresh) else big for off in [1, 2, 4, 8, 16]: - other = cutlass.Int32(nvvm.shfl_sync(0xFFFFFFFF, cand, off, 31, kind=nvvm.Shfl.BFLY)) - cand = cand if cand < other else other - warm_b = cand if cand < big else cutlass.Int32(0) - # bwd warmup: smallest chunk prefix of [wend, nc) that saturates - idx = wend + lidx - ok_f = idx < num_chunks_b - v = mChunkVals[cv_base + (idx if ok_f else cutlass.Int32(0)), head_idx] + other = cutlass.Int32(nvvm.shfl_sync(0xFFFFFFFF, candidate, off, 31, kind=nvvm.Shfl.BFLY)) + candidate = candidate if candidate < other else other + warmup_before = candidate if candidate < big else cutlass.Int32(0) + # bwd warmup: smallest chunk prefix of [write_end, nc) that saturates + idx = write_end + lane_idx + ok_f = idx < batch_num_chunks + v = mChunkVals[chunk_value_base + (idx if ok_f else cutlass.Int32(0)), head_idx] acc = v if ok_f else cutlass.Float32(0.0) for off in [1, 2, 4, 8, 16]: o = cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, acc, off, 0, kind=nvvm.Shfl.UP)) - acc = acc + (o if lidx >= cutlass.Int32(off) else cutlass.Float32(0.0)) - cand = lidx + cutlass.Int32(1) if (ok_f and acc <= log2_thresh) else big + acc = acc + (o if lane_idx >= cutlass.Int32(off) else cutlass.Float32(0.0)) + candidate = lane_idx + cutlass.Int32(1) if (ok_f and acc <= log2_thresh) else big for off in [1, 2, 4, 8, 16]: - other = cutlass.Int32(nvvm.shfl_sync(0xFFFFFFFF, cand, off, 31, kind=nvvm.Shfl.BFLY)) - cand = cand if cand < other else other - warm_f = cand if cand < big else cutlass.Int32(0) - if lidx == 0: - packed = warm_b + warm_f * cutlass.Int32(256) - packed = packed if (warm_b > 0 and warm_f > 0) else cutlass.Int32(0) - sWarm[j] = packed + other = cutlass.Int32(nvvm.shfl_sync(0xFFFFFFFF, candidate, off, 31, kind=nvvm.Shfl.BFLY)) + candidate = candidate if candidate < other else other + warmup_after = candidate if candidate < big else cutlass.Int32(0) + if lane_idx == 0: + packed = warmup_before + warmup_after * cutlass.Int32(256) + packed = packed if (warmup_before > 0 and warmup_after > 0) else cutlass.Int32(0) + sWarmup[j] = packed j = j + cutlass.Int32(WARPS) nvvm.barrier_cta_sync() if tidx == 0: - prev_cut = cutlass.Int32(0) # wstart of the open item, chunk units - cur_cstart = cutlass.Int32(0) # cstart of the open item, chunk units + prev_cut = cutlass.Int32(0) # write_start of the open item, chunk units + current_compute_start = cutlass.Int32(0) # compute_start of the open item, chunk units jj = cutlass.Int32(1) while jj < num_blocks: - r = sWarm[jj] + r = sWarmup[jj] if r != 0: - wend = jj * span - warm_b = r % cutlass.Int32(256) - warm_f = r // cutlass.Int32(256) - cend = wend + warm_f - cend = cend if cend < num_chunks_b else num_chunks_b - emit_item(mStaging, mCount, batch_idx, head_idx, prev_cut, wend, cur_cstart, cend, batch_start, batch_end) - cur_cstart = wend - warm_b - prev_cut = wend + write_end = jj * span + warmup_before = r % cutlass.Int32(256) + warmup_after = r // cutlass.Int32(256) + compute_end = write_end + warmup_after + compute_end = compute_end if compute_end < batch_num_chunks else batch_num_chunks + emit_item(mStaging, mCount, batch_idx, head_idx, prev_cut, write_end, current_compute_start, compute_end, batch_start, batch_end) + current_compute_start = write_end - warmup_before + prev_cut = write_end jj = jj + cutlass.Int32(1) - emit_item(mStaging, mCount, batch_idx, head_idx, prev_cut, num_chunks_b, cur_cstart, num_chunks_b, batch_start, batch_end) + emit_item(mStaging, mCount, batch_idx, head_idx, prev_cut, batch_num_chunks, current_compute_start, batch_num_chunks, batch_start, batch_end) @cute.jit @@ -566,8 +564,8 @@ def gen_item_bounds(b_t: cutlass.Constexpr[int], n_heads_out, mCuSeqlens, item): head_idx = item % n_heads_out batch_start = cutlass.Int32(mCuSeqlens[batch_idx]) batch_end = cutlass.Int32(mCuSeqlens[batch_idx + 1]) - num_chunks_b = cute.ceil_div(batch_end - batch_start, b_t) - return batch_idx, head_idx, batch_start, batch_end, num_chunks_b + batch_num_chunks = cute.ceil_div(batch_end - batch_start, b_t) + return batch_idx, head_idx, batch_start, batch_end, batch_num_chunks @cute.jit @@ -584,13 +582,13 @@ def write_item( """Final-table row ``dst`` from source item ``src``: the walk's staged row, or (``gen``) the synthesized uncut item ``(0, nc, 0, nc)``.""" if cutlass.const_expr(gen): - batch_idx, head_idx, batch_start, batch_end, num_chunks_b = gen_item_bounds(b_t, n_heads_out, mCuSeqlens, src) + batch_idx, head_idx, batch_start, batch_end, batch_num_chunks = gen_item_bounds(b_t, n_heads_out, mCuSeqlens, src) mWorkItems[dst, 0] = batch_idx mWorkItems[dst, 1] = head_idx mWorkItems[dst, 2] = cutlass.Int32(0) - mWorkItems[dst, 3] = num_chunks_b + mWorkItems[dst, 3] = batch_num_chunks mWorkItems[dst, 4] = cutlass.Int32(0) - mWorkItems[dst, 5] = num_chunks_b + mWorkItems[dst, 5] = batch_num_chunks mWorkItems[dst, 6] = batch_start mWorkItems[dst, 7] = batch_end else: @@ -601,10 +599,10 @@ def write_item( @cute.jit def order_body( gen: cutlass.Constexpr[bool], - has_sched: cutlass.Constexpr[bool], + has_scheduler: cutlass.Constexpr[bool], b_t: cutlass.Constexpr[int], n_threads: cutlass.Constexpr[int], - order_elems: cutlass.Constexpr[int], + order_elements: cutlass.Constexpr[int], tidx, n_heads_out: cutlass.Int32, n_tiles: cutlass.Int32, @@ -612,27 +610,28 @@ def order_body( mStaging: cute.Tensor | None, mCount: cute.Tensor, mWorkItems: cute.Tensor, - mSched: cute.Tensor | None, + mScheduler: cute.Tensor | None, sKey, sIdx, sSpread, ): """LPT ordering body over ``n_threads`` CTA threads and caller-owned SMEM - staging (``sKey``/``sIdx`` of ``n_threads * order_elems`` Int32 cells + - a 2-cell ``sSpread``): bitonic-sort the items by span ``cend - cstart``, - longest first, into the final table. Sorts the walk's staged items, or - with ``gen`` synthesizes the uncut whole-sequence item per (batch, head) - from ``cu_seqlens`` directly — the no-cuts table. Thread 0 also zeroes - every ``sched_ctr`` cell (the main kernels' ticket rings, dirty on exit). - Runs on the standalone :func:`order_kernel` CTA, or fused into a main - kernel's CTA 0 prologue. Internally CTA-wide-barriers; every thread of + staging (``sKey``/``sIdx`` of ``n_threads * order_elements`` Int32 cells + + a 2-cell ``sSpread``): bitonic-sort the items by span + ``compute_end - compute_start``, longest first, into the final table. + Sorts the walk's staged items, or with ``gen`` synthesizes the uncut + whole-sequence item per (batch, head) from ``cu_seqlens`` directly — + the no-cuts table. Thread 0 also zeroes every ``scheduler_counter`` cell + (the main kernels' ticket rings, dirty on exit). + Runs fused into each main kernel's single-CTA prologue launch. + Internally CTA-wide-barriers; every thread of the calling CTA must reach it.""" - capacity = cutlass.const_expr(n_threads * order_elems) - if cutlass.const_expr(has_sched): + capacity = cutlass.const_expr(n_threads * order_elements) + if cutlass.const_expr(has_scheduler): if tidx == 0: si = cutlass.Int32(0) - while si < mSched.shape[0]: - mSched[si] = cutlass.Int32(0) + while si < mScheduler.shape[0]: + mScheduler[si] = cutlass.Int32(0) si = si + cutlass.Int32(1) if cutlass.const_expr(gen): n = n_tiles @@ -655,12 +654,12 @@ def order_body( nvvm.barrier_cta_sync() kmin = cutlass.Int32(2147483647) kmax = cutlass.Int32(-2147483648) - for e in cutlass.range_constexpr(order_elems): + for e in cutlass.range_constexpr(order_elements): i = tidx + cutlass.Int32(e * n_threads) if i < n: if cutlass.const_expr(gen): - batch_idx, head_idx, batch_start, batch_end, num_chunks_b = gen_item_bounds(b_t, n_heads_out, mCuSeqlens, i) - key = num_chunks_b + batch_idx, head_idx, batch_start, batch_end, batch_num_chunks = gen_item_bounds(b_t, n_heads_out, mCuSeqlens, i) + key = batch_num_chunks else: key = mStaging[i, 5] - mStaging[i, 4] sKey[i] = key @@ -685,7 +684,7 @@ def order_body( while k <= b_pad: j = k // cutlass.Int32(2) while j > 0: - for e in cutlass.range_constexpr(order_elems): + for e in cutlass.range_constexpr(order_elements): i = tidx + cutlass.Int32(e * n_threads) if i < b_pad: l = i ^ j @@ -704,7 +703,7 @@ def order_body( nvvm.barrier_cta_sync() j = j // cutlass.Int32(2) k = k * cutlass.Int32(2) - for e in cutlass.range_constexpr(order_elems): + for e in cutlass.range_constexpr(order_elements): i = tidx + cutlass.Int32(e * n_threads) if i < n: src = sIdx[i] @@ -719,7 +718,7 @@ def launch( safe_gate: cutlass.Constexpr[bool], gate_channels: cutlass.Constexpr[int], overhead_chunks: cutlass.Constexpr[int], - has_sched: cutlass.Constexpr[bool], + has_scheduler: cutlass.Constexpr[bool], n_heads_out: cutlass.Int32, n_tiles: cutlass.Int32, num_sms: cutlass.Int32, @@ -735,7 +734,7 @@ def launch( mStaging: cute.Tensor | None, mWorkItems: cute.Tensor, mCount: cute.Tensor, - mSched: cute.Tensor | None, + mScheduler: cute.Tensor | None, n_scan_ctas: cutlass.Int32, n_walk_ctas: cutlass.Int32, stream: cuda.CUstream, @@ -828,7 +827,7 @@ class TableRecipe(NamedTuple): n_walk_ctas: int -def run_table(r, gate, a_log, dt_bias, cu_seqlens, chunk_scratch, item_scratch, work_items, work_count, sched_ctr, stream) -> None: +def run_table(r, gate, a_log, dt_bias, cu_seqlens, chunk_scratch, item_scratch, work_items, work_count, scheduler_counter, stream) -> None: """The lowered split-table launch: no validation, no key build. Only buffers move between calls; every scalar comes from the recipe.""" r.compiled( @@ -847,7 +846,7 @@ def run_table(r, gate, a_log, dt_bias, cu_seqlens, chunk_scratch, item_scratch, item_scratch if r.split else None, work_items, work_count, - sched_ctr, + scheduler_counter, r.n_scan_ctas, r.n_walk_ctas, cuda.CUstream(int(stream)), @@ -872,7 +871,7 @@ def build_split_table( a_log=None, dt_bias=None, gate_lower_bound=None, - sched_ctr=None, + scheduler_counter=None, split=True, stream, ) -> "TableRecipe": @@ -887,21 +886,18 @@ def build_split_table( ``gate_lower_bound * sigmoid(exp(a_log) * (g + dt_bias))`` per element; scalar (GDN) ``-exp(a_log[h]) * softplus(g + dt_bias[h])`` per head. - With ``split=False`` the scan and walk never launch: the order kernel - alone synthesizes the no-cuts table — the uncut whole-sequence item per - (batch, head), LPT-sorted by sequence length — and ``ideal_chunks`` / - ``chunk_scratch`` / ``item_scratch`` / the gate contents are unused. - Batch-invariant mode and coarse checkpoint cadences (cuts may not cross - a checkpoint period) take this path. + ``split=False`` is rejected: the no-cuts table (the uncut whole-sequence + item per (batch, head), LPT-sorted) is synthesized by each main kernel's + own prologue (``order_gen``), so batch-invariant mode and coarse + checkpoint cadences (cuts may not cross a checkpoint period) never + involve this function. ``work_items`` and ``item_scratch`` are ``(max_items, - WORK_ITEM_FIELDS)`` int32 with ``max_items >= max_work_items(...)`` - (``>= n_tiles`` rows and no ``item_scratch`` with ``split=False``); - ``work_count`` is ``(1,)`` int32. Every cell of ``sched_ctr`` when + WORK_ITEM_FIELDS)`` int32 with ``max_items >= max_work_items(...)``; + ``work_count`` is ``(1,)`` int32. Every cell of ``scheduler_counter`` when passed — the main kernels' int32 ticket rings, ``(2,)`` per kernel launch that consumes this table, dirty on exit — is zeroed by the - order kernel, which runs in both modes; the count is zeroed by the - scan (split) or written by the order kernel (non-split). + order phase; the count is zeroed by the scan. ``chunk_scratch`` is ``(>= chunk_scratch_rows(total_tokens, B, b_t), HO)`` fp32 (contents managed here). Runs entirely on device — no host synchronization.""" @@ -920,41 +916,27 @@ def build_split_table( gate_scale_log2 = float(gate_lower_bound) * RCP_LN2 if safe_gate and gate_channels > 0 else 0.0 n_heads_out = gate.shape[1] batch_size = cu_seqlens.shape[0] - 1 - if split: - if ideal_chunks is None or chunk_scratch is None or item_scratch is None: - raise ValueError("split=True requires ideal_chunks, chunk_scratch, and item_scratch") - if gate_channels and gate_channels % WARP_SIZE != 0: - raise ValueError(f"per-channel gate dim must be a multiple of {WARP_SIZE}, got {gate_channels}") - if gate_channels and gate_channels % 128 == 0 and data_ptr(gate) % 16 != 0: - raise ValueError("per-channel gate base must be 16-byte aligned (vectorized scan loads)") - n_walk_ctas = batch_size * n_heads_out - need_rows = chunk_scratch_rows(gate.shape[0], batch_size, b_t) - n_scan_ctas = -(-need_rows // (SCAN_WARPS * SCAN_ROWS_PER_WARP)) - if len(chunk_scratch.shape) != 2 or chunk_scratch.shape[0] < need_rows or chunk_scratch.shape[1] != n_heads_out: - raise ValueError(f"chunk_scratch must be (>= {need_rows}, {n_heads_out}) fp32, got {tuple(chunk_scratch.shape)}") - if tuple(item_scratch.shape) != tuple(work_items.shape) or work_items.shape[1] != WORK_ITEM_FIELDS: - raise ValueError( - f"item_scratch must match work_items (max_items, {WORK_ITEM_FIELDS}) int32, got {tuple(item_scratch.shape)} vs {tuple(work_items.shape)}" - ) - else: - if work_items.shape[0] < n_tiles or work_items.shape[1] != WORK_ITEM_FIELDS: - raise ValueError(f"work_items must be (>= {n_tiles}, {WORK_ITEM_FIELDS}) int32, got {tuple(work_items.shape)}") - log_gate = False - safe_gate = False - gate_channels = 0 - a_log = None - dt_bias = None - gate_scale_log2 = 0.0 - gate = None - chunk_scratch = None - item_scratch = None - ideal_chunks = 0 - n_scan_ctas = 0 - n_walk_ctas = 0 + if not split: + raise ValueError("split=False has no host-side stage: each main kernel's prologue synthesizes the no-cuts table (order_gen)") + if ideal_chunks is None or chunk_scratch is None or item_scratch is None: + raise ValueError("split=True requires ideal_chunks, chunk_scratch, and item_scratch") + if gate_channels and gate_channels % WARP_SIZE != 0: + raise ValueError(f"per-channel gate dim must be a multiple of {WARP_SIZE}, got {gate_channels}") + if gate_channels and gate_channels % 128 == 0 and data_ptr(gate) % 16 != 0: + raise ValueError("per-channel gate base must be 16-byte aligned (vectorized scan loads)") + n_walk_ctas = batch_size * n_heads_out + need_rows = chunk_scratch_rows(gate.shape[0], batch_size, b_t) + n_scan_ctas = -(-need_rows // (SCAN_WARPS * SCAN_ROWS_PER_WARP)) + if len(chunk_scratch.shape) != 2 or chunk_scratch.shape[0] < need_rows or chunk_scratch.shape[1] != n_heads_out: + raise ValueError(f"chunk_scratch must be (>= {need_rows}, {n_heads_out}) fp32, got {tuple(chunk_scratch.shape)}") + if tuple(item_scratch.shape) != tuple(work_items.shape) or work_items.shape[1] != WORK_ITEM_FIELDS: + raise ValueError( + f"item_scratch must match work_items (max_items, {WORK_ITEM_FIELDS}) int32, got {tuple(item_scratch.shape)} vs {tuple(work_items.shape)}" + ) overhead_chunks = max(1, OVERHEAD_TOKENS // b_t) cu_stream = cuda.CUstream(int(stream)) - key = (bool(split), b_t, n_heads_out, bool(log_gate), bool(safe_gate), gate_channels, sched_ctr is not None, str(cu_seqlens.dtype)) + key = (bool(split), b_t, n_heads_out, bool(log_gate), bool(safe_gate), gate_channels, scheduler_counter is not None, str(cu_seqlens.dtype)) if key not in compiled_cache: dt_bias_c = None @@ -980,7 +962,7 @@ def build_split_table( bool(safe_gate), gate_channels, overhead_chunks, - sched_ctr is not None, + scheduler_counter is not None, cutlass.Int32(n_heads_out), cutlass.Int32(n_tiles), cutlass.Int32(num_sms), @@ -996,7 +978,7 @@ def build_split_table( item_scratch_c, work_items_c, work_count_c, - from_dlpack(sched_ctr, assumed_align=4).mark_layout_dynamic() if sched_ctr is not None else None, + from_dlpack(scheduler_counter, assumed_align=4).mark_layout_dynamic() if scheduler_counter is not None else None, cutlass.Int32(n_scan_ctas), cutlass.Int32(n_walk_ctas), cu_stream, @@ -1018,7 +1000,7 @@ def build_split_table( item_scratch, work_items, work_count, - sched_ctr, + scheduler_counter, n_scan_ctas, n_walk_ctas, cu_stream, diff --git a/python/cudnn/linear_attention/frost/common/thd.py b/python/cudnn/linear_attention/frost/common/thd.py index 0ad38786e..2e8caa8a8 100644 --- a/python/cudnn/linear_attention/frost/common/thd.py +++ b/python/cudnn/linear_attention/frost/common/thd.py @@ -80,7 +80,7 @@ def emit_checkpoint_seq_descs( """Per-BATCH descriptor array for the per-chunk checkpoint tensor with the head axis as a descriptor dimension (``(dv, dk, chunk, head)``). Derives the per-sequence checkpoint offsets from the TOKEN ``cu_seqlens`` on the fly - (``count_b = (seqlen_b - 1) // every_n + 1``, running-prefix-summed) — an + (``count_b = (batch_seqlen - 1) // every_n + 1``, running-prefix-summed) — an address fold no coordinate transform can express — and caps GLOBAL_DIM[``seq_ord``] to ``count_b``. The head index is a load coordinate. Runs on one electing thread; the calling warp elects and diff --git a/python/cudnn/linear_attention/frost/engine.py b/python/cudnn/linear_attention/frost/engine.py index a20dd52ef..3e2bba36f 100644 --- a/python/cudnn/linear_attention/frost/engine.py +++ b/python/cudnn/linear_attention/frost/engine.py @@ -86,5 +86,5 @@ def execute(self, graph, variant_pack, ctx) -> None: if not ok: raise ValueError(dense_layout_message(self.compiled.plan_name, ports, offender)) views = variant_pack.operands(self.indices) - workspace = Workspace.over(variant_pack, self.compiled.ws_bytes, type(self.compiled).__name__) + workspace = Workspace.over(variant_pack, self.compiled.workspace_size, type(self.compiled).__name__) self.compiled.run(views, workspace, ctx.stream) diff --git a/python/cudnn/linear_attention/frost/gdn2_engine.py b/python/cudnn/linear_attention/frost/gdn2_engine.py index 0b7067fb1..e7d734e69 100644 --- a/python/cudnn/linear_attention/frost/gdn2_engine.py +++ b/python/cudnn/linear_attention/frost/gdn2_engine.py @@ -4,7 +4,7 @@ """FROST GDN-2 engine: GDN2 nodes on the chunked prefill kernel (``kernel/gdn2_prefill_f16.py``, Blackwell SM100/SM103, bf16/fp16, BT=16). Forward + backward (GDN2_BWD on ``kernel/gdn2_bprop_f16.py`` with a -checkpoint regen on ``kernel/gdn2_recompute_f16.py``); the only GDN-2 engine — no cuTile +checkpoint recompute on ``kernel/gdn2_recompute_f16.py``); the only GDN-2 engine — no cuTile fallback.""" from __future__ import annotations @@ -29,13 +29,13 @@ def build_gdn2(graph): raise ValueError("build_gdn2: graph does not contain exactly one GDN2/GDN2_BWD node") node = nodes[0] if node.node_type.name == "GDN2_BWD": - from .kernel import gdn2_bprop_f16 as bwd_mod - from .kernel import gdn2_recompute_f16 as regen_mod + from .kernel import gdn2_bprop_f16 as bwd_module + from .kernel import gdn2_recompute_f16 as recompute_module - return CompiledGdn2Bwd(node, bwd_mod, regen_mod) - from .kernel import gdn2_prefill_f16 as kernel_mod + return CompiledGdn2Bwd(node, bwd_module, recompute_module) + from .kernel import gdn2_prefill_f16 as kernel_module - return CompiledGdn2(node, kernel_mod) + return CompiledGdn2(node, kernel_module) class Gdn2FrostEngine(BaseEngine): @@ -52,9 +52,9 @@ def check_support(self, graph) -> None: facts = graph._facts_for(analyze) frost_la_gate("Gdn2FrostEngine", facts, "GDN2") - ckpt = facts.checkpoint_every_n_tokens - if ckpt and (facts.is_bwd or ckpt % 16 != 0): - raise NotImplementedError(f"Gdn2FrostEngine: checkpoint_every_n_tokens must be a positive multiple of 16 on the GDN-2 node (got {ckpt})") + checkpoint = facts.checkpoint_every_n_tokens + if checkpoint and (facts.is_bwd or checkpoint % 16 != 0): + raise NotImplementedError(f"Gdn2FrostEngine: checkpoint_every_n_tokens must be a positive multiple of 16 on the GDN-2 node (got {checkpoint})") if not facts.gates_at_ho: raise NotImplementedError(f"Gdn2FrostEngine: g/beta/w must carry HO = max(q, v) heads ({facts.h_o})") fp32 = cudnn.data_type.FLOAT @@ -106,11 +106,11 @@ def build_plan(self, graph, plan, ctx=None) -> CompiledPlan: class CompiledGdn2: """Compiled FROST GDN-2 plan: a callable over the resolved node buffers.""" - def __init__(self, node, kernel_mod): + def __init__(self, node, kernel_module): from .common.split_k import WORK_ITEM_FIELDS, build_split_table, chunk_scratch_rows, compute_ideal_chunks, max_work_items, run_table self.node = node - self.kernel = kernel_mod + self.kernel = kernel_module self.build_split_table = build_split_table self.run_table = run_table self.table = None @@ -122,20 +122,20 @@ def __init__(self, node, kernel_mod): self.safe_gate = bool(node.params.get("safe_gate", False)) self.use_beta_sigmoid = bool(node.params.get("use_beta_sigmoid", False)) glb = node.params.get("gate_lower_bound") - self.gate_lower_bound = float(glb) if glb is not None else kernel_mod.DEFAULT_GATE_LOWER_BOUND + self.gate_lower_bound = float(glb) if glb is not None else kernel_module.DEFAULT_GATE_LOWER_BOUND self.has_final_state = "final_state" in node.outputs self.has_state_checkpoints = "state_checkpoints" in node.outputs - self.ckpt = int(node.params.get("checkpoint_every_n_tokens", 0) or 0) + self.checkpoint = int(node.params.get("checkpoint_every_n_tokens", 0) or 0) self.batch_invariant = bool(node.params.get("batch_invariant", False)) q, g = node.inputs["q"], node.inputs["g"] - self.b_t = kernel_mod.CFG.B_T - self.split = self.ckpt % self.b_t == 0 and not self.batch_invariant + self.b_t = kernel_module.CFG.B_T + self.split = self.checkpoint % self.b_t == 0 and not self.batch_invariant total = q.dim[0] HO = g.dim[1] B = node.inputs["cu_seqlens"].dim[0] - 1 layout = WorkspaceLayout() - self.off_sched = layout.add(8) + self.off_scheduler = layout.add(8) self.num_sm = multiprocessor_count(current_device()) self.n_tiles = B * HO self.n_heads_out = HO @@ -153,12 +153,12 @@ def __init__(self, node, kernel_mod): self.off_chunk_scratch = layout.add(self.chunk_scratch_rows * HO * 4) from .common.host import tensormap_workspace_bytes - self.tensormap_bytes = tensormap_workspace_bytes(kernel_mod, B) + self.tensormap_bytes = tensormap_workspace_bytes(kernel_module, B) self.off_tensormaps = layout.add(self.tensormap_bytes, align=128) self.needs_table = self.split - self.ws_bytes = layout.size + self.workspace_size = layout.size regions = [ - (self.off_sched, "int32", (2,)), + (self.off_scheduler, "int32", (2,)), (self.off_work_items, "int32", (self.work_item_rows, WORK_ITEM_FIELDS)), (self.off_work_count, "int32", (1,)), ] @@ -171,7 +171,7 @@ def __init__(self, node, kernel_mod): self.carve = carve_plan("Gdn2FrostEngine (GDN2)", regions) def workspace_bytes(self) -> int: - return self.ws_bytes + return self.workspace_size def bind(self, names) -> None: pos = {name: i for i, name in enumerate(names)} @@ -206,14 +206,14 @@ def run(self, views, workspace, stream) -> None: stream = stream if stream is not None else 0 if self.split: - sched_ctr, work_items, work_count, item_scratch, chunk_scratch, tensormaps = workspace.carve(self.carve) + scheduler_counter, work_items, work_count, item_scratch, chunk_scratch, tensormaps = workspace.carve(self.carve) else: - sched_ctr, work_items, work_count, tensormaps = workspace.carve(self.carve) + scheduler_counter, work_items, work_count, tensormaps = workspace.carve(self.carve) item_scratch = chunk_scratch = None if self.kcache is not None and (self.table is not None or not self.needs_table): if self.needs_table: - self.run_table(self.table, g, a_log, dt_bias, cu, chunk_scratch, item_scratch, work_items, work_count, sched_ctr, stream) + self.run_table(self.table, g, a_log, dt_bias, cu, chunk_scratch, item_scratch, work_items, work_count, scheduler_counter, stream) self.kernel.run_prefill( self.kcache, q, @@ -231,10 +231,10 @@ def run(self, views, workspace, stream) -> None: state_checkpoints, work_items, work_count, - sched_ctr, + scheduler_counter, item_scratch, tensormaps, - self.ckpt if self.has_state_checkpoints else 0, + self.checkpoint if self.has_state_checkpoints else 0, self.scale, stream, ) @@ -259,14 +259,14 @@ def run(self, views, workspace, stream) -> None: a_log=a_log, dt_bias=dt_bias, gate_lower_bound=self.gate_lower_bound if self.safe_gate else None, - sched_ctr=sched_ctr, + scheduler_counter=scheduler_counter, split=self.split, stream=stream, ) - ckpt_kwargs = {} + checkpoint_kwargs = {} if self.has_state_checkpoints: - ckpt_kwargs = dict(checkpoint_every_n_tokens=self.ckpt, output_state_checkpoints=state_checkpoints) + checkpoint_kwargs = dict(checkpoint_every_n_tokens=self.checkpoint, output_state_checkpoints=state_checkpoints) self.kcache = self.kernel.chunk_gdn2_sm100( q, k, @@ -287,10 +287,10 @@ def run(self, views, workspace, stream) -> None: use_beta_sigmoid=self.use_beta_sigmoid, work_items=work_items, work_count=work_count, - sched_ctr=sched_ctr, + scheduler_counter=scheduler_counter, work_item_scratch=item_scratch, tensormap_workspace=tensormaps, - **ckpt_kwargs, + **checkpoint_kwargs, stream=stream, ) return None @@ -301,17 +301,17 @@ class CompiledGdn2Bwd: regenerated per-chunk checkpoint series when the graph carries no ``state_checkpoints`` input, plus GVA/GQA head scratch for dQ/dK/dV.""" - def __init__(self, node, bwd_mod, regen_mod): + def __init__(self, node, bwd_module, recompute_module): from .common.split_k import WORK_ITEM_FIELDS, build_split_table, chunk_scratch_rows, compute_ideal_chunks, max_work_items, run_table self.node = node - self.bwd = bwd_mod - self.regen = regen_mod + self.bwd = bwd_module + self.recompute = recompute_module self.build_split_table = build_split_table self.run_table = run_table self.table = None self.kcache = None - self.regen_cache = None + self.recompute_cache = None self.plan_name = "Gdn2FrostEngine (GDN2_BWD)" from .common.gate_bwd import GATE_BWD_BLOCKS, channel_gate_bwd from .common.head_reduce import head_group_reduce @@ -325,13 +325,13 @@ def __init__(self, node, bwd_mod, regen_mod): self.safe_gate = bool(node.params.get("safe_gate", False)) self.use_beta_sigmoid = bool(node.params.get("use_beta_sigmoid", False)) glb = node.params.get("gate_lower_bound") - self.gate_lower_bound = float(glb) if glb is not None else bwd_mod.DEFAULT_GATE_LOWER_BOUND + self.gate_lower_bound = float(glb) if glb is not None else bwd_module.DEFAULT_GATE_LOWER_BOUND self.gate_bwd_blocks = GATE_BWD_BLOCKS self.has_state_checkpoints = "state_checkpoints" in node.inputs self.has_dstate0 = "d_initial_state" in node.outputs q, g, v = node.inputs["q"], node.inputs["g"], node.inputs["v"] - self.b_t = bwd_mod.CFG.B_T + self.b_t = bwd_module.CFG.B_T total = q.dim[0] HQ, HV = q.dim[1], v.dim[1] HO = g.dim[1] @@ -340,9 +340,10 @@ def __init__(self, node, bwd_mod, regen_mod): self.io_name = "float16" if node.inputs["q"].get_data_type().name == "HALF" else "bfloat16" self.n_heads_out, self.total = HO, total layout = WorkspaceLayout() - self.off_sched = layout.add(16) + self.off_scheduler = layout.add(16) self.num_sm = multiprocessor_count(current_device()) - self.bwd_dyn_sched = B * HO <= self.num_sm + # dynamic always: static costs 2.5-10.5% at multi-wave tile counts and never wins (lyris job 2752338) + self.bwd_dynamic_scheduling = True self.batch_invariant = bool(node.params.get("batch_invariant", False)) # cuts never in batch-invariant mode self.split = not self.batch_invariant @@ -362,8 +363,8 @@ def __init__(self, node, bwd_mod, regen_mod): if not self.has_state_checkpoints: self.state_checkpoints_rows = max(total // self.b_t + B, 1) self.off_state_checkpoints = layout.add(self.state_checkpoints_rows * HO * K * V * 2) - self.regen_tm_bytes = tensormap_workspace_bytes(regen_mod, B) - self.off_regen_tensormaps = layout.add(self.regen_tm_bytes, align=128) + self.recompute_tensormap_bytes = tensormap_workspace_bytes(recompute_module, B) + self.off_recompute_tensormaps = layout.add(self.recompute_tensormap_bytes, align=128) HK = node.inputs["k"].dim[1] self.fold_dq = HQ < HO self.fold_dk = HK < HO @@ -377,25 +378,25 @@ def __init__(self, node, bwd_mod, regen_mod): if self.safe_gate: self.off_gate_part_a = layout.add(self.gate_bwd_blocks * HO * K * 4) self.off_gate_part_dt = layout.add(self.gate_bwd_blocks * HO * K * 4) - self.bwd_tm_bytes = tensormap_workspace_bytes(bwd_mod, B) - self.off_bwd_tensormaps = layout.add(self.bwd_tm_bytes, align=128) - self.order_in_regen = not self.has_state_checkpoints + self.bwd_tensormap_bytes = tensormap_workspace_bytes(bwd_module, B) + self.off_bwd_tensormaps = layout.add(self.bwd_tensormap_bytes, align=128) + self.order_in_recompute = not self.has_state_checkpoints self.needs_table = self.split - self.ws_bytes = layout.size + self.workspace_size = layout.size regions = [ - ("sched_regen", self.off_sched, "int32", (2,)), - ("sched_bwd", self.off_sched + 8, "int32", (2,)), - ("sched_all", self.off_sched, "int32", (4,)), + ("scheduler_recompute", self.off_scheduler, "int32", (2,)), + ("scheduler_bwd", self.off_scheduler + 8, "int32", (2,)), + ("scheduler_all", self.off_scheduler, "int32", (4,)), ("work_items", self.off_work_items, "int32", (self.work_item_rows, WORK_ITEM_FIELDS)), ("work_count", self.off_work_count, "int32", (1,)), - ("bwd_tensormaps", self.off_bwd_tensormaps, "int64", (self.bwd_tm_bytes // 8,)), + ("bwd_tensormaps", self.off_bwd_tensormaps, "int64", (self.bwd_tensormap_bytes // 8,)), ] if self.split: regions.append(("item_scratch", self.off_item_scratch, "int32", (self.work_item_rows, WORK_ITEM_FIELDS))) regions.append(("chunk_scratch", self.off_chunk_scratch, "float32", (self.chunk_scratch_rows, HO))) if not self.has_state_checkpoints: regions.append(("state_checkpoints", self.off_state_checkpoints, self.io_name, (self.state_checkpoints_rows, HO, V, K))) - regions.append(("regen_tensormaps", self.off_regen_tensormaps, "int64", (self.regen_tm_bytes // 8,))) + regions.append(("recompute_tensormaps", self.off_recompute_tensormaps, "int64", (self.recompute_tensormap_bytes // 8,))) if self.fold_dq: regions.append(("dq_ho", self.off_dq_ho, self.io_name, (total, HO, K))) if self.fold_dk: @@ -409,7 +410,7 @@ def __init__(self, node, bwd_mod, regen_mod): self.carve = carve_plan("Gdn2FrostEngine (GDN2_BWD)", [(off, dt, shape) for _name, off, dt, shape in regions]) def workspace_bytes(self) -> int: - return self.ws_bytes + return self.workspace_size def bind(self, names) -> None: pos = {name: i for i, name in enumerate(names)} @@ -462,8 +463,8 @@ def run(self, views, workspace, stream) -> None: stream = stream if stream is not None else 0 region = dict(zip(self.carve_names, workspace.carve(self.carve))) - sched_regen = region["sched_regen"] - sched_bwd = region["sched_bwd"] + scheduler_recompute = region["scheduler_recompute"] + scheduler_bwd = region["scheduler_bwd"] work_items = region["work_items"] work_count = region["work_count"] @@ -479,15 +480,15 @@ def run(self, views, workspace, stream) -> None: region.get("item_scratch"), work_items, work_count, - region["sched_all"], + region["scheduler_all"], stream, ) if self.has_state_checkpoints: checkpoint_series = state_checkpoints else: checkpoint_series = region["state_checkpoints"] - self.regen.run_recompute( - self.regen_cache, + self.recompute.run_recompute( + self.recompute_cache, k, v, g, @@ -501,10 +502,10 @@ def run(self, views, workspace, stream) -> None: checkpoint_series, work_items, work_count, - sched_regen, - region["sched_all"], + scheduler_recompute, + region["scheduler_all"], region.get("item_scratch"), - region["regen_tensormaps"], + region["recompute_tensormaps"], self.b_t, stream, ) @@ -532,9 +533,9 @@ def run(self, views, workspace, stream) -> None: dstate_in, work_items, work_count, - sched_bwd if self.bwd_dyn_sched else None, - region["sched_all"] if not self.order_in_regen else None, - region.get("item_scratch") if not self.order_in_regen else None, + scheduler_bwd if self.bwd_dynamic_scheduling else None, + region["scheduler_all"] if not self.order_in_recompute else None, + region.get("item_scratch") if not self.order_in_recompute else None, region["bwd_tensormaps"], self.scale, stream, @@ -570,7 +571,7 @@ def run(self, views, workspace, stream) -> None: a_log=a_log, dt_bias=dt_bias, gate_lower_bound=self.gate_lower_bound if self.safe_gate else None, - sched_ctr=region["sched_all"], + scheduler_counter=region["scheduler_all"], split=self.split, stream=stream, ) @@ -579,7 +580,7 @@ def run(self, views, workspace, stream) -> None: checkpoint_series = state_checkpoints else: checkpoint_series = region["state_checkpoints"] - self.regen_cache = self.regen.chunk_gdn2_recompute_sm100( + self.recompute_cache = self.recompute.chunk_gdn2_recompute_sm100( k, v, g, @@ -598,11 +599,11 @@ def run(self, views, workspace, stream) -> None: use_beta_sigmoid=self.use_beta_sigmoid, work_items=work_items, work_count=work_count, - sched_ctr=sched_regen, - sched_all=region["sched_all"], + scheduler_counter=scheduler_recompute, + scheduler_all=region["scheduler_all"], work_item_scratch=region.get("item_scratch"), order_in_prologue=True, - tensormap_workspace=region["regen_tensormaps"], + tensormap_workspace=region["recompute_tensormaps"], stream=stream, ) @@ -642,10 +643,10 @@ def run(self, views, workspace, stream) -> None: use_beta_sigmoid=self.use_beta_sigmoid, work_items=work_items, work_count=work_count, - sched_ctr=sched_bwd if self.bwd_dyn_sched else None, - sched_all=region["sched_all"] if not self.order_in_regen else None, - work_item_scratch=region.get("item_scratch") if not self.order_in_regen else None, - order_in_prologue=not self.order_in_regen, + scheduler_counter=scheduler_bwd if self.bwd_dynamic_scheduling else None, + scheduler_all=region["scheduler_all"] if not self.order_in_recompute else None, + work_item_scratch=region.get("item_scratch") if not self.order_in_recompute else None, + order_in_prologue=not self.order_in_recompute, tensormap_workspace=region["bwd_tensormaps"], stream=stream, ) diff --git a/python/cudnn/linear_attention/frost/gdn_engine.py b/python/cudnn/linear_attention/frost/gdn_engine.py index 197f93d25..939d7922b 100644 --- a/python/cudnn/linear_attention/frost/gdn_engine.py +++ b/python/cudnn/linear_attention/frost/gdn_engine.py @@ -29,13 +29,13 @@ def build_gdn(graph): raise ValueError("build_gdn: graph does not contain exactly one GDN/GDN_BWD node") node = nodes[0] if node.node_type.name == "GDN_BWD": - from .kernel import gdn_bprop_f16 as bwd_mod - from .kernel import gdn_recompute_f16 as regen_mod + from .kernel import gdn_bprop_f16 as bwd_module + from .kernel import gdn_recompute_f16 as recompute_module - return CompiledGdnBwd(node, bwd_mod, regen_mod) - from .kernel import gdn_prefill_f16 as kernel_mod + return CompiledGdnBwd(node, bwd_module, recompute_module) + from .kernel import gdn_prefill_f16 as kernel_module - return CompiledGdn(node, kernel_mod) + return CompiledGdn(node, kernel_module) class GdnFrostEngine(BaseEngine): @@ -52,9 +52,9 @@ def check_support(self, graph) -> None: facts = graph._facts_for(analyze) frost_la_gate("GdnFrostEngine", facts, "GDN") - ckpt = facts.checkpoint_every_n_tokens - if ckpt and (facts.is_bwd or ckpt % 64 != 0): - raise NotImplementedError(f"GdnFrostEngine: checkpoint_every_n_tokens must be a positive multiple of 64 on the GDN node (got {ckpt})") + checkpoint = facts.checkpoint_every_n_tokens + if checkpoint and (facts.is_bwd or checkpoint % 64 != 0): + raise NotImplementedError(f"GdnFrostEngine: checkpoint_every_n_tokens must be a positive multiple of 64 on the GDN node (got {checkpoint})") if not facts.gates_at_ho: raise NotImplementedError(f"GdnFrostEngine: g/beta must carry HO = max(q, v) heads ({facts.h_o})") fp32 = cudnn.data_type.FLOAT @@ -103,11 +103,11 @@ def build_plan(self, graph, plan, ctx=None) -> CompiledPlan: class CompiledGdn: """Compiled FROST GDN plan: a callable over the resolved node buffers.""" - def __init__(self, node, kernel_mod): + def __init__(self, node, kernel_module): from .common.split_k import WORK_ITEM_FIELDS, build_split_table, chunk_scratch_rows, compute_ideal_chunks, max_work_items, run_table self.node = node - self.kernel = kernel_mod + self.kernel = kernel_module self.build_split_table = build_split_table self.run_table = run_table self.table = None @@ -123,24 +123,24 @@ def __init__(self, node, kernel_mod): self.use_beta_sigmoid = bool(node.params.get("use_beta_sigmoid", False)) q, v, g = node.inputs["q"], node.inputs["v"], node.inputs["g"] - self.b_t = kernel_mod.CFG.B_T + self.b_t = kernel_module.CFG.B_T total = q.dim[0] HO = g.dim[1] HQ, HK = q.dim[1], node.inputs["k"].dim[1] self.io_name = "float16" if q.get_data_type().name == "HALF" else "bfloat16" B = node.inputs["cu_seqlens"].dim[0] - 1 self.has_final_state = "final_state" in node.outputs - self.ckpt = int(node.params.get("checkpoint_every_n_tokens", 0) or 0) + self.checkpoint = int(node.params.get("checkpoint_every_n_tokens", 0) or 0) self.has_state_checkpoints = "state_checkpoints" in node.outputs self.batch_invariant = bool(node.params.get("batch_invariant", False)) - self.split = self.ckpt % self.b_t == 0 and not self.batch_invariant + self.split = self.checkpoint % self.b_t == 0 and not self.batch_invariant layout = WorkspaceLayout() from .common.host import tensormap_workspace_bytes - self.tensormap_words = tensormap_workspace_bytes(kernel_mod, B) // 8 + self.tensormap_words = tensormap_workspace_bytes(kernel_module, B) // 8 self.off_tensormaps = layout.add(self.tensormap_words * 8) - self.off_sched = layout.add(8) + self.off_scheduler = layout.add(8) self.num_sm = multiprocessor_count(current_device()) self.n_tiles = B * HO self.n_heads_out = HO @@ -162,10 +162,10 @@ def __init__(self, node, kernel_mod): self.off_inv_q = layout.add(total * HQ * 4) self.off_inv_k = layout.add(total * HK * 4) self.needs_table = self.split - self.ws_bytes = layout.size + self.workspace_size = layout.size regions = [ (self.off_tensormaps, "int64", (self.tensormap_words,)), - (self.off_sched, "int32", (2,)), + (self.off_scheduler, "int32", (2,)), (self.off_work_items, "int32", (self.work_item_rows, WORK_ITEM_FIELDS)), (self.off_work_count, "int32", (1,)), ] @@ -184,7 +184,7 @@ def __init__(self, node, kernel_mod): self.carve = carve_plan("GdnFrostEngine (GDN)", regions) def workspace_bytes(self) -> int: - return self.ws_bytes + return self.workspace_size def bind(self, names) -> None: pos = {name: i for i, name in enumerate(names)} @@ -217,9 +217,9 @@ def run(self, views, workspace, stream) -> None: stream = stream if stream is not None else 0 carved = workspace.carve(self.carve) if self.split: - tensormaps, sched_ctr, work_items, work_count, item_scratch, chunk_scratch, *l2n = carved + tensormaps, scheduler_counter, work_items, work_count, item_scratch, chunk_scratch, *l2n = carved else: - tensormaps, sched_ctr, work_items, work_count, *l2n = carved + tensormaps, scheduler_counter, work_items, work_count, *l2n = carved item_scratch = chunk_scratch = None if self.use_qk_l2norm: q_n, k_n, inv_q, inv_k = l2n @@ -228,7 +228,7 @@ def run(self, views, workspace, stream) -> None: if self.kcache is not None and (self.table is not None or not self.needs_table): if self.needs_table: - self.run_table(self.table, g, a_log, dt_bias, cu, chunk_scratch, item_scratch, work_items, work_count, sched_ctr, stream) + self.run_table(self.table, g, a_log, dt_bias, cu, chunk_scratch, item_scratch, work_items, work_count, scheduler_counter, stream) self.kernel.run_prefill( self.kcache, q, @@ -243,10 +243,10 @@ def run(self, views, workspace, stream) -> None: state_checkpoints, work_items, work_count, - sched_ctr, + scheduler_counter, item_scratch, tensormaps, - self.ckpt, + self.checkpoint, self.scale, stream, a_log=a_log if self.safe_gate else None, @@ -272,7 +272,7 @@ def run(self, views, workspace, stream) -> None: safe_gate=self.safe_gate, a_log=a_log, dt_bias=dt_bias, - sched_ctr=sched_ctr, + scheduler_counter=scheduler_counter, split=self.split, stream=stream, ) @@ -290,8 +290,8 @@ def run(self, views, workspace, stream) -> None: self.scale, work_items=work_items, work_count=work_count, - sched_ctr=sched_ctr, - checkpoint_every_n_tokens=self.ckpt, + scheduler_counter=scheduler_counter, + checkpoint_every_n_tokens=self.checkpoint, output_state_checkpoints=state_checkpoints, log_gate=True, safe_gate=self.safe_gate, @@ -312,17 +312,17 @@ class CompiledGdnBwd: the node's ``state_checkpoints`` input, or regenerates them with the recompute (checkpoint-only) kernel when the port is absent.""" - def __init__(self, node, bwd_mod, regen_mod): + def __init__(self, node, bwd_module, recompute_module): from .common.split_k import WORK_ITEM_FIELDS, build_split_table, chunk_scratch_rows, compute_ideal_chunks, max_work_items, run_table self.node = node - self.bwd = bwd_mod - self.regen = regen_mod + self.bwd = bwd_module + self.recompute = recompute_module self.build_split_table = build_split_table self.run_table = run_table self.table = None self.kcache = None - self.regen_cache = None + self.recompute_cache = None self.plan_name = "GdnFrostEngine (GDN_BWD)" from .common.gate_bwd import scalar_gate_bwd, scalar_gate_blocks from .common.head_reduce import head_group_reduce @@ -340,7 +340,7 @@ def __init__(self, node, bwd_mod, regen_mod): self.use_beta_sigmoid = bool(node.params.get("use_beta_sigmoid", False)) q, v, g = node.inputs["q"], node.inputs["v"], node.inputs["g"] - self.b_t = bwd_mod.CFG.B_T + self.b_t = bwd_module.CFG.B_T total = q.dim[0] self.gate_bwd_blocks = scalar_gate_blocks(total) K, V = q.dim[-1], v.dim[-1] @@ -351,19 +351,19 @@ def __init__(self, node, bwd_mod, regen_mod): self.io_name = "float16" if node.inputs["q"].get_data_type().name == "HALF" else "bfloat16" self.num_sm = multiprocessor_count(current_device()) - self.bwd_dyn_sched = B * HO <= self.num_sm + self.bwd_dynamic_scheduling = True self.batch_invariant = bool(node.params.get("batch_invariant", False)) # cuts never in batch-invariant mode self.split = not self.batch_invariant layout = WorkspaceLayout() - self.off_sched = layout.add(16) - self.tensormap_words = tensormap_workspace_bytes(bwd_mod, B) // 8 + self.off_scheduler = layout.add(16) + self.tensormap_words = tensormap_workspace_bytes(bwd_module, B) // 8 self.off_tensormaps = layout.add(self.tensormap_words * 8) if not self.has_state_checkpoints: self.state_checkpoints_rows = max(total // self.b_t + B, 1) self.off_state_checkpoints = layout.add(self.state_checkpoints_rows * HO * K * V * 2) - self.regen_tensormap_words = tensormap_workspace_bytes(regen_mod, B) // 8 - self.off_regen_tensormaps = layout.add(self.regen_tensormap_words * 8) + self.recompute_tensormap_words = tensormap_workspace_bytes(recompute_module, B) // 8 + self.off_recompute_tensormaps = layout.add(self.recompute_tensormap_words * 8) HK = node.inputs["k"].dim[1] self.fold_dq = HQ < HO self.fold_dk = HK < HO @@ -396,14 +396,14 @@ def __init__(self, node, bwd_mod, regen_mod): self.chunk_scratch_rows = chunk_scratch_rows(total, B, self.b_t) self.off_chunk_scratch = layout.add(self.chunk_scratch_rows * HO * 4) self.needs_table = self.split - self.regen_orders = not self.has_state_checkpoints + self.recompute_orders = not self.has_state_checkpoints self.bwd_orders = self.has_state_checkpoints self.n_heads_out, self.total = HO, total - self.ws_bytes = layout.size + self.workspace_size = layout.size regions = [ - ("sched_regen", self.off_sched, "int32", (2,)), - ("sched_bwd", self.off_sched + 8, "int32", (2,)), - ("sched_all", self.off_sched, "int32", (4,)), + ("scheduler_recompute", self.off_scheduler, "int32", (2,)), + ("scheduler_bwd", self.off_scheduler + 8, "int32", (2,)), + ("scheduler_all", self.off_scheduler, "int32", (4,)), ("tensormaps", self.off_tensormaps, "int64", (self.tensormap_words,)), ("work_items", self.off_work_items, "int32", (self.work_item_rows, WORK_ITEM_FIELDS)), ("work_count", self.off_work_count, "int32", (1,)), @@ -413,7 +413,7 @@ def __init__(self, node, bwd_mod, regen_mod): regions.append(("chunk_scratch", self.off_chunk_scratch, "float32", (self.chunk_scratch_rows, HO))) if not self.has_state_checkpoints: regions.append(("state_checkpoints", self.off_state_checkpoints, self.io_name, (self.state_checkpoints_rows, HO, V, K))) - regions.append(("regen_tensormaps", self.off_regen_tensormaps, "int64", (self.regen_tensormap_words,))) + regions.append(("recompute_tensormaps", self.off_recompute_tensormaps, "int64", (self.recompute_tensormap_words,))) if self.fold_dq: regions.append(("dq_ho", self.off_dq_ho, self.io_name, (total, HO, K))) if self.fold_dk: @@ -432,7 +432,7 @@ def __init__(self, node, bwd_mod, regen_mod): self.carve = carve_plan("GdnFrostEngine (GDN_BWD)", [(off, dt, shape) for _name, off, dt, shape in regions]) def workspace_bytes(self) -> int: - return self.ws_bytes + return self.workspace_size def bind(self, names) -> None: pos = {name: i for i, name in enumerate(names)} @@ -481,8 +481,8 @@ def run(self, views, workspace, stream) -> None: stream = stream if stream is not None else 0 region = dict(zip(self.carve_names, workspace.carve(self.carve))) - sched_regen = region["sched_regen"] - sched_bwd = region["sched_bwd"] + scheduler_recompute = region["scheduler_recompute"] + scheduler_bwd = region["scheduler_bwd"] work_items = region["work_items"] work_count = region["work_count"] if self.use_qk_l2norm: @@ -501,15 +501,15 @@ def run(self, views, workspace, stream) -> None: region.get("item_scratch"), work_items, work_count, - region["sched_all"], + region["scheduler_all"], stream, ) if self.has_state_checkpoints: checkpoint_series = state_checkpoints else: checkpoint_series = region["state_checkpoints"] - self.regen.run_recompute( - self.regen_cache, + self.recompute.run_recompute( + self.recompute_cache, k, v, g, @@ -520,10 +520,10 @@ def run(self, views, workspace, stream) -> None: checkpoint_series, work_items, work_count, - sched_regen, - region["sched_all"] if self.regen_orders else None, - region.get("item_scratch") if self.regen_orders else None, - region["regen_tensormaps"], + scheduler_recompute, + region["scheduler_all"] if self.recompute_orders else None, + region.get("item_scratch") if self.recompute_orders else None, + region["recompute_tensormaps"], self.b_t, stream, a_log=a_log if self.safe_gate else None, @@ -551,8 +551,8 @@ def run(self, views, workspace, stream) -> None: dstate_in, work_items, work_count, - sched_bwd if self.bwd_dyn_sched else None, - region["sched_all"] if self.bwd_orders else None, + scheduler_bwd if self.bwd_dynamic_scheduling else None, + region["scheduler_all"] if self.bwd_orders else None, region.get("item_scratch") if self.bwd_orders else None, region["tensormaps"], self.scale, @@ -588,7 +588,7 @@ def run(self, views, workspace, stream) -> None: safe_gate=self.safe_gate, a_log=a_log, dt_bias=dt_bias, - sched_ctr=region["sched_all"], + scheduler_counter=region["scheduler_all"], split=self.split, stream=stream, ) @@ -597,7 +597,7 @@ def run(self, views, workspace, stream) -> None: checkpoint_series = state_checkpoints else: checkpoint_series = region["state_checkpoints"] - self.regen_cache = self.regen.chunk_gdn_recompute_sm100( + self.recompute_cache = self.recompute.chunk_gdn_recompute_sm100( k, v, g, @@ -613,12 +613,12 @@ def run(self, views, workspace, stream) -> None: use_beta_sigmoid=self.use_beta_sigmoid, work_items=work_items, work_count=work_count, - sched_ctr=sched_regen, - sched_all=region["sched_all"] if self.regen_orders else None, - work_item_scratch=region.get("item_scratch") if self.regen_orders else None, - order_in_prologue=self.regen_orders, + scheduler_counter=scheduler_recompute, + scheduler_all=region["scheduler_all"] if self.recompute_orders else None, + work_item_scratch=region.get("item_scratch") if self.recompute_orders else None, + order_in_prologue=self.recompute_orders, log_gate=True, - workspace=region["regen_tensormaps"], + workspace=region["recompute_tensormaps"], stream=stream, ) @@ -653,8 +653,8 @@ def run(self, views, workspace, stream) -> None: use_beta_sigmoid=self.use_beta_sigmoid, work_items=work_items, work_count=work_count, - sched_ctr=sched_bwd if self.bwd_dyn_sched else None, - sched_all=region["sched_all"] if self.bwd_orders else None, + scheduler_counter=scheduler_bwd if self.bwd_dynamic_scheduling else None, + scheduler_all=region["scheduler_all"] if self.bwd_orders else None, work_item_scratch=region.get("item_scratch") if self.bwd_orders else None, order_in_prologue=self.bwd_orders, log_gate=True, diff --git a/python/cudnn/linear_attention/frost/kda_engine.py b/python/cudnn/linear_attention/frost/kda_engine.py index e661f0fda..7378372a8 100644 --- a/python/cudnn/linear_attention/frost/kda_engine.py +++ b/python/cudnn/linear_attention/frost/kda_engine.py @@ -30,13 +30,13 @@ def build_kda(graph): raise ValueError("build_kda: graph does not contain exactly one KDA/KDA_BWD node") node = nodes[0] if node.node_type.name == "KDA_BWD": - from .kernel import kda_bprop_f16 as bwd_mod - from .kernel import kda_recompute_f16 as regen_mod + from .kernel import kda_bprop_f16 as bwd_module + from .kernel import kda_recompute_f16 as recompute_module - return CompiledKdaBwd(node, bwd_mod, regen_mod) - from .kernel import kda_prefill_f16 as kernel_mod + return CompiledKdaBwd(node, bwd_module, recompute_module) + from .kernel import kda_prefill_f16 as kernel_module - return CompiledKda(node, kernel_mod) + return CompiledKda(node, kernel_module) class KdaFrostEngine(BaseEngine): @@ -53,9 +53,9 @@ def check_support(self, graph) -> None: facts = graph._facts_for(analyze) frost_la_gate("KdaFrostEngine", facts, "KDA") - ckpt = facts.checkpoint_every_n_tokens - if ckpt and (facts.is_bwd or ckpt % 16 != 0): - raise NotImplementedError(f"KdaFrostEngine: checkpoint_every_n_tokens must be a positive multiple of 16 on the KDA node (got {ckpt})") + checkpoint = facts.checkpoint_every_n_tokens + if checkpoint and (facts.is_bwd or checkpoint % 16 != 0): + raise NotImplementedError(f"KdaFrostEngine: checkpoint_every_n_tokens must be a positive multiple of 16 on the KDA node (got {checkpoint})") if not facts.gates_at_ho: raise NotImplementedError(f"KdaFrostEngine: g/beta must carry HO = max(q, v) heads ({facts.h_o})") fp32 = cudnn.data_type.FLOAT @@ -104,11 +104,11 @@ def build_plan(self, graph, plan, ctx=None) -> CompiledPlan: class CompiledKda: """Compiled FROST KDA plan: a callable over the resolved node buffers.""" - def __init__(self, node, kernel_mod): + def __init__(self, node, kernel_module): from .common.split_k import WORK_ITEM_FIELDS, build_split_table, chunk_scratch_rows, compute_ideal_chunks, max_work_items, run_table self.node = node - self.kernel = kernel_mod + self.kernel = kernel_module self.build_split_table = build_split_table self.run_table = run_table self.table = None @@ -120,20 +120,20 @@ def __init__(self, node, kernel_mod): self.use_beta_sigmoid = bool(node.params.get("use_beta_sigmoid", False)) self.safe_gate = bool(node.params.get("safe_gate", False)) glb = node.params.get("gate_lower_bound") - self.gate_lower_bound = float(glb) if glb is not None else kernel_mod.DEFAULT_GATE_LOWER_BOUND + self.gate_lower_bound = float(glb) if glb is not None else kernel_module.DEFAULT_GATE_LOWER_BOUND self.has_final_state = "final_state" in node.outputs self.has_state_checkpoints = "state_checkpoints" in node.outputs - self.ckpt = int(node.params.get("checkpoint_every_n_tokens", 0) or 0) + self.checkpoint = int(node.params.get("checkpoint_every_n_tokens", 0) or 0) self.batch_invariant = bool(node.params.get("batch_invariant", False)) q, g = node.inputs["q"], node.inputs["g"] - self.b_t = kernel_mod.CFG.B_T - self.split = self.ckpt % self.b_t == 0 and not self.batch_invariant + self.b_t = kernel_module.CFG.B_T + self.split = self.checkpoint % self.b_t == 0 and not self.batch_invariant total = q.dim[0] HO = g.dim[1] B = node.inputs["cu_seqlens"].dim[0] - 1 layout = WorkspaceLayout() - self.off_sched = layout.add(8) + self.off_scheduler = layout.add(8) self.num_sm = multiprocessor_count(current_device()) self.n_tiles = B * HO self.n_heads_out = HO @@ -151,12 +151,12 @@ def __init__(self, node, kernel_mod): self.off_chunk_scratch = layout.add(self.chunk_scratch_rows * HO * 4) from .common.host import tensormap_workspace_bytes - self.tensormap_bytes = tensormap_workspace_bytes(kernel_mod, B) + self.tensormap_bytes = tensormap_workspace_bytes(kernel_module, B) self.off_tensormaps = layout.add(self.tensormap_bytes, align=128) self.needs_table = self.split - self.ws_bytes = layout.size + self.workspace_size = layout.size regions = [ - (self.off_sched, "int32", (2,)), + (self.off_scheduler, "int32", (2,)), (self.off_work_items, "int32", (self.work_item_rows, WORK_ITEM_FIELDS)), (self.off_work_count, "int32", (1,)), ] @@ -169,7 +169,7 @@ def __init__(self, node, kernel_mod): self.carve = carve_plan("KdaFrostEngine (KDA)", regions) def workspace_bytes(self) -> int: - return self.ws_bytes + return self.workspace_size def bind(self, names) -> None: pos = {name: i for i, name in enumerate(names)} @@ -202,14 +202,14 @@ def run(self, views, workspace, stream) -> None: stream = stream if stream is not None else 0 if self.split: - sched_ctr, work_items, work_count, item_scratch, chunk_scratch, tensormaps = workspace.carve(self.carve) + scheduler_counter, work_items, work_count, item_scratch, chunk_scratch, tensormaps = workspace.carve(self.carve) else: - sched_ctr, work_items, work_count, tensormaps = workspace.carve(self.carve) + scheduler_counter, work_items, work_count, tensormaps = workspace.carve(self.carve) item_scratch = chunk_scratch = None if self.kcache is not None and (self.table is not None or not self.needs_table): if self.needs_table: - self.run_table(self.table, g, a_log, dt_bias, cu, chunk_scratch, item_scratch, work_items, work_count, sched_ctr, stream) + self.run_table(self.table, g, a_log, dt_bias, cu, chunk_scratch, item_scratch, work_items, work_count, scheduler_counter, stream) self.kernel.run_prefill( self.kcache, q, @@ -226,10 +226,10 @@ def run(self, views, workspace, stream) -> None: state_checkpoints, work_items, work_count, - sched_ctr, + scheduler_counter, item_scratch, tensormaps, - self.ckpt if self.has_state_checkpoints else 0, + self.checkpoint if self.has_state_checkpoints else 0, self.scale, stream, ) @@ -254,14 +254,14 @@ def run(self, views, workspace, stream) -> None: a_log=a_log, dt_bias=dt_bias, gate_lower_bound=self.gate_lower_bound if self.safe_gate else None, - sched_ctr=sched_ctr, + scheduler_counter=scheduler_counter, split=self.split, stream=stream, ) - ckpt_kwargs = {} + checkpoint_kwargs = {} if self.has_state_checkpoints: - ckpt_kwargs = dict(checkpoint_every_n_tokens=self.ckpt, output_state_checkpoints=state_checkpoints) + checkpoint_kwargs = dict(checkpoint_every_n_tokens=self.checkpoint, output_state_checkpoints=state_checkpoints) self.kcache = self.kernel.chunk_kda_sm100( q, k, @@ -281,10 +281,10 @@ def run(self, views, workspace, stream) -> None: dt_bias=dt_bias, work_items=work_items, work_count=work_count, - sched_ctr=sched_ctr, + scheduler_counter=scheduler_counter, work_item_scratch=item_scratch, tensormap_workspace=tensormaps, - **ckpt_kwargs, + **checkpoint_kwargs, stream=stream, ) return None @@ -297,17 +297,17 @@ class CompiledKdaBwd: gradients land in HO-head scratch and are reduced back to the native head counts.""" - def __init__(self, node, bwd_mod, regen_mod): + def __init__(self, node, bwd_module, recompute_module): from .common.split_k import WORK_ITEM_FIELDS, build_split_table, chunk_scratch_rows, compute_ideal_chunks, max_work_items, run_table self.node = node - self.bwd = bwd_mod - self.regen = regen_mod + self.bwd = bwd_module + self.recompute = recompute_module self.build_split_table = build_split_table self.run_table = run_table self.table = None self.kcache = None - self.regen_cache = None + self.recompute_cache = None self.plan_name = "KdaFrostEngine (KDA_BWD)" from .common.gate_bwd import GATE_BWD_BLOCKS, channel_gate_bwd from .common.head_reduce import head_group_reduce @@ -321,13 +321,13 @@ def __init__(self, node, bwd_mod, regen_mod): self.safe_gate = bool(node.params.get("safe_gate", False)) self.use_beta_sigmoid = bool(node.params.get("use_beta_sigmoid", False)) glb = node.params.get("gate_lower_bound") - self.gate_lower_bound = float(glb) if glb is not None else bwd_mod.DEFAULT_GATE_LOWER_BOUND + self.gate_lower_bound = float(glb) if glb is not None else bwd_module.DEFAULT_GATE_LOWER_BOUND self.gate_bwd_blocks = GATE_BWD_BLOCKS self.has_state_checkpoints = "state_checkpoints" in node.inputs self.has_dstate0 = "d_initial_state" in node.outputs q, g, v = node.inputs["q"], node.inputs["g"], node.inputs["v"] - self.b_t = bwd_mod.CFG.B_T + self.b_t = bwd_module.CFG.B_T total = q.dim[0] HQ, HV = q.dim[1], v.dim[1] HO = g.dim[1] @@ -336,9 +336,10 @@ def __init__(self, node, bwd_mod, regen_mod): self.io_name = "float16" if node.inputs["q"].get_data_type().name == "HALF" else "bfloat16" self.n_heads_out, self.total = HO, total layout = WorkspaceLayout() - self.off_sched = layout.add(16) + self.off_scheduler = layout.add(16) self.num_sm = multiprocessor_count(current_device()) - self.bwd_dyn_sched = B * HO <= self.num_sm + # dynamic always: static costs 2.5-10.5% at multi-wave tile counts and never wins (lyris job 2752338) + self.bwd_dynamic_scheduling = True self.batch_invariant = bool(node.params.get("batch_invariant", False)) # cuts never in batch-invariant mode self.split = not self.batch_invariant @@ -358,8 +359,8 @@ def __init__(self, node, bwd_mod, regen_mod): if not self.has_state_checkpoints: self.state_checkpoints_rows = max(total // self.b_t + B, 1) self.off_state_checkpoints = layout.add(self.state_checkpoints_rows * HO * K * V * 2) - self.regen_tm_bytes = tensormap_workspace_bytes(regen_mod, B) - self.off_regen_tensormaps = layout.add(self.regen_tm_bytes, align=128) + self.recompute_tensormap_bytes = tensormap_workspace_bytes(recompute_module, B) + self.off_recompute_tensormaps = layout.add(self.recompute_tensormap_bytes, align=128) HK = node.inputs["k"].dim[1] self.fold_dq = HQ < HO self.fold_dk = HK < HO @@ -373,24 +374,24 @@ def __init__(self, node, bwd_mod, regen_mod): if self.safe_gate: self.off_gate_part_a = layout.add(self.gate_bwd_blocks * HO * K * 4) self.off_gate_part_dt = layout.add(self.gate_bwd_blocks * HO * K * 4) - self.bwd_tm_bytes = tensormap_workspace_bytes(bwd_mod, B) - self.off_bwd_tensormaps = layout.add(self.bwd_tm_bytes, align=128) + self.bwd_tensormap_bytes = tensormap_workspace_bytes(bwd_module, B) + self.off_bwd_tensormaps = layout.add(self.bwd_tensormap_bytes, align=128) self.needs_table = self.split - self.ws_bytes = layout.size + self.workspace_size = layout.size regions = [ - ("sched_regen", self.off_sched, "int32", (2,)), - ("sched_bwd", self.off_sched + 8, "int32", (2,)), - ("sched_all", self.off_sched, "int32", (4,)), + ("scheduler_recompute", self.off_scheduler, "int32", (2,)), + ("scheduler_bwd", self.off_scheduler + 8, "int32", (2,)), + ("scheduler_all", self.off_scheduler, "int32", (4,)), ("work_items", self.off_work_items, "int32", (self.work_item_rows, WORK_ITEM_FIELDS)), ("work_count", self.off_work_count, "int32", (1,)), - ("bwd_tensormaps", self.off_bwd_tensormaps, "int64", (self.bwd_tm_bytes // 8,)), + ("bwd_tensormaps", self.off_bwd_tensormaps, "int64", (self.bwd_tensormap_bytes // 8,)), ] if self.split: regions.append(("item_scratch", self.off_item_scratch, "int32", (self.work_item_rows, WORK_ITEM_FIELDS))) regions.append(("chunk_scratch", self.off_chunk_scratch, "float32", (self.chunk_scratch_rows, HO))) if not self.has_state_checkpoints: regions.append(("state_checkpoints", self.off_state_checkpoints, self.io_name, (self.state_checkpoints_rows, HO, V, K))) - regions.append(("regen_tensormaps", self.off_regen_tensormaps, "int64", (self.regen_tm_bytes // 8,))) + regions.append(("recompute_tensormaps", self.off_recompute_tensormaps, "int64", (self.recompute_tensormap_bytes // 8,))) if self.fold_dq: regions.append(("dq_ho", self.off_dq_ho, self.io_name, (total, HO, K))) if self.fold_dk: @@ -404,7 +405,7 @@ def __init__(self, node, bwd_mod, regen_mod): self.carve = carve_plan("KdaFrostEngine (KDA_BWD)", [(off, dt, shape) for _name, off, dt, shape in regions]) def workspace_bytes(self) -> int: - return self.ws_bytes + return self.workspace_size def bind(self, names) -> None: pos = {name: i for i, name in enumerate(names)} @@ -453,8 +454,8 @@ def run(self, views, workspace, stream) -> None: stream = stream if stream is not None else 0 region = dict(zip(self.carve_names, workspace.carve(self.carve))) - sched_regen = region["sched_regen"] - sched_bwd = region["sched_bwd"] + scheduler_recompute = region["scheduler_recompute"] + scheduler_bwd = region["scheduler_bwd"] work_items = region["work_items"] work_count = region["work_count"] @@ -470,15 +471,15 @@ def run(self, views, workspace, stream) -> None: region.get("item_scratch"), work_items, work_count, - region["sched_all"], + region["scheduler_all"], stream, ) if self.has_state_checkpoints: checkpoint_series = state_checkpoints else: checkpoint_series = region["state_checkpoints"] - self.regen.run_recompute( - self.regen_cache, + self.recompute.run_recompute( + self.recompute_cache, k, v, g, @@ -491,10 +492,10 @@ def run(self, views, workspace, stream) -> None: checkpoint_series, work_items, work_count, - sched_regen, - region["sched_all"], + scheduler_recompute, + region["scheduler_all"], region.get("item_scratch"), - region["regen_tensormaps"], + region["recompute_tensormaps"], self.b_t, stream, ) @@ -520,8 +521,8 @@ def run(self, views, workspace, stream) -> None: dstate_in, work_items, work_count, - sched_bwd if self.bwd_dyn_sched else None, - region["sched_all"] if self.has_state_checkpoints else None, + scheduler_bwd if self.bwd_dynamic_scheduling else None, + region["scheduler_all"] if self.has_state_checkpoints else None, region.get("item_scratch") if self.has_state_checkpoints else None, region["bwd_tensormaps"], self.scale, @@ -558,7 +559,7 @@ def run(self, views, workspace, stream) -> None: a_log=a_log, dt_bias=dt_bias, gate_lower_bound=self.gate_lower_bound if self.safe_gate else None, - sched_ctr=region["sched_all"], + scheduler_counter=region["scheduler_all"], split=self.split, stream=stream, ) @@ -567,7 +568,7 @@ def run(self, views, workspace, stream) -> None: checkpoint_series = state_checkpoints else: checkpoint_series = region["state_checkpoints"] - self.regen_cache = self.regen.chunk_kda_recompute_sm100( + self.recompute_cache = self.recompute.chunk_kda_recompute_sm100( k, v, g, @@ -585,11 +586,11 @@ def run(self, views, workspace, stream) -> None: use_beta_sigmoid=self.use_beta_sigmoid, work_items=work_items, work_count=work_count, - sched_ctr=sched_regen, - sched_all=region["sched_all"], + scheduler_counter=scheduler_recompute, + scheduler_all=region["scheduler_all"], work_item_scratch=region.get("item_scratch"), order_in_prologue=True, - tensormap_workspace=region["regen_tensormaps"], + tensormap_workspace=region["recompute_tensormaps"], stream=stream, ) @@ -627,8 +628,8 @@ def run(self, views, workspace, stream) -> None: use_beta_sigmoid=self.use_beta_sigmoid, work_items=work_items, work_count=work_count, - sched_ctr=sched_bwd if self.bwd_dyn_sched else None, - sched_all=region["sched_all"] if self.has_state_checkpoints else None, + scheduler_counter=scheduler_bwd if self.bwd_dynamic_scheduling else None, + scheduler_all=region["scheduler_all"] if self.has_state_checkpoints else None, work_item_scratch=region.get("item_scratch") if self.has_state_checkpoints else None, order_in_prologue=self.has_state_checkpoints, tensormap_workspace=region["bwd_tensormaps"], diff --git a/python/cudnn/linear_attention/frost/kernel/gdn2_bprop_config.py b/python/cudnn/linear_attention/frost/kernel/gdn2_bprop_config.py index 3d2bccf39..8ea96656b 100644 --- a/python/cudnn/linear_attention/frost/kernel/gdn2_bprop_config.py +++ b/python/cudnn/linear_attention/frost/kernel/gdn2_bprop_config.py @@ -44,16 +44,10 @@ class Cfg: EPILOGUE_WARP_ID: int = 15 # dq/dk/dv TMA stores only # --- register split --- - # WG1's drain runs at the register ceiling (spilled at 152) while WG0 - # sits near ~100 live regs, so WG0 funds WG1. Constraints: compute - # groups can't go below the 128-reg launch base (setmaxregister is - # INCREASE-only there), and warps 12-15 are one warpgroup so they must - # share a single setmaxregister value (56; super/epilogue peak ~R49). - # dht configs keep ~15 in-loop WG2 spills at 136 (152 doesn't fit). NUM_REGS_COMPUTE_GROUP_0: int = 128 NUM_REGS_COMPUTE_GROUP_1: int = 184 NUM_REGS_COMPUTE_GROUP_2: int = 136 - NUM_REGS_OTHER: int = 64 # warpgroup-uniform; +8 donated by CG1 + NUM_REGS_OTHER: int = 64 THREADS_PER_WARP: int = 32 diff --git a/python/cudnn/linear_attention/frost/kernel/gdn2_bprop_f16.py b/python/cudnn/linear_attention/frost/kernel/gdn2_bprop_f16.py index be5552bd9..283ad2bee 100644 --- a/python/cudnn/linear_attention/frost/kernel/gdn2_bprop_f16.py +++ b/python/cudnn/linear_attention/frost/kernel/gdn2_bprop_f16.py @@ -1,54 +1,3 @@ -"""Chunked Gated DeltaNet v2 (GDN-2) BPROP kernel for Blackwell SM100/SM103 -(Cutlass DSL), BT=16 tiling with a per-key-channel decay. Framework-neutral -entry ``chunk_gdn2_bwd_sm100``. - -Algorithm overview (per chunk c, iterated c = NT-1 .. 0; within-chunk -log2-domain gate cumsum G[t,d], eG = 2^G, eGl = 2^G[BT-1]): - Inputs : Q/K[BT,DK], V/dO[BT,DV], S = checkpoint[c-1] (state ENTERING chunk c, KV), - Gate[BT,DK], Beta[BT,DK] (per-key erase), W[BT,DV] (per-value write) - State : dH[DV,DK] (state gradient, fp32 TMEM, accumulated backward) - - Operands (WG0, prefill recompute): K_decay = eG.(Beta.K) (erase key, Beta - folded), K_inv = K/eG, K_restore = (eGl/eG).K, Q_decay = eG.Q, diag(eGl). - - Forward recompute: T_inv = (I + strict_tril(K_decay@K_inv^T))^-1 (register Neumann, - Beta pre-folded); A = tril_incl(Q_decay@K_inv^T); Y = W.V - S^T K_decay; U = T_inv@Y. - - Backward math: - dU = K_restore@dH + A^T@dO (Q_decay carries scale, so A does too) - dY = T_inv^T@dU - dV = W.dY dW_out = V.dY (elementwise, WG1) - dA = tril_incl(dO@U^T) (unscaled) dM = dY@U^T dM_strict = +strict(dM) - dQ = eG.scale.(dO@S^T + dA@K_inv) - dK = Beta.eG.dK_decay + dK_inv/eG + (eGl/eG).dK_restore where (sign-flipped parts) - dK_decay part = dY@S^T + dM_strict@K_inv (= -dK_decay) - dK_inv part = dA^T@(scale.Q_decay) - dM_strict^T@K_decay (= dK_inv; one TMEM - acc, the minus rides the staged -dM_strict tile) - dK_restore part = U@dH^T (= +dK_restore) - dBeta[t,d] = k_n.eG.dK_decay = -k_n.eG.dK_decay part (per-channel, WG2) - dGate[t,d] = q_n.dQ_pre + Beta.dBeta + k_n.(dK_inv_part/eG - - (eGl/eG).dK_restore_part) - dGate_last[d] = eGl.sum_v(dH.S) + sum_t k_n.(eGl/eG).dK_restore_part - dGate = suffix-sum(dGate + dGate_last at row BT-1) (WG2 in-register reverse cumsum) - dH <- diag-GEMM(eGl).dH + (scale.Q_decay)^T@dO - K_decay^T@dY - -ABI: state_checkpoints `[total_checkpoints, HO, DV, DK]` (VK, k contiguous - the GDN checkpoint layout) io -dtype, the plain per-chunk series (entry `c` = state entering chunk c, so row 0 -is the initial state or zeros); beta `[T, HO, DK]` / w `[T, HO, DV]` io dtype -(post-sigmoid); dq/dk/dv io at HO heads; dgate `[T, HO, DK]` fp32 (natural-log -gate domain); dbeta/dw io dtype like beta/w; d_initial_state / d_final_state -fp32 `[N, HO, DV, DK]` (V-major, matching the prefill states). - -Warp assignments (16 warps = 512 threads): - warps 0-3 : WG0 - Gate prefix scan + decay/restore operands (all chunks) - warps 4-7 : WG1 - value-side TMEM staging, restages, dstate capture, dV/dW_out - warps 8-11 : WG2 - dQ/dK part drain, dGate/dBeta assembly, reverse cumsum - warp 12 : super-MMA - register KK/A/dA/dM + Neumann inverse - warp 13 : tcgen05-MMA - the backward schedule - warp 14 : TMA load - Q/K/V/Gate/dO/state(checkpoint) loads + Beta/W tiles - warp 15 : epilogue - dQ/dK/dV/dGate/dBeta/dW_out TMA stores -""" - from dataclasses import dataclass from functools import lru_cache from typing import NamedTuple, Optional, Type @@ -60,7 +9,7 @@ import cutlass.cute as cute from cutlass.cute.runtime import from_dlpack -from ..common.split_k import ORDER_CAPACITY, ORDER_ELEMS, ORDER_THREADS, decode_work_item, order_body +from ..common.split_k import ORDER_CAPACITY, ORDER_ELEMENTS, ORDER_THREADS, decode_work_item, order_body from ..common.host import get_dtype from cudnn.frost.buffers import data_ptr from cudnn.frost.device import current_device, multiprocessor_count @@ -75,9 +24,10 @@ ) from cudnn.frost.tile_dsl.handles import MmaDesc, SmemTile, tma_slice_runtime_desc from cudnn.frost.tile_dsl.mma import mma_ss, mma_step, mma_ts_step -from cudnn.frost.tile_dsl.swizzle import swizzle_lin_S, swizzle_xor_128b +from cudnn.frost.tile_dsl.swizzle import swizzle_xor_128b, swizzle_xor_32b from cudnn.frost.tile_dsl.tma import tma_load_tile, tma_store_commit, tma_store_tile, tma_store_wait, tma_tensormap_acquire from cudnn.frost.tile_dsl.pointwise import ( + sigmoid, f16x2_to_f32, fadd2, ffma2, @@ -90,22 +40,12 @@ ) LOG2_E: float = 1.4426950408889634 - - DEFAULT_GATE_LOWER_BOUND: float = -5.0 - - L2_NORM_EPS: float = 1.0e-12 class Gdn2BwdBars(NamedTuple): - """Every inter-warp handoff as an ``MBarrier`` over its ring. Consumers - track ``(idx, phase)`` inline; the producer tag selects the arrive - lowering (``TMA_LOAD``/``MMA_COMMIT``/``THREAD``). - - Buffers read by both the MMA warp and a compute/warp group carry mixed - arrive counts (one MMA commit + N thread arrivers) so the producer only - reuses the slot once every reader is done.""" + """Every inter-warp handoff as an ``MBarrier`` over its ring.""" mb_q_ready: MBarrier mb_q_done: MBarrier @@ -125,9 +65,9 @@ class Gdn2BwdBars(NamedTuple): mb_state_ready: MBarrier mb_state_done: MBarrier mb_state_cg0_done: MBarrier - mb_state_inp_ready: MBarrier - mb_state_inp_done: MBarrier - mb_state_inp_cg2_done: MBarrier + mb_state_input_ready: MBarrier + mb_state_input_done: MBarrier + mb_state_input_cg2_done: MBarrier mb_k_decay_inv_ready: MBarrier mb_q_decay_k_restore_ready: MBarrier @@ -143,17 +83,17 @@ class Gdn2BwdBars(NamedTuple): mb_dm_done: MBarrier mb_state_k_acc_ready: MBarrier - mb_y_inp_ready: MBarrier + mb_y_input_ready: MBarrier mb_u_acc_ready: MBarrier mb_u_smem_ready: MBarrier mb_du_acc_ready: MBarrier - mb_du_inp_ready: MBarrier + mb_du_input_ready: MBarrier mb_dy_acc_ready: MBarrier - mb_neg_dy_inp_ready: MBarrier + mb_neg_dy_input_ready: MBarrier mb_dy_smem_ready: MBarrier mb_dy_smem_done: MBarrier mb_dstate_acc_ready: MBarrier - mb_dstate_inp_ready: MBarrier + mb_dstate_input_ready: MBarrier mb_dstate_smem_ready: MBarrier mb_dstate_smem_done: MBarrier mb_dstate_smem_cg2_done: MBarrier @@ -182,13 +122,12 @@ class Gdn2BwdBars(NamedTuple): mb_dstate0_acc_stored: MBarrier mb_tmem_done: MBarrier - mb_sched_ready: MBarrier - mb_sched_done: MBarrier + mb_scheduler_ready: MBarrier + mb_scheduler_done: MBarrier def make_gdn2_bwd_bars(cfg) -> Gdn2BwdBars: - """Bars factory. MUST be called from inside ``kernel`` (allocates the - mbarrier rings in SMEM ahead of the data buffers).""" + """Gdn2BwdBars factory.""" def alloc(n): return cutlass.Array(cutlass.Int64, n, space=cutlass.AddressSpace.smem, alignment=8) @@ -218,9 +157,9 @@ def alloc(n): mb_state_ready=MBarrier(alloc(cfg.smem_state_stages), stages=cfg.smem_state_stages, init_count=1, producer=Producer.TMA_LOAD), mb_state_done=MBarrier(alloc(cfg.smem_state_stages), stages=cfg.smem_state_stages, init_count=MMA, producer=Producer.MMA_COMMIT), mb_state_cg0_done=MBarrier(alloc(cfg.smem_state_stages), stages=cfg.smem_state_stages, init_count=CG0, producer=Producer.THREAD), - mb_state_inp_ready=MBarrier(alloc(2), stages=2, init_count=CG0, producer=Producer.THREAD), - mb_state_inp_done=MBarrier(alloc(2), stages=2, init_count=MMA, producer=Producer.MMA_COMMIT), - mb_state_inp_cg2_done=MBarrier(alloc(2), stages=2, init_count=CG2, producer=Producer.THREAD), + mb_state_input_ready=MBarrier(alloc(2), stages=2, init_count=CG0, producer=Producer.THREAD), + mb_state_input_done=MBarrier(alloc(2), stages=2, init_count=MMA, producer=Producer.MMA_COMMIT), + mb_state_input_cg2_done=MBarrier(alloc(2), stages=2, init_count=CG2, producer=Producer.THREAD), mb_k_decay_inv_ready=MBarrier(alloc(cfg.smem_decay_stages), stages=cfg.smem_decay_stages, init_count=CG0, producer=Producer.THREAD), mb_q_decay_k_restore_ready=MBarrier(alloc(cfg.smem_decay_stages), stages=cfg.smem_decay_stages, init_count=CG0, producer=Producer.THREAD), mb_decay_done=MBarrier(alloc(cfg.smem_decay_stages), stages=cfg.smem_decay_stages, init_count=MMA, producer=Producer.MMA_COMMIT), @@ -233,17 +172,17 @@ def alloc(n): mb_da_done=MBarrier(alloc(cfg.smem_intermediate_stages), stages=cfg.smem_intermediate_stages, init_count=MMA, producer=Producer.MMA_COMMIT), mb_dm_done=MBarrier(alloc(cfg.smem_intermediate_stages), stages=cfg.smem_intermediate_stages, init_count=MMA, producer=Producer.MMA_COMMIT), mb_state_k_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA, producer=Producer.MMA_COMMIT), - mb_y_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1, producer=Producer.THREAD), + mb_y_input_ready=MBarrier(alloc(1), stages=1, init_count=CG1, producer=Producer.THREAD), mb_u_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA, producer=Producer.MMA_COMMIT), mb_u_smem_ready=MBarrier(alloc(1), stages=1, init_count=CG1, producer=Producer.THREAD), mb_du_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA, producer=Producer.MMA_COMMIT), - mb_du_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1, producer=Producer.THREAD), + mb_du_input_ready=MBarrier(alloc(1), stages=1, init_count=CG1, producer=Producer.THREAD), mb_dy_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA, producer=Producer.MMA_COMMIT), - mb_neg_dy_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1, producer=Producer.THREAD), + mb_neg_dy_input_ready=MBarrier(alloc(1), stages=1, init_count=CG1, producer=Producer.THREAD), mb_dy_smem_ready=MBarrier(alloc(1), stages=1, init_count=CG1, producer=Producer.THREAD), mb_dy_smem_done=MBarrier(alloc(1), stages=1, init_count=MMA, producer=Producer.MMA_COMMIT), mb_dstate_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA, producer=Producer.MMA_COMMIT), - mb_dstate_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1, producer=Producer.THREAD), + mb_dstate_input_ready=MBarrier(alloc(1), stages=1, init_count=CG1, producer=Producer.THREAD), mb_dstate_smem_ready=MBarrier(alloc(1), stages=1, init_count=CG1, producer=Producer.THREAD), mb_dstate_smem_done=MBarrier(alloc(1), stages=1, init_count=MMA, producer=Producer.MMA_COMMIT), mb_dstate_smem_cg2_done=MBarrier(alloc(1), stages=1, init_count=CG2, producer=Producer.THREAD), @@ -268,43 +207,46 @@ def alloc(n): mb_dwo_tmastg_done=MBarrier(alloc(cfg.smem_dwo_stages), stages=cfg.smem_dwo_stages, init_count=WARP, producer=Producer.THREAD), mb_dstate0_acc_stored=MBarrier(alloc(1), stages=1, init_count=CG1, producer=Producer.THREAD), mb_tmem_done=MBarrier(alloc(1), stages=1, init_count=CG1 + CG2, producer=Producer.THREAD), - mb_sched_ready=MBarrier(alloc(cfg.sched_stages), stages=cfg.sched_stages, init_count=1, producer=Producer.THREAD), - mb_sched_done=MBarrier(alloc(cfg.sched_stages), stages=cfg.sched_stages, init_count=15, producer=Producer.THREAD), + mb_scheduler_ready=MBarrier(alloc(cfg.scheduler_stages), stages=cfg.scheduler_stages, init_count=1, producer=Producer.THREAD), + mb_scheduler_done=MBarrier(alloc(cfg.scheduler_stages), stages=cfg.scheduler_stages, init_count=15, producer=Producer.THREAD), ) -# ---- Dynamic tile scheduler ------------------------------------------------------ +# ---- Dynamic tile scheduler ---------------------------------------------------------- @cute.jit -def sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ctas): +def scheduler_publish_next(cfg, bars, sScheduler, mScheduler, scheduler_state, tile_idx, num_ctas, tail_base, tail_row, elect_one): """TMA-warp side: pull the next tile off the global ticket, publish it.""" - if cutlass.const_expr(cfg.dyn_sched): - bars.mb_sched_done[sched_state.idx].wait(sched_state.phase) - if nvvm.elect_sync(): - fetched = cutlass.Int32(nvvm.atomicrmw("add", mSched.iterator, cutlass.Int32(1), mem_order="relaxed", syncscope="gpu")) - sSched[sched_state.idx] = num_ctas + fetched + if cutlass.const_expr(cfg.dynamic_scheduling): + sentinel = cutlass.Int32(1 << 28) + pinned = tail_row if tile_idx < tail_base else sentinel + bars.mb_scheduler_done[scheduler_state.idx].wait(scheduler_state.phase) + if elect_one: + fetched = cutlass.Int32(nvvm.atomicrmw("add", mScheduler.iterator, cutlass.Int32(1), mem_order="relaxed", syncscope="gpu")) + granted = num_ctas + fetched + sScheduler[scheduler_state.idx] = granted if granted < tail_base else pinned nvvm.bar_warp_sync(cute.arch.FULL_MASK) - next_tile = sSched[sched_state.idx] - if nvvm.elect_sync(): - bars.mb_sched_ready[sched_state.idx].arrive() - return next_tile, advance(sched_state, cfg.sched_stages) - return tile_idx + num_ctas, sched_state + next_tile = sScheduler[scheduler_state.idx] + if elect_one: + bars.mb_scheduler_ready[scheduler_state.idx].arrive() + return next_tile, advance(scheduler_state, cfg.scheduler_stages) + return tile_idx + num_ctas, scheduler_state @cute.jit -def sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas): +def scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one): """Consumer side: read the TMA warp's published next tile.""" - if cutlass.const_expr(cfg.dyn_sched): - bars.mb_sched_ready[sched_state.idx].wait(sched_state.phase) - next_tile = sSched[sched_state.idx] - if nvvm.elect_sync(): - bars.mb_sched_done[sched_state.idx].arrive() - return next_tile, advance(sched_state, cfg.sched_stages) - return tile_idx + num_ctas, sched_state + if cutlass.const_expr(cfg.dynamic_scheduling): + bars.mb_scheduler_ready[scheduler_state.idx].wait(scheduler_state.phase) + next_tile = sScheduler[scheduler_state.idx] + if elect_one: + bars.mb_scheduler_done[scheduler_state.idx].arrive() + return next_tile, advance(scheduler_state, cfg.scheduler_stages) + return tile_idx + num_ctas, scheduler_state -# ---- Warp bodies ----------------------------------------------------------------- +# ---- Warp bodies --------------------------------------------------------------------- @cute.jit @@ -315,8 +257,8 @@ def epilogue_warp( num_ctas, cu_seqlens, mWorkItems, - sSched, - lane, + sScheduler, + lane_idx, sK_inv_raw, sQ_decay_raw, sDo_raw, @@ -338,30 +280,29 @@ def epilogue_warp( ) -> None: """Epilogue warp role (warp 15): register-MMA A/dA tiles and the gradient TMA stores, in chunk order with a one-behind store ladder.""" - elect_one = nvvm.elect_sync() - nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) - # ---- ldmatrix/stmatrix lane decode ------------------------------------------- - rhs_row_coord = lane % 8 + (cutlass.Int32(8) if (lane // 16) else cutlass.Int32(0)) - rhs_col_offset = cutlass.Int32(8) if ((lane // 8) % 2) else cutlass.Int32(0) - lhs_row_coord = lane % 8 + (cutlass.Int32(8) if ((lane // 8) % 2) else cutlass.Int32(0)) - lhs_col_offset = cutlass.Int32(8) if ((lane // 8) // 2) else cutlass.Int32(0) - stsm_row_coord = lane & 7 - stsm_col_coord = cutlass.Int32(0) - if (lane // 8) & 1: - stsm_row_coord = stsm_row_coord + cutlass.Int32(8) - if lane // 8 >= 2: - stsm_col_coord = cutlass.Int32(8) - stsm_idx = swizzle_lin_S(stsm_row_coord * cfg.b_t + stsm_col_coord, bbits=1, mbase=3, sshift=3) - row_lo = lane // 4 + elect_one = nvvm.elect_sync() + # ---- ldmatrix/stmatrix lane decode ----------------------------------------------- + b_row_coord = lane_idx % 8 + (cutlass.Int32(8) if (lane_idx // 16) else cutlass.Int32(0)) + b_col_offset = cutlass.Int32(8) if ((lane_idx // 8) % 2) else cutlass.Int32(0) + a_row_coord = lane_idx % 8 + (cutlass.Int32(8) if ((lane_idx // 8) % 2) else cutlass.Int32(0)) + a_col_offset = cutlass.Int32(8) if ((lane_idx // 8) // 2) else cutlass.Int32(0) + intermediate_row_coord = lane_idx & 7 + intermediate_col_coord = cutlass.Int32(0) + if (lane_idx // 8) & 1: + intermediate_row_coord = intermediate_row_coord + cutlass.Int32(8) + if lane_idx // 8 >= 2: + intermediate_col_coord = cutlass.Int32(8) + intermediate_idx = intermediate_row_coord * cfg.b_t + swizzle_xor_32b(intermediate_row_coord, intermediate_col_coord) + row_lo = lane_idx // 4 row_hi = row_lo + cutlass.Int32(8) - # hoisted tril bitmask: bit i = row >= col for accum index i + # tril bitmask: bit i = row >= col for accum index i tril_incl_mask = cutlass.Int32(0) for accum_idx in cutlass.range_constexpr(8): row_coord = row_hi if cutlass.const_expr(accum_idx % 4 >= 2) else row_lo - col_coord = (accum_idx // 4) * 8 + 2 * (lane % 4) + col_coord = (accum_idx // 4) * 8 + 2 * (lane_idx % 4) if cutlass.const_expr(accum_idx % 2 == 1): col_coord = col_coord + cutlass.Int32(1) tril_incl_mask = tril_incl_mask | (cutlass.Int32(1 << accum_idx) if row_coord >= col_coord else cutlass.Int32(0)) @@ -441,11 +382,13 @@ def epilogue_warp( dgate_index = PipelineState.start(phase=0) db_index = PipelineState.start(phase=0) dwo_index = PipelineState.start(phase=0) - sched_state = PipelineState.start(phase=0) + scheduler_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) FIRST_STATE_CHUNK = 0 if cfg.use_initial_state else 1 while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) head_o = head_idx slot = batch_idx * cutlass.Int32(TENSOR_MAP_QWORDS) if elect_one: @@ -461,13 +404,13 @@ def epilogue_warp( tma_tensormap_acquire(desc_dgate_slot) tma_tensormap_acquire(desc_db_slot) tma_tensormap_acquire(desc_dwo_slot) - sk_nt = cend - wstart + sk_nt = compute_end - write_start pend_start = cutlass.Int32(0) pend_writes = cutlass.Boolean(False) for rev_idx in cutlass.range(sk_nt, unroll=1): - chunk_idx = cend - cutlass.Int32(1) - rev_idx + chunk_idx = compute_end - cutlass.Int32(1) - rev_idx chunk_start = chunk_idx * cfg.b_t - writes = chunk_idx < wend + writes = chunk_idx < write_end gc = gbase + rev_idx decay_stage = gc % cfg.smem_decay_stages intermediate_stage = gc % cfg.smem_intermediate_stages @@ -477,31 +420,31 @@ def epilogue_warp( sDo_ptr = sDo_raw.data_ptr() + raw_stage * (cfg.d_v * cfg.b_t) sIntermediate_ptr = sIntermediate_raw.data_ptr() + intermediate_stage * (cfg.intermediate_tiles * cfg.b_t * cfg.b_t) - # ---- A = tril_incl(Q_decay @ K_inv^T) -------------------------------- + # ---- A = tril(Q decay @ K inv^T, 0) -------------------------------------- bars.mb_a_done[intermediate_stage].wait(((gc // cfg.smem_intermediate_stages) + 1) % 2) bars.mb_q_decay_k_restore_ready[decay_stage].wait((gc // cfg.smem_decay_stages) % 2) a_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): a_acc[accum_idx] = cutlass.Float32(0.0) - for k_block in cutlass.range_constexpr(cfg.d_k // 16): - a_col = k_block * 16 + lhs_col_offset + for i in cutlass.range_constexpr(cfg.d_k // 16): + a_col = i * 16 + a_col_offset a_seg = a_col // 64 - a_frag = nvvm.ldmatrix( - sQ_decay_ptr + a_seg * (cfg.b_t * 64) + lhs_row_coord * 64 + swizzle_xor_128b(lhs_row_coord, a_col - a_seg * 64, elem_bytes=2), + q_decay_frag = nvvm.ldmatrix( + sQ_decay_ptr + a_seg * (cfg.b_t * 64) + a_row_coord * 64 + swizzle_xor_128b(a_row_coord, a_col - a_seg * 64, elem_bytes=2), 4, nvvm.MMALayout.ROW, ) - b_col = k_block * 16 + rhs_col_offset + b_col = i * 16 + b_col_offset b_seg = b_col // 64 - b_frag = nvvm.ldmatrix( - sK_inv_ptr + b_seg * (cfg.b_t * 64) + rhs_row_coord * 64 + swizzle_xor_128b(rhs_row_coord, b_col - b_seg * 64, elem_bytes=2), + k_inv_frag = nvvm.ldmatrix( + sK_inv_ptr + b_seg * (cfg.b_t * 64) + b_row_coord * 64 + swizzle_xor_128b(b_row_coord, b_col - b_seg * 64, elem_bytes=2), 4, nvvm.MMALayout.ROW, ) mma_step( a_acc, - (a_frag[0], a_frag[1], a_frag[2], a_frag[3]), - (b_frag[0], b_frag[1], b_frag[2], b_frag[3]), + (q_decay_frag[0], q_decay_frag[1], q_decay_frag[2], q_decay_frag[3]), + (k_inv_frag[0], k_inv_frag[1], k_inv_frag[2], k_inv_frag[3]), k_step=0, M=16, N=16, @@ -510,7 +453,7 @@ def epilogue_warp( for accum_idx in cutlass.range_constexpr(8): a_acc[accum_idx] = a_acc[accum_idx] if (tril_incl_mask >> accum_idx) & 1 else cutlass.Float32(0.0) nvvm.stmatrix( - sIntermediate_ptr + stsm_idx, + sIntermediate_ptr + intermediate_idx, [ fp32_to_fp16(a_acc[0], a_acc[1], dtype=cfg.io_dtype), fp32_to_fp16(a_acc[2], a_acc[3], dtype=cfg.io_dtype), @@ -523,45 +466,43 @@ def epilogue_warp( nvvm.fence_proxy("async.shared", space="cta") bars.mb_a_ready[intermediate_stage].arrive() - # ---- dA = tril_incl(dO @ U^T) ---------------------------------------- + # ---- dA = tril(dO @ U^T, 0) ---------------------------------------------- bars.mb_u_smem_ready.wait(u_index.phase) u_index = advance(u_index, 1) da_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): da_acc[accum_idx] = cutlass.Float32(0.0) - for k_block in cutlass.range_constexpr(cfg.d_v // 16): - a_col = k_block * 16 + lhs_col_offset + for i in cutlass.range_constexpr(cfg.d_v // 16): + a_col = i * 16 + a_col_offset a_seg = a_col // 64 - a_frag = nvvm.ldmatrix( - sDo_ptr + a_seg * (cfg.b_t * 64) + lhs_row_coord * 64 + swizzle_xor_128b(lhs_row_coord, a_col - a_seg * 64, elem_bytes=2), + do_frag = nvvm.ldmatrix( + sDo_ptr + a_seg * (cfg.b_t * 64) + a_row_coord * 64 + swizzle_xor_128b(a_row_coord, a_col - a_seg * 64, elem_bytes=2), 4, nvvm.MMALayout.ROW, ) - b_col = k_block * 16 + rhs_col_offset + b_col = i * 16 + b_col_offset b_seg = b_col // 64 - b_frag = nvvm.ldmatrix( - sU_raw.data_ptr() + b_seg * (cfg.b_t * 64) + rhs_row_coord * 64 + swizzle_xor_128b(rhs_row_coord, b_col - b_seg * 64, elem_bytes=2), + u_frag = nvvm.ldmatrix( + sU_raw.data_ptr() + b_seg * (cfg.b_t * 64) + b_row_coord * 64 + swizzle_xor_128b(b_row_coord, b_col - b_seg * 64, elem_bytes=2), 4, nvvm.MMALayout.ROW, ) mma_step( da_acc, - (a_frag[0], a_frag[1], a_frag[2], a_frag[3]), - (b_frag[0], b_frag[1], b_frag[2], b_frag[3]), + (do_frag[0], do_frag[1], do_frag[2], do_frag[3]), + (u_frag[0], u_frag[1], u_frag[2], u_frag[3]), k_step=0, M=16, N=16, ab_dtype=cfg.io_dtype, ) - # fence: the dO/U ldmatrix reads must complete before this release - # licenses the TMA reload (sDo) nvvm.fence_proxy("async.shared", space="cta") bars.mb_do_done[raw_stage].arrive() for accum_idx in cutlass.range_constexpr(8): da_acc[accum_idx] = da_acc[accum_idx] if (tril_incl_mask >> accum_idx) & 1 else cutlass.Float32(0.0) bars.mb_da_done[intermediate_stage].wait(((gc // cfg.smem_intermediate_stages) + 1) % 2) nvvm.stmatrix( - sIntermediate_ptr + 2 * (cfg.b_t * cfg.b_t) + stsm_idx, + sIntermediate_ptr + 2 * (cfg.b_t * cfg.b_t) + intermediate_idx, [ fp32_to_fp16(da_acc[0], da_acc[1], dtype=cfg.io_dtype), fp32_to_fp16(da_acc[2], da_acc[3], dtype=cfg.io_dtype), @@ -575,7 +516,8 @@ def epilogue_warp( bars.mb_da_ready[intermediate_stage].arrive() raw_index = advance(raw_index, cfg.smem_raw_stages) - # ---- dQ/dK/dGate/dBeta/dV/dW_out: previous chunk, one-behind store ladder ---- + # ---- dQ/dK/dGate/dBeta/dV/dW out: previous chunk, one-behind store + # ladder ------------------------------------------------------------------- if rev_idx > 0: bars.mb_dq_tmastg_ready[dq_index.idx].wait(dq_index.phase) if pend_writes: @@ -634,7 +576,7 @@ def epilogue_warp( pend_start = chunk_start pend_writes = writes - # ---- tile tail: drain the last chunk's dQ/dK/dGate/dBeta/dV/dW_out ------- + # ---- tile tail: store last chunk dQ/dK/dGate/dBeta/dV/dW out ----------------- if sk_nt > 0: bars.mb_dq_tmastg_ready[dq_index.idx].wait(dq_index.phase) if pend_writes: @@ -692,7 +634,7 @@ def epilogue_warp( dwo_index = advance(dwo_index, cfg.smem_dwo_stages) gbase += sk_nt - tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one) @cute.jit @@ -703,8 +645,8 @@ def super_mma_warp( num_ctas, cu_seqlens, mWorkItems, - sSched, - lane, + sScheduler, + lane_idx, sK_decay_raw, sK_inv_raw, sU_raw, @@ -715,41 +657,44 @@ def super_mma_warp( """Super-MMA warp role (warp 12): builds the Neumann T_inv and strict-tril dM staging tiles, in chunk order.""" nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) + elect_one = nvvm.elect_sync() sdy_index = PipelineState.start(phase=0) - # ---- ldmatrix lane decode ---------------------------------------------------- - rhs_row_coord = lane % 8 + (cutlass.Int32(8) if (lane // 16) else cutlass.Int32(0)) - rhs_col_offset = cutlass.Int32(8) if ((lane // 8) % 2) else cutlass.Int32(0) - lhs_row_coord = lane % 8 + (cutlass.Int32(8) if ((lane // 8) % 2) else cutlass.Int32(0)) - lhs_col_offset = cutlass.Int32(8) if ((lane // 8) // 2) else cutlass.Int32(0) - stsm_row_coord = lane & 7 - stsm_col_coord = cutlass.Int32(0) - if (lane // 8) & 1: - stsm_row_coord = stsm_row_coord + cutlass.Int32(8) - if lane // 8 >= 2: - stsm_col_coord = cutlass.Int32(8) - stsm_idx = swizzle_lin_S(stsm_row_coord * cfg.b_t + stsm_col_coord, bbits=1, mbase=3, sshift=3) - row_lo = lane // 4 + # ---- ldmatrix/stmatrix lane decode ----------------------------------------------- + b_row_coord = lane_idx % 8 + (cutlass.Int32(8) if (lane_idx // 16) else cutlass.Int32(0)) + b_col_offset = cutlass.Int32(8) if ((lane_idx // 8) % 2) else cutlass.Int32(0) + a_row_coord = lane_idx % 8 + (cutlass.Int32(8) if ((lane_idx // 8) % 2) else cutlass.Int32(0)) + a_col_offset = cutlass.Int32(8) if ((lane_idx // 8) // 2) else cutlass.Int32(0) + intermediate_row_coord = lane_idx & 7 + intermediate_col_coord = cutlass.Int32(0) + if (lane_idx // 8) & 1: + intermediate_row_coord = intermediate_row_coord + cutlass.Int32(8) + if lane_idx // 8 >= 2: + intermediate_col_coord = cutlass.Int32(8) + intermediate_idx = intermediate_row_coord * cfg.b_t + swizzle_xor_32b(intermediate_row_coord, intermediate_col_coord) + row_lo = lane_idx // 4 row_hi = row_lo + cutlass.Int32(8) - # hoisted tril bitmasks: bit i = row > col / row == col for accum index i + # tril bitmasks: bit i = row > col / row == col for accum index i tril_strict_mask = cutlass.Int32(0) eye_mask = cutlass.Int32(0) for accum_idx in cutlass.range_constexpr(8): row_coord = row_hi if cutlass.const_expr(accum_idx % 4 >= 2) else row_lo - col_coord = (accum_idx // 4) * 8 + 2 * (lane % 4) + col_coord = (accum_idx // 4) * 8 + 2 * (lane_idx % 4) if cutlass.const_expr(accum_idx % 2 == 1): col_coord = col_coord + cutlass.Int32(1) tril_strict_mask = tril_strict_mask | (cutlass.Int32(1 << accum_idx) if row_coord > col_coord else cutlass.Int32(0)) eye_mask = eye_mask | (cutlass.Int32(1 << accum_idx) if row_coord == col_coord else cutlass.Int32(0)) gbase = cutlass.Int32(0) - sched_state = PipelineState.start(phase=0) + scheduler_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) FIRST_STATE_CHUNK = 0 if cfg.use_initial_state else 1 while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) - sk_nt = cend - wstart + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) + sk_nt = compute_end - write_start for rev_idx in cutlass.range(sk_nt, unroll=1): gc = gbase + rev_idx decay_stage = gc % cfg.smem_decay_stages @@ -760,38 +705,38 @@ def super_mma_warp( bars.mb_t_inv_done[intermediate_stage].wait(((gc // cfg.smem_intermediate_stages) + 1) % 2) - # ---- KK = K_decay @ K_inv^T ------------------------------------------ + # ---- KK = K decay @ K inv^T ---------------------------------------------- bars.mb_k_decay_inv_ready[decay_stage].wait((gc // cfg.smem_decay_stages) % 2) - kk_lhs_row = lhs_row_coord + kk_a_row = a_row_coord kk_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): kk_acc[accum_idx] = cutlass.Float32(0.0) - for k_block in cutlass.range_constexpr(cfg.d_k // 16): - a_col = k_block * 16 + lhs_col_offset + for i in cutlass.range_constexpr(cfg.d_k // 16): + a_col = i * 16 + a_col_offset a_seg = a_col // 64 - a_frag = nvvm.ldmatrix( - sK_decay_ptr + a_seg * (cfg.b_t * 64) + kk_lhs_row * 64 + swizzle_xor_128b(kk_lhs_row, a_col - a_seg * 64, elem_bytes=2), + k_decay_frag = nvvm.ldmatrix( + sK_decay_ptr + a_seg * (cfg.b_t * 64) + kk_a_row * 64 + swizzle_xor_128b(kk_a_row, a_col - a_seg * 64, elem_bytes=2), 4, nvvm.MMALayout.ROW, ) - b_col = k_block * 16 + rhs_col_offset + b_col = i * 16 + b_col_offset b_seg = b_col // 64 - b_frag = nvvm.ldmatrix( - sK_inv_ptr + b_seg * (cfg.b_t * 64) + rhs_row_coord * 64 + swizzle_xor_128b(rhs_row_coord, b_col - b_seg * 64, elem_bytes=2), + k_inv_frag = nvvm.ldmatrix( + sK_inv_ptr + b_seg * (cfg.b_t * 64) + b_row_coord * 64 + swizzle_xor_128b(b_row_coord, b_col - b_seg * 64, elem_bytes=2), 4, nvvm.MMALayout.ROW, ) mma_step( kk_acc, - (a_frag[0], a_frag[1], a_frag[2], a_frag[3]), - (b_frag[0], b_frag[1], b_frag[2], b_frag[3]), + (k_decay_frag[0], k_decay_frag[1], k_decay_frag[2], k_decay_frag[3]), + (k_inv_frag[0], k_inv_frag[1], k_inv_frag[2], k_inv_frag[3]), k_step=0, M=16, N=16, ab_dtype=cfg.io_dtype, ) - # ---- L = tril(KK, -1) ------------------------------------------------ + # ---- L = tril(KK, -1) ---------------------------------------------------- l_regs = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): lower = kk_acc[accum_idx] if (tril_strict_mask >> accum_idx) & 1 else cutlass.Float32(0.0) @@ -802,6 +747,7 @@ def super_mma_warp( l_a3 = fp32_to_fp16(l_regs[6], l_regs[7], dtype=cfg.io_dtype) l_values = cutlass.Vector.from_elements((l_a0, l_a1, l_a2, l_a3), cutlass.Int32).bitcast(cfg.io_dtype).to(cutlass.Float32) + # ---- T^-1 = I - L, then three Neumann doubling rounds -------------------- tinv_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): eye = cutlass.Float32(1.0) if (eye_mask >> accum_idx) & 1 else cutlass.Float32(0.0) @@ -809,7 +755,8 @@ def super_mma_warp( lpow_a0, lpow_a1, lpow_a2, lpow_a3 = l_a0, l_a1, l_a2, l_a3 mov_lpow0, mov_lpow1, mov_lpow2, mov_lpow3 = movmatrix_16b(l_a0), movmatrix_16b(l_a1), movmatrix_16b(l_a2), movmatrix_16b(l_a3) - for _round in cutlass.range_constexpr(3): + for neumann_round in cutlass.range_constexpr(3): + # ---- Lpow = Lpow @ Lpow ---------------------------------------------- sq_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): sq_acc[accum_idx] = cutlass.Float32(0.0) @@ -827,6 +774,7 @@ def super_mma_warp( lpow_a2 = fp32_to_fp16(sq_acc[4], sq_acc[5], dtype=cfg.io_dtype) lpow_a3 = fp32_to_fp16(sq_acc[6], sq_acc[7], dtype=cfg.io_dtype) mov_lpow0, mov_lpow1, mov_lpow2, mov_lpow3 = movmatrix_16b(lpow_a0), movmatrix_16b(lpow_a1), movmatrix_16b(lpow_a2), movmatrix_16b(lpow_a3) + # ---- T^-1 += T^-1 @ Lpow --------------------------------------------- upd_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): upd_acc[accum_idx] = cutlass.Float32(0.0) @@ -853,7 +801,7 @@ def super_mma_warp( tinv_acc[6], tinv_acc[7] = fadd2(tinv_lo3, tinv_hi3, upd_acc[6], upd_acc[7]) nvvm.stmatrix( - sIntermediate_ptr + 1 * (cfg.b_t * cfg.b_t) + stsm_idx, + sIntermediate_ptr + 1 * (cfg.b_t * cfg.b_t) + intermediate_idx, [ fp32_to_fp16(tinv_acc[0], tinv_acc[1], dtype=cfg.io_dtype), fp32_to_fp16(tinv_acc[2], tinv_acc[3], dtype=cfg.io_dtype), @@ -866,32 +814,32 @@ def super_mma_warp( nvvm.fence_proxy("async.shared", space="cta") bars.mb_t_inv_ready[intermediate_stage].arrive() - # ---- dM = dY @ U^T --------------------------------------------------- + # ---- dM = dY @ U^T ------------------------------------------------------- bars.mb_dm_done[intermediate_stage].wait(((gc // cfg.smem_intermediate_stages) + 1) % 2) bars.mb_dy_smem_ready.wait(sdy_index.phase) sdy_index = advance(sdy_index, 1) dm_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): dm_acc[accum_idx] = cutlass.Float32(0.0) - for k_block in cutlass.range_constexpr(cfg.d_v // 16): - a_col = k_block * 16 + lhs_col_offset + for i in cutlass.range_constexpr(cfg.d_v // 16): + a_col = i * 16 + a_col_offset a_seg = a_col // 64 - a_frag = nvvm.ldmatrix( - sDy_raw.data_ptr() + a_seg * (cfg.b_t * 64) + lhs_row_coord * 64 + swizzle_xor_128b(lhs_row_coord, a_col - a_seg * 64, elem_bytes=2), + dy_frag = nvvm.ldmatrix( + sDy_raw.data_ptr() + a_seg * (cfg.b_t * 64) + a_row_coord * 64 + swizzle_xor_128b(a_row_coord, a_col - a_seg * 64, elem_bytes=2), 4, nvvm.MMALayout.ROW, ) - b_col = k_block * 16 + rhs_col_offset + b_col = i * 16 + b_col_offset b_seg = b_col // 64 - b_frag = nvvm.ldmatrix( - sU_raw.data_ptr() + b_seg * (cfg.b_t * 64) + rhs_row_coord * 64 + swizzle_xor_128b(rhs_row_coord, b_col - b_seg * 64, elem_bytes=2), + u_frag = nvvm.ldmatrix( + sU_raw.data_ptr() + b_seg * (cfg.b_t * 64) + b_row_coord * 64 + swizzle_xor_128b(b_row_coord, b_col - b_seg * 64, elem_bytes=2), 4, nvvm.MMALayout.ROW, ) mma_step( dm_acc, - (a_frag[0], a_frag[1], a_frag[2], a_frag[3]), - (b_frag[0], b_frag[1], b_frag[2], b_frag[3]), + (dy_frag[0], dy_frag[1], dy_frag[2], dy_frag[3]), + (u_frag[0], u_frag[1], u_frag[2], u_frag[3]), k_step=0, M=16, N=16, @@ -904,16 +852,16 @@ def super_mma_warp( w1 = fp32_to_fp16(dm_strict_regs[2], dm_strict_regs[3], dtype=cfg.io_dtype) w2 = fp32_to_fp16(dm_strict_regs[4], dm_strict_regs[5], dtype=cfg.io_dtype) w3 = fp32_to_fp16(dm_strict_regs[6], dm_strict_regs[7], dtype=cfg.io_dtype) - nvvm.stmatrix(sIntermediate_ptr + 3 * (cfg.b_t * cfg.b_t) + stsm_idx, [w0, w1, w2, w3], nvvm.MMALayout.ROW, shape=nvvm.StoreShape.M8N8) + nvvm.stmatrix(sIntermediate_ptr + 3 * (cfg.b_t * cfg.b_t) + intermediate_idx, [w0, w1, w2, w3], nvvm.MMALayout.ROW, shape=nvvm.StoreShape.M8N8) nw0 = fp32_to_fp16(-dm_strict_regs[0], -dm_strict_regs[1], dtype=cfg.io_dtype) nw1 = fp32_to_fp16(-dm_strict_regs[2], -dm_strict_regs[3], dtype=cfg.io_dtype) nw2 = fp32_to_fp16(-dm_strict_regs[4], -dm_strict_regs[5], dtype=cfg.io_dtype) nw3 = fp32_to_fp16(-dm_strict_regs[6], -dm_strict_regs[7], dtype=cfg.io_dtype) - nvvm.stmatrix(sIntermediate_ptr + 4 * (cfg.b_t * cfg.b_t) + stsm_idx, [nw0, nw1, nw2, nw3], nvvm.MMALayout.ROW, shape=nvvm.StoreShape.M8N8) + nvvm.stmatrix(sIntermediate_ptr + 4 * (cfg.b_t * cfg.b_t) + intermediate_idx, [nw0, nw1, nw2, nw3], nvvm.MMALayout.ROW, shape=nvvm.StoreShape.M8N8) nvvm.fence_proxy("async.shared", space="cta") bars.mb_dm_ready[intermediate_stage].arrive() gbase += sk_nt - tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one) @cute.jit @@ -924,20 +872,20 @@ def tcgen05_mma_warp( num_ctas, cu_seqlens, mWorkItems, - sSched, + sScheduler, tmem_hold, - sState_alt, - sState_direct, - sK_decay_lead16, - sK_inv_amaj, - sK_restore_lead16, - sDo_lead16, - sDo_amaj, + sState_trans, + sState, + sK_decay, + sK_inv_trans, + sK_restore, + sDo, + sDo_trans, sQ_decay_trans, sK_decay_trans, - sU_lead16, - sDy_lead16, - sDstate_alt, + sU, + sDy, + sDstate_trans, sIntermediate, sState_scale_diag, bars, @@ -952,7 +900,7 @@ def tcgen05_mma_warp( tmem_base = tmem_hold.load() bpe = cfg.io_dtype.width // 8 - # ---- chunk-invariant GEMM descriptors ---------------------------------------- + # ---- chunk-invariant GEMM descriptors -------------------------------------------- idesc_mv_nt = nvvm.Tcgen05InstrDesc.build( c_dtype=cutlass.Float32, a_dtype=cfg.io_dtype, @@ -968,7 +916,7 @@ def tcgen05_mma_warp( m_dim=cfg.d_v, a_major=1, ) - bmm_state_k_desc = MmaDesc( + bmm_state_k_decay_desc = MmaDesc( M=cfg.d_v, N=cfg.b_t, K=cfg.d_k, @@ -981,14 +929,14 @@ def tcgen05_mma_warp( idesc=idesc_state_k_at, kind=nvvm.Tcgen05MMAKind.F16, ) - idesc_state_k_kmaj = nvvm.Tcgen05InstrDesc.build( + idesc_state_k = nvvm.Tcgen05InstrDesc.build( c_dtype=cutlass.Float32, a_dtype=cfg.io_dtype, b_dtype=cfg.io_dtype, n_dim=cfg.b_t, m_dim=cfg.d_v, ) - bmm_state_k_kmaj_desc = MmaDesc( + bmm_state_k_decay_desc = MmaDesc( M=cfg.d_v, N=cfg.b_t, K=cfg.d_k, @@ -998,10 +946,10 @@ def tcgen05_mma_warp( btranspose=False, atranspose=False, cta_group=1, - idesc=idesc_state_k_kmaj, + idesc=idesc_state_k, kind=nvvm.Tcgen05MMAKind.F16, ) - bmm_dvinter_desc = MmaDesc( + bmm_dstate_k_restore_desc = MmaDesc( M=cfg.d_v, N=cfg.b_t, K=cfg.d_k, @@ -1022,7 +970,7 @@ def tcgen05_mma_warp( a_major=1, b_major=1, ) - bmm_du_at_desc = MmaDesc( + bmm_do_a_desc = MmaDesc( M=cfg.d_v, N=cfg.b_t, K=cfg.b_t, @@ -1044,7 +992,7 @@ def tcgen05_mma_warp( a_major=1, b_major=1, ) - bmm_dstate_q_at_desc = MmaDesc( + bmm_do_q_decay_desc = MmaDesc( M=cfg.d_v, N=cfg.d_k, K=cfg.b_t, @@ -1057,7 +1005,7 @@ def tcgen05_mma_warp( idesc=idesc_dstate_q_at, kind=nvvm.Tcgen05MMAKind.F16, ) - bmm_qk_ts_desc = MmaDesc( + bmm_y_t_inv_desc = MmaDesc( M=cfg.d_v, N=cfg.b_t, K=cfg.b_t, @@ -1077,7 +1025,7 @@ def tcgen05_mma_warp( m_dim=cfg.d_v, b_major=1, ) - bmm_qk_ts_t_desc = MmaDesc( + bmm_du_t_inv_trans_desc = MmaDesc( M=cfg.d_v, N=cfg.b_t, K=cfg.b_t, @@ -1096,7 +1044,7 @@ def tcgen05_mma_warp( n_dim=16, m_dim=cfg.d_v, ) - bmm_diag_desc = MmaDesc( + bmm_dstate_diag_desc = MmaDesc( M=cfg.d_v, N=16, K=16, @@ -1116,7 +1064,7 @@ def tcgen05_mma_warp( m_dim=cfg.d_v, b_major=1, ) - bmm_dstate_k_desc = MmaDesc( + bmm_dy_k_decay_desc = MmaDesc( M=cfg.d_v, N=cfg.d_k, K=cfg.b_t, @@ -1135,19 +1083,7 @@ def tcgen05_mma_warp( n_dim=cfg.b_t, m_dim=cfg.d_k, ) - bmm_state_desc = MmaDesc( - M=cfg.d_k, - N=cfg.b_t, - K=cfg.d_v, - bpe_a=bpe, - bpe_b=bpe, - tile_k_hw=16, - btranspose=False, - cta_group=1, - idesc=idesc_state_ts, - kind=nvvm.Tcgen05MMAKind.F16, - ) - bmm_state_ts_desc = MmaDesc( + bmm_state_dy_desc = MmaDesc( M=cfg.d_k, N=cfg.b_t, K=cfg.d_v, @@ -1159,6 +1095,7 @@ def tcgen05_mma_warp( idesc=idesc_state_ts, kind=nvvm.Tcgen05MMAKind.F16, ) + bmm_state_do_desc = bmm_state_dy_desc idesc_dstate_at = nvvm.Tcgen05InstrDesc.build( c_dtype=cutlass.Float32, a_dtype=cfg.io_dtype, @@ -1167,7 +1104,7 @@ def tcgen05_mma_warp( m_dim=cfg.d_k, a_major=1, ) - bmm_dstate_at_desc = MmaDesc( + bmm_dstate_u_desc = MmaDesc( M=cfg.d_k, N=cfg.b_t, K=cfg.d_v, @@ -1187,18 +1124,6 @@ def tcgen05_mma_warp( n_dim=cfg.b_t, m_dim=cfg.d_k, ) - bmm_dgrad_ts_desc = MmaDesc( - M=cfg.d_k, - N=cfg.b_t, - K=cfg.b_t, - bpe_a=bpe, - bpe_b=bpe, - tile_k_hw=16, - btranspose=False, - cta_group=1, - idesc=idesc_dgp, - kind=nvvm.Tcgen05MMAKind.F16, - ) idesc_dgp_at = nvvm.Tcgen05InstrDesc.build( c_dtype=cutlass.Float32, a_dtype=cfg.io_dtype, @@ -1207,7 +1132,7 @@ def tcgen05_mma_warp( m_dim=cfg.d_k, a_major=1, ) - bmm_dgrad_at_desc = MmaDesc( + bmm_k_inv_dm_desc = MmaDesc( M=cfg.d_k, N=cfg.b_t, K=cfg.b_t, @@ -1220,6 +1145,7 @@ def tcgen05_mma_warp( idesc=idesc_dgp_at, kind=nvvm.Tcgen05MMAKind.F16, ) + bmm_k_inv_da_desc = bmm_k_inv_dm_desc idesc_dgp_at_t = nvvm.Tcgen05InstrDesc.build( c_dtype=cutlass.Float32, a_dtype=cfg.io_dtype, @@ -1229,7 +1155,7 @@ def tcgen05_mma_warp( a_major=1, b_major=1, ) - bmm_dgrad_at_t_desc = MmaDesc( + bmm_k_decay_dm_desc = MmaDesc( M=cfg.d_k, N=cfg.b_t, K=cfg.b_t, @@ -1242,11 +1168,12 @@ def tcgen05_mma_warp( idesc=idesc_dgp_at_t, kind=nvvm.Tcgen05MMAKind.F16, ) + bmm_q_decay_da_desc = bmm_k_decay_dm_desc state_index = PipelineState.start(phase=0) - y_inp_index = PipelineState.start(phase=0) - dstate_inp_index = PipelineState.start(phase=0) - du_inp_index = PipelineState.start(phase=0) + y_input_index = PipelineState.start(phase=0) + dstate_input_index = PipelineState.start(phase=0) + du_input_index = PipelineState.start(phase=0) neg_dy_index = PipelineState.start(phase=0) u_smem_index = PipelineState.start(phase=0) dstate_smem_index = PipelineState.start(phase=0) @@ -1257,30 +1184,32 @@ def tcgen05_mma_warp( intermediate_seg = (cfg.intermediate_tiles * cfg.b_t * cfg.b_t * (cfg.io_dtype.width // 8)) >> 4 intermediate_slot = (cfg.b_t * cfg.b_t * (cfg.io_dtype.width // 8)) >> 4 diag_seg = ((cfg.d_k // 16) * 256 * (cfg.io_dtype.width // 8)) >> 4 - d_do_amaj0 = sDo_amaj[0].desc() + d_do_trans0 = sDo_trans[0].desc() d_qd_trans0 = sQ_decay_trans[0].desc() d_kd_trans0 = sK_decay_trans[0].desc() - d_ki_amaj0 = sK_inv_amaj[0].desc() + d_ki_trans0 = sK_inv_trans[0].desc() d_int0 = sIntermediate[0].desc() - d_kd_lead0 = sK_decay_lead16[0].desc() - d_do_lead0 = sDo_lead16[0].desc() - d_kr_lead0 = sK_restore_lead16[0].desc() + d_kd0 = sK_decay[0].desc() + d_do0 = sDo[0].desc() + d_kr0 = sK_restore[0].desc() d_diag0 = sState_scale_diag[0].desc() - d_dstate_alt0 = sDstate_alt[0].desc() - d_u_lead0 = sU_lead16[0].desc() - d_dy_lead0 = sDy_lead16[0].desc() + d_dstate_trans0 = sDstate_trans[0].desc() + d_u0 = sU[0].desc() + d_dy0 = sDy[0].desc() assert cfg.smem_state_stages == 1 - d_state_alt0 = sState_alt[0].desc() - d_state_direct0 = sState_direct[0].desc() + d_state_trans0 = sState_trans[0].desc() + d_state0 = sState[0].desc() dstate0_index = PipelineState.start(phase=0) gbase = cutlass.Int32(0) - sched_state = PipelineState.start(phase=0) + scheduler_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) FIRST_STATE_CHUNK = 0 if cfg.use_initial_state else 1 while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) - sk_nt = cend - wstart + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) + sk_nt = compute_end - write_start for rev_idx in cutlass.range(sk_nt, unroll=1): gc = gbase + rev_idx decay_stage = gc % cfg.smem_decay_stages @@ -1292,27 +1221,27 @@ def tcgen05_mma_warp( has_dstate = cutlass.Boolean(True) raw_stage_idx = gc % cfg.smem_raw_stages - # ---- stage-derived operand descriptors ------------------------------- + # ---- stage-derived operand descriptors ----------------------------------- decay_op_off = decay_stage * op_seg - d_do_amaj = d_do_amaj0 + raw_stage_idx * do_seg + d_do_trans = d_do_trans0 + raw_stage_idx * do_seg d_qd_trans = d_qd_trans0 + decay_op_off d_kd_trans = d_kd_trans0 + decay_op_off - d_ki_amaj = d_ki_amaj0 + decay_op_off + d_ki_trans = d_ki_trans0 + decay_op_off d_int = d_int0 + intermediate_stage * intermediate_seg d_int_tinv = d_int + intermediate_slot d_int_da = d_int + 2 * intermediate_slot d_int_dm = d_int + 3 * intermediate_slot d_int_ndm = d_int + 4 * intermediate_slot - chunk_idx = cend - cutlass.Int32(1) - rev_idx + chunk_idx = compute_end - cutlass.Int32(1) - rev_idx - # ---- state_k = state(S) @ K_decay^T ---------------------------------- + # ---- k state = state(S) @ K decay^T -------------------------------------- bars.mb_k_decay_inv_ready[decay_stage].wait(decay_phase) if chunk_idx >= FIRST_STATE_CHUNK: bars.mb_state_ready[state_index.idx].wait(state_index.phase) mma_ss( - bmm_state_k_kmaj_desc, - d_state_direct0, - d_kd_lead0 + decay_op_off, + bmm_state_k_decay_desc, + d_state0, + d_kd0 + decay_op_off, nvvm.make_tmem_ptr((tmem_base + cfg.tmem_state_k_acc_offset), cutlass.Float32), accumulate=False, ) @@ -1321,68 +1250,68 @@ def tcgen05_mma_warp( bars.mb_state_done[state_index.idx].arrive(cta_group=1) state_index = advance(state_index, cfg.smem_state_stages) - # ---- dQ inter = state(T) @ dO^T -------------------------------------- + # ---- dQ inter = state(T) @ dO^T ------------------------------------------ bars.mb_dqk_acc_done.wait(parts_done_index.phase) parts_done_index = advance(parts_done_index, 1) - bars.mb_state_inp_ready[gc % 2].wait((gc // 2) % 2) + bars.mb_state_input_ready[gc % 2].wait((gc // 2) % 2) bars.mb_do_ready[raw_stage_idx].wait((gc // cfg.smem_raw_stages) % 2) if chunk_idx >= FIRST_STATE_CHUNK: - a_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_state_inp_offset + (gc % 2) * (cfg.d_v // 2)), cutlass.Int8) - b_desc = d_do_lead0 + raw_stage_idx * do_seg + a_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_state_input_offset + (gc % 2) * (cfg.d_v // 2)), cutlass.Int8) + b_desc = d_do0 + raw_stage_idx * do_seg c_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dq_acc_offset), cutlass.Float32) - for sub in cutlass.range_constexpr(bmm_state_ts_desc.num_subtiles_B): - for k in cutlass.range_constexpr(bmm_state_ts_desc.sps_B): + for i in cutlass.range_constexpr(bmm_state_do_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_state_do_desc.sps_B): mma_ts_step( - bmm_state_ts_desc, - a_ptr.subview(sub * bmm_state_ts_desc.sps_B * bmm_state_ts_desc.tmem_advance_A), - b_desc + sub * (bmm_state_ts_desc.smem_subtile_B >> 4), + bmm_state_do_desc, + a_ptr.subview(i * bmm_state_do_desc.sps_B * bmm_state_do_desc.tmem_advance_A), + b_desc + i * (bmm_state_do_desc.smem_subtile_B >> 4), c_ptr, k, - cutlass.Boolean(sub + k > 0), + cutlass.Boolean(i + k > 0), ) - # ---- dU inter = dstate_inp(T) @ K_restore ---------------------------- + # ---- dU inter = dstate input(T) @ K restore ------------------------------ bars.mb_q_decay_k_restore_ready[decay_stage].wait(decay_phase) if has_dstate: - bars.mb_dstate_inp_ready.wait(dstate_inp_index.phase) - a_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dstate_inp_offset), cutlass.Int8) - b_desc = d_kr_lead0 + decay_op_off + bars.mb_dstate_input_ready.wait(dstate_input_index.phase) + a_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dstate_input_offset), cutlass.Int8) + b_desc = d_kr0 + decay_op_off c_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_du_acc_offset), cutlass.Float32) - for sub in cutlass.range_constexpr(bmm_dvinter_desc.num_subtiles_B): - for k in cutlass.range_constexpr(bmm_dvinter_desc.sps_B): + for i in cutlass.range_constexpr(bmm_dstate_k_restore_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_dstate_k_restore_desc.sps_B): mma_ts_step( - bmm_dvinter_desc, - a_ptr.subview(sub * bmm_dvinter_desc.sps_B * bmm_dvinter_desc.tmem_advance_A), - b_desc + sub * (bmm_dvinter_desc.smem_subtile_B >> 4), + bmm_dstate_k_restore_desc, + a_ptr.subview(i * bmm_dstate_k_restore_desc.sps_B * bmm_dstate_k_restore_desc.tmem_advance_A), + b_desc + i * (bmm_dstate_k_restore_desc.smem_subtile_B >> 4), c_ptr, k, - cutlass.Boolean(sub + k > 0), + cutlass.Boolean(i + k > 0), ) - # ---- dstate decay = dstate_inp(T) @ diag(eGl) ------------------------ + # ---- dstate decay = dstate input(T) @ diag(eGl) -------------------------- if has_dstate: desc_diag = d_diag0 + decay_stage * diag_seg - for k_block in cutlass.range_constexpr(cfg.d_k // 16): - a_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dstate_inp_offset) + k_block * 8, cutlass.Int8) - b_desc = desc_diag.advance_start_address(k_block * 256 * 2) - c_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dstate_acc_offset) + k_block * 16, cutlass.Float32) - for sub in cutlass.range_constexpr(bmm_diag_desc.num_subtiles_B): - for k in cutlass.range_constexpr(bmm_diag_desc.sps_B): + for i in cutlass.range_constexpr(cfg.d_k // 16): + a_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dstate_input_offset) + i * 8, cutlass.Int8) + b_desc = desc_diag.advance_start_address(i * 256 * 2) + c_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dstate_acc_offset) + i * 16, cutlass.Float32) + for i in cutlass.range_constexpr(bmm_dstate_diag_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_dstate_diag_desc.sps_B): mma_ts_step( - bmm_diag_desc, - a_ptr.subview(sub * bmm_diag_desc.sps_B * bmm_diag_desc.tmem_advance_A), - b_desc + sub * (bmm_diag_desc.smem_subtile_B >> 4), + bmm_dstate_diag_desc, + a_ptr.subview(i * bmm_dstate_diag_desc.sps_B * bmm_dstate_diag_desc.tmem_advance_A), + b_desc + i * (bmm_dstate_diag_desc.smem_subtile_B >> 4), c_ptr, k, - cutlass.Boolean(sub + k > 0), + cutlass.Boolean(i + k > 0), ) - dstate_inp_index = advance(dstate_inp_index, 1) + dstate_input_index = advance(dstate_input_index, 1) - # ---- dU intra += dO^T(S) @ A ----------------------------------------- + # ---- dU intra += dO^T(S) @ A --------------------------------------------- bars.mb_a_ready[intermediate_stage].wait(intermediate_phase) mma_ss( - bmm_du_at_desc, - d_do_amaj, + bmm_do_a_desc, + d_do_trans, d_int, nvvm.make_tmem_ptr((tmem_base + cfg.tmem_du_acc_offset), cutlass.Float32), accumulate=has_dstate, @@ -1391,66 +1320,66 @@ def tcgen05_mma_warp( bars.mb_du_acc_ready.arrive(cta_group=1) bars.mb_a_done[intermediate_stage].arrive(cta_group=1) - # ---- dstate Q-term += dO^T(S) @ Q_decay ------------------------------ + # ---- dstate Q-term += dO^T(S) @ Q decay ---------------------------------- mma_ss( - bmm_dstate_q_at_desc, - d_do_amaj, + bmm_do_q_decay_desc, + d_do_trans, d_qd_trans, nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dstate_acc_offset), cutlass.Float32), accumulate=has_dstate, ) - # ---- U = Y(T) @ T_inv ------------------------------------------------ + # ---- U = Y(T) @ T^-1 ----------------------------------------------------- bars.mb_t_inv_ready[intermediate_stage].wait(intermediate_phase) - bars.mb_y_inp_ready.wait(y_inp_index.phase) - y_inp_index = advance(y_inp_index, 1) - a_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_y_inp_offset), cutlass.Int8) + bars.mb_y_input_ready.wait(y_input_index.phase) + y_input_index = advance(y_input_index, 1) + a_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_y_input_offset), cutlass.Int8) b_desc = d_int_tinv c_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_u_acc_offset), cutlass.Float32) - for sub in cutlass.range_constexpr(bmm_qk_ts_desc.num_subtiles_B): - for k in cutlass.range_constexpr(bmm_qk_ts_desc.sps_B): + for i in cutlass.range_constexpr(bmm_y_t_inv_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_y_t_inv_desc.sps_B): mma_ts_step( - bmm_qk_ts_desc, - a_ptr.subview(sub * bmm_qk_ts_desc.sps_B * bmm_qk_ts_desc.tmem_advance_A), - b_desc + sub * (bmm_qk_ts_desc.smem_subtile_B >> 4), + bmm_y_t_inv_desc, + a_ptr.subview(i * bmm_y_t_inv_desc.sps_B * bmm_y_t_inv_desc.tmem_advance_A), + b_desc + i * (bmm_y_t_inv_desc.smem_subtile_B >> 4), c_ptr, k, - cutlass.Boolean(sub + k > 0), + cutlass.Boolean(i + k > 0), ) if elect_one: bars.mb_u_acc_ready.arrive(cta_group=1) bars.mb_do_mma_done[raw_stage_idx].arrive(cta_group=1) - # ---- dY = dU(T) @ T_inv ---------------------------------------------- - bars.mb_du_inp_ready.wait(du_inp_index.phase) - du_inp_index = advance(du_inp_index, 1) - a_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_du_inp_offset), cutlass.Int8) + # ---- dY = dU(T) @ T^-1 --------------------------------------------------- + bars.mb_du_input_ready.wait(du_input_index.phase) + du_input_index = advance(du_input_index, 1) + a_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_du_input_offset), cutlass.Int8) b_desc = d_int_tinv c_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dy_acc_offset), cutlass.Float32) - for sub in cutlass.range_constexpr(bmm_qk_ts_t_desc.num_subtiles_B): - for k in cutlass.range_constexpr(bmm_qk_ts_t_desc.sps_B): + for i in cutlass.range_constexpr(bmm_du_t_inv_trans_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_du_t_inv_trans_desc.sps_B): mma_ts_step( - bmm_qk_ts_t_desc, - a_ptr.subview(sub * bmm_qk_ts_t_desc.sps_B * bmm_qk_ts_t_desc.tmem_advance_A), - b_desc + sub * (bmm_qk_ts_t_desc.smem_subtile_B >> 4), + bmm_du_t_inv_trans_desc, + a_ptr.subview(i * bmm_du_t_inv_trans_desc.sps_B * bmm_du_t_inv_trans_desc.tmem_advance_A), + b_desc + i * (bmm_du_t_inv_trans_desc.smem_subtile_B >> 4), c_ptr, k, - cutlass.Boolean(sub + k > 0), + cutlass.Boolean(i + k > 0), ) if elect_one: bars.mb_dy_acc_ready.arrive(cta_group=1) bars.mb_t_inv_done[intermediate_stage].arrive(cta_group=1) - # ---- dK_restore part = dstate(S) @ U^T ------------------------------- + # ---- dK restore part = dstate(S) @ U^T ----------------------------------- bars.mb_u_smem_ready.wait(u_smem_index.phase) u_smem_index = advance(u_smem_index, 1) if has_dstate: bars.mb_dstate_smem_ready.wait(dstate_smem_index.phase) dstate_smem_index = advance(dstate_smem_index, 1) mma_ss( - bmm_dstate_at_desc, - d_dstate_alt0, - d_u_lead0, + bmm_dstate_u_desc, + d_dstate_trans0, + d_u0, nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dk_restore_acc_offset), cutlass.Float32), accumulate=False, ) @@ -1458,18 +1387,18 @@ def tcgen05_mma_warp( bars.mb_dk_restore_part_acc_ready.arrive(cta_group=1) bars.mb_dstate_smem_done.arrive(cta_group=1) - # ---- dstate K-term += -dY(T) @ K_decay ------------------------------- - bars.mb_neg_dy_inp_ready.wait(neg_dy_index.phase) + # ---- dstate K-term += -dY(T) @ K decay ----------------------------------- + bars.mb_neg_dy_input_ready.wait(neg_dy_index.phase) neg_dy_index = advance(neg_dy_index, 1) - a_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_neg_dy_inp_offset), cutlass.Int8) + a_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_neg_dy_input_offset), cutlass.Int8) b_desc = d_kd_trans c_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dstate_acc_offset), cutlass.Float32) - for sub in cutlass.range_constexpr(bmm_dstate_k_desc.num_subtiles_B): - for k in cutlass.range_constexpr(bmm_dstate_k_desc.sps_B): + for i in cutlass.range_constexpr(bmm_dy_k_decay_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_dy_k_decay_desc.sps_B): mma_ts_step( - bmm_dstate_k_desc, - a_ptr.subview(sub * bmm_dstate_k_desc.sps_B * bmm_dstate_k_desc.tmem_advance_A), - b_desc + sub * (bmm_dstate_k_desc.smem_subtile_B >> 4), + bmm_dy_k_decay_desc, + a_ptr.subview(i * bmm_dy_k_decay_desc.sps_B * bmm_dy_k_decay_desc.tmem_advance_A), + b_desc + i * (bmm_dy_k_decay_desc.smem_subtile_B >> 4), c_ptr, k, cutlass.Boolean(True), @@ -1477,29 +1406,29 @@ def tcgen05_mma_warp( if elect_one: bars.mb_dstate_acc_ready.arrive(cta_group=1) - # ---- dK_inv part = scale.Q_decay^T(S) @ dA --------------------------- + # ---- dK inv part = scale.Q decay^T(S) @ dA ------------------------------- bars.mb_da_ready[intermediate_stage].wait(intermediate_phase) mma_ss( - bmm_dgrad_at_t_desc, + bmm_q_decay_da_desc, d_qd_trans, d_int_da, nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dk_inv_acc_offset), cutlass.Float32), accumulate=False, ) - # ---- dQ attn += K_inv^T(S) @ dA^T ------------------------------------ + # ---- dQ attn += K inv^T(S) @ dA^T ---------------------------------------- if chunk_idx >= FIRST_STATE_CHUNK: mma_ss( - bmm_dgrad_at_desc, - d_ki_amaj, + bmm_k_inv_da_desc, + d_ki_trans, d_int_da, nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dq_acc_offset), cutlass.Float32), accumulate=True, ) if chunk_idx < FIRST_STATE_CHUNK: mma_ss( - bmm_dgrad_at_desc, - d_ki_amaj, + bmm_k_inv_da_desc, + d_ki_trans, d_int_da, nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dq_acc_offset), cutlass.Float32), accumulate=False, @@ -1508,30 +1437,30 @@ def tcgen05_mma_warp( bars.mb_dq_acc_ready.arrive(cta_group=1) bars.mb_da_done[intermediate_stage].arrive(cta_group=1) - # ---- dK_decay part = state(T) @ dY^T --------------------------------- + # ---- dK decay part = state(T) @ dY^T ------------------------------------- + bars.mb_dy_smem_ready.wait(gc % 2) if chunk_idx >= FIRST_STATE_CHUNK: - bars.mb_dy_smem_ready.wait(gc % 2) - a_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_state_inp_offset + (gc % 2) * (cfg.d_v // 2)), cutlass.Int8) - b_desc = d_dy_lead0 + a_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_state_input_offset + (gc % 2) * (cfg.d_v // 2)), cutlass.Int8) + b_desc = d_dy0 c_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dk_decay_acc_offset), cutlass.Float32) - for sub in cutlass.range_constexpr(bmm_state_ts_desc.num_subtiles_B): - for k in cutlass.range_constexpr(bmm_state_ts_desc.sps_B): + for i in cutlass.range_constexpr(bmm_state_dy_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_state_dy_desc.sps_B): mma_ts_step( - bmm_state_ts_desc, - a_ptr.subview(sub * bmm_state_ts_desc.sps_B * bmm_state_ts_desc.tmem_advance_A), - b_desc + sub * (bmm_state_ts_desc.smem_subtile_B >> 4), + bmm_state_dy_desc, + a_ptr.subview(i * bmm_state_dy_desc.sps_B * bmm_state_dy_desc.tmem_advance_A), + b_desc + i * (bmm_state_dy_desc.smem_subtile_B >> 4), c_ptr, k, - cutlass.Boolean(sub + k > 0), + cutlass.Boolean(i + k > 0), ) if elect_one: bars.mb_dy_smem_done.arrive(cta_group=1) - bars.mb_state_inp_done[gc % 2].arrive(cta_group=1) + bars.mb_state_input_done[gc % 2].arrive(cta_group=1) - # ---- dK_inv part += K_decay^T(S) @ -dM_strict ------------------------ + # ---- dK inv part += K decay^T(S) @ -dM strict ---------------------------- bars.mb_dm_ready[intermediate_stage].wait(intermediate_phase) mma_ss( - bmm_dgrad_at_t_desc, + bmm_k_decay_dm_desc, d_kd_trans, d_int_ndm, nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dk_inv_acc_offset), cutlass.Float32), @@ -1540,19 +1469,19 @@ def tcgen05_mma_warp( if elect_one: bars.mb_dk_inv_part_acc_ready.arrive(cta_group=1) - # ---- dK_decay part += K_inv^T(S) @ dM_strict^T ----------------------- + # ---- dK decay part += K inv^T(S) @ dM strict^T --------------------------- if chunk_idx >= FIRST_STATE_CHUNK: mma_ss( - bmm_dgrad_at_desc, - d_ki_amaj, + bmm_k_inv_dm_desc, + d_ki_trans, d_int_dm, nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dk_decay_acc_offset), cutlass.Float32), accumulate=True, ) if chunk_idx < FIRST_STATE_CHUNK: mma_ss( - bmm_dgrad_at_desc, - d_ki_amaj, + bmm_k_inv_dm_desc, + d_ki_trans, d_int_dm, nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dk_decay_acc_offset), cutlass.Float32), accumulate=False, @@ -1562,11 +1491,12 @@ def tcgen05_mma_warp( bars.mb_dm_done[intermediate_stage].arrive(cta_group=1) bars.mb_decay_done[decay_stage].arrive(cta_group=1) - # ---- tile end: WG1's dstate0 drain gates the next tile's dstate reuse ---- - bars.mb_dstate0_acc_stored.wait(dstate0_index.phase) - dstate0_index = advance(dstate0_index, 1) + # ---- tile end: WG1's dstate0 store gates the next tile's dstate reuse -------- + if sk_nt > 0: + bars.mb_dstate0_acc_stored.wait(dstate0_index.phase) + dstate0_index = advance(dstate0_index, 1) gbase += sk_nt - tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one) bars.mb_tmem_done[0].wait(0) nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_1) nvvm.tcgen05_dealloc( @@ -1584,8 +1514,8 @@ def tmaldg_warp( num_ctas, cu_seqlens, mWorkItems, - mSched, - sSched, + mScheduler, + sScheduler, q_tx_bytes, k_tx_bytes, gate_tx_bytes, @@ -1613,10 +1543,17 @@ def tmaldg_warp( ) -> None: """TMA-LDG warp role (warp 14): persistent scheduler loop issuing every G->S operand load.""" - elect_one = nvvm.elect_sync() - nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) + raw_index = PipelineState.start(phase=1) + state_index = PipelineState.start(phase=1) + scheduler_state = PipelineState.start(phase=1) + tail_count = ((total_tiles - cutlass.Int32(1)) % num_ctas) + cutlass.Int32(1) + tail_base = (total_tiles - tail_count) if tail_count * 2 >= num_ctas else total_tiles + tail_row = tail_base + cute.arch.smid() + tail_row = tail_row if tail_row < total_tiles else cutlass.Int32(1 << 28) + + elect_one = nvvm.elect_sync() sQ_tma = SmemTile( base=sQ_raw, elems_per_stage=(cfg.d_k * cfg.b_t), @@ -1705,14 +1642,15 @@ def tmaldg_warp( tma_granu_elems=64, tma_subtile_stride_elems=cfg.d_k * 64, ) - raw_index = PipelineState.start(phase=1) - state_index = PipelineState.start(phase=1) - sched_state = PipelineState.start(phase=1) tile_idx = cutlass.Int32(bidx) FIRST_STATE_CHUNK = 0 if cfg.use_initial_state else 1 while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) - next_tile, sched_state = sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ctas) + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) + next_tile, scheduler_state = scheduler_publish_next( + cfg, bars, sScheduler, mScheduler, scheduler_state, tile_idx, num_ctas, tail_base, tail_row, elect_one + ) head_o = head_idx head_q = head_idx if cfg.q_ratio == 1 else head_idx // cutlass.Int32(cfg.q_ratio) head_k = head_idx if cfg.k_ratio == 1 else head_idx // cutlass.Int32(cfg.k_ratio) @@ -1735,40 +1673,40 @@ def tmaldg_warp( tma_tensormap_acquire(desc_beta_slot) tma_tensormap_acquire(desc_w_slot) tma_tensormap_acquire(desc_checkpoint_slot) - sk_nt = cend - wstart + sk_nt = compute_end - write_start for rev_idx in cutlass.range(sk_nt, unroll=1): - chunk_idx = cend - cutlass.Int32(1) - rev_idx + chunk_idx = compute_end - cutlass.Int32(1) - rev_idx chunk_start = chunk_idx * cfg.b_t - # ---- Q load ---------------------------------------------------------- + # ---- Q load -------------------------------------------------------------- bars.mb_q_done[raw_index.idx].wait(raw_index.phase) if elect_one: bars.mb_q_ready[raw_index.idx].arrive(n_bytes=q_tx_bytes) q_slice = tma_slice_runtime_desc(desc_q_slot, cutlass.Int32(0), head_q, chunk_start) tma_load_tile(sQ_tma[raw_index.idx], q_slice, bars.mb_q_ready[raw_index.idx].smem_ptr, acquire=False) - # ---- K load ---------------------------------------------------------- + # ---- K load -------------------------------------------------------------- bars.mb_k_done[raw_index.idx].wait(raw_index.phase) if elect_one: bars.mb_k_ready[raw_index.idx].arrive(n_bytes=k_tx_bytes) k_slice = tma_slice_runtime_desc(desc_k_slot, cutlass.Int32(0), head_k, chunk_start) tma_load_tile(sK_tma[raw_index.idx], k_slice, bars.mb_k_ready[raw_index.idx].smem_ptr, acquire=False) - # ---- Gate load ------------------------------------------------------- + # ---- Gate load: GMEM -> SMEM --------------------------------------------- bars.mb_gate_done[raw_index.idx].wait(raw_index.phase) if elect_one: bars.mb_gate_ready[raw_index.idx].arrive(n_bytes=gate_tx_bytes) gate_slice = tma_slice_runtime_desc(desc_gate_slot, cutlass.Int32(0), head_o, chunk_start) tma_load_tile(sGate_tma[raw_index.idx], gate_slice, bars.mb_gate_ready[raw_index.idx].smem_ptr, acquire=False) - # ---- Beta load ------------------------------------------------------- + # ---- Beta load: GMEM -> SMEM --------------------------------------------- bars.mb_beta_done[raw_index.idx].wait(raw_index.phase) if elect_one: bars.mb_beta_ready[raw_index.idx].arrive(n_bytes=beta_tx_bytes) beta_slice = tma_slice_runtime_desc(desc_beta_slot, cutlass.Int32(0), head_o, chunk_start) tma_load_tile(sBeta_tma[raw_index.idx], beta_slice, bars.mb_beta_ready[raw_index.idx].smem_ptr, acquire=False) - # ---- entering state -- + # ---- entering state ------------------------------------------------------ if chunk_idx >= FIRST_STATE_CHUNK: state_idx = state_index.idx bars.mb_state_cg0_done[state_idx].wait(state_index.phase) @@ -1779,7 +1717,7 @@ def tmaldg_warp( state_slice = tma_slice_runtime_desc(desc_checkpoint_slot, cutlass.Int32(0), cutlass.Int32(0), chunk_idx, head_o) tma_load_tile(sState_tma[state_idx], state_slice, bars.mb_state_ready[state_idx].smem_ptr, acquire=False) - # ---- dO load --------------------------------------------------------- + # ---- dO load ------------------------------------------------------------- bars.mb_do_done[raw_index.idx].wait(raw_index.phase) bars.mb_do_mma_done[raw_index.idx].wait(raw_index.phase) if elect_one: @@ -1787,14 +1725,14 @@ def tmaldg_warp( do_slice = tma_slice_runtime_desc(desc_do_slot, cutlass.Int32(0), head_o, chunk_start) tma_load_tile(sDo_tma[raw_index.idx], do_slice, bars.mb_do_ready[raw_index.idx].smem_ptr, acquire=False) - # ---- V load ---------------------------------------------------------- + # ---- V load -------------------------------------------------------------- bars.mb_v_done[raw_index.idx].wait(raw_index.phase) if elect_one: bars.mb_v_ready[raw_index.idx].arrive(n_bytes=v_tx_bytes) v_slice = tma_slice_runtime_desc(desc_v_slot, cutlass.Int32(0), head_v, chunk_start) tma_load_tile(sV_tma[raw_index.idx], v_slice, bars.mb_v_ready[raw_index.idx].smem_ptr, acquire=False) - # ---- W load ---------------------------------------------------------- + # ---- W load -------------------------------------------------------------- bars.mb_w_done[raw_index.idx].wait(raw_index.phase) if elect_one: bars.mb_w_ready[raw_index.idx].arrive(n_bytes=w_tx_bytes) @@ -1809,9 +1747,7 @@ def gate_scale(cfg, raw_gate: cutlass.Float32) -> cutlass.Float32: """Map raw gate to the log2-domain decay increment.""" if cutlass.const_expr(cfg.safe_gate): - half = cutlass.Float32(0.5) - sigmoid = cute.math.tanh(raw_gate * half, approx=True) * half + half - return cfg.gate_scale_log2 * sigmoid + return cfg.gate_scale_log2 * sigmoid(raw_gate) # Default ABI: Gate arrives in natural-log space return raw_gate * cutlass.Float32(LOG2_E) @@ -1824,8 +1760,8 @@ def compute0_warp_group( num_ctas, cu_seqlens, mWorkItems, - sSched, - lane, + sScheduler, + lane_idx, tmem_hold, warp_idx, scale, @@ -1853,31 +1789,34 @@ def compute0_warp_group( stashes the per-row Q/K inverse norms for WG2's dGate assembly and copies H -> TMEM f16 at the chunk tail.""" nvvm.setmaxregister(cfg.num_regs_compute_group_0, nvvm.SetMaxRegisterAction.INCREASE) + elect_one = nvvm.elect_sync() cg0_warp = warp_idx - cfg.compute_group_0_warp_ids[0] nvvm.barrier_cta_sync(cfg.tmem_lifecycle_barrier_id, thread_count=cfg.tmem_user_threads) tmem_base = tmem_hold.load() tmem_col = tmem_base & 0xFFFF tmem_row = tmem_base >> 16 tmem_sp = warp_idx % (cfg.d_v // cfg.threads_per_warp) - value_dim = tmem_sp * cfg.threads_per_warp + lane - state_copy_addr = (tmem_row + tmem_sp * cfg.threads_per_warp) << 16 + value_dim = tmem_sp * cfg.threads_per_warp + lane_idx + state_copy_addr = tmem_row << 16 state_index = PipelineState.start(phase=0) - cg0_prefix_dim = cg0_warp * cfg.threads_per_warp + lane + cg0_channel_dim = cg0_warp * cfg.threads_per_warp + lane_idx cg0_a_log_exp = cutlass.Float32(1.0) cg0_dt_bias_value = cutlass.Float32(0.0) gbase = cutlass.Int32(0) - sched_state = PipelineState.start(phase=0) + scheduler_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) FIRST_STATE_CHUNK = 0 if cfg.use_initial_state else 1 while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) - sk_nt = cend - wstart + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) + sk_nt = compute_end - write_start if cutlass.const_expr(cfg.safe_gate): if sk_nt > 0: cg0_a_log_exp = cute.math.exp2(mA_log[head_idx].to(cutlass.Float32) * LOG2_E, fastmath=True) - cg0_dt_bias_value = mDt_bias[head_idx, cg0_prefix_dim].to(cutlass.Float32) + cg0_dt_bias_value = mDt_bias[head_idx, cg0_channel_dim].to(cutlass.Float32) for rev_idx in cutlass.range(sk_nt, unroll=1): - chunk_idx = cend - cutlass.Int32(1) - rev_idx + chunk_idx = compute_end - cutlass.Int32(1) - rev_idx gc = gbase + rev_idx chunk_start = chunk_idx * cfg.b_t decay_stage = gc % cfg.smem_decay_stages @@ -1899,45 +1838,37 @@ def compute0_warp_group( bars.mb_beta_ready[raw_stage].wait((gc // cfg.smem_raw_stages) % 2) row_group_start = cg0_warp * (cfg.b_t // len(cfg.compute_group_0_warp_ids)) - lane_row_group = lane // 8 - lane_in_row_group = lane - lane_row_group * 8 + lane_row_group = lane_idx // 8 + lane_in_row_group = lane_idx - lane_row_group * 8 decay_row = row_group_start + lane_row_group g_prefix_ptr = sGate_ptr - prefix_dim = cg0_warp * cfg.threads_per_warp + lane - # ---- gate prefix scan: cumulative log-gate per key channel ----------- + channel_dim = cg0_warp * cfg.threads_per_warp + lane_idx + # ---- gate prefix scan: cumulative log-gate per key channel --------------- gate_raw = cutlass.Array(cutlass.Float32, cfg.b_t, alignment=16) for row in cutlass.range_constexpr(cfg.b_t): - f32_segment = prefix_dim // 32 - f32_segment_dim = prefix_dim - f32_segment * 32 + f32_segment = channel_dim // 32 + f32_segment_dim = channel_dim - f32_segment * 32 prefix_idx = f32_segment * (cfg.b_t * 32) + row * 32 + swizzle_xor_128b(row, f32_segment_dim, elem_bytes=4) gate_raw[row] = (sGate_ptr + prefix_idx).load() g_prefix_regs = cutlass.Array(cutlass.Float32, cfg.b_t, alignment=16) if cutlass.const_expr(cfg.safe_gate): - valid_rows = seqlen_b - chunk_idx * cutlass.Int32(cfg.b_t) - valid_mask = cutlass.vector.create_mask([cfg.b_t], [valid_rows]) - for row_pair in cutlass.range_constexpr(cfg.b_t // 2): - row0 = row_pair * 2 - row1 = row0 + 1 - gate0 = cg0_a_log_exp * (gate_raw[row0] + cg0_dt_bias_value) - gate1 = cg0_a_log_exp * (gate_raw[row1] + cg0_dt_bias_value) - gate0 = gate_scale( - cfg, - gate0, - ) - gate1 = gate_scale( - cfg, - gate1, - ) - gate_pair = cutlass.Vector.from_elements((gate0, gate1), cutlass.Float32) - gate_pair = cutlass.vector.where(valid_mask[row0 : row1 + 1], gate_pair, 0.0) - g_prefix_regs[row0] = gate_pair[0] - g_prefix_regs[row1] = gate_pair[1] + for row in cutlass.range_constexpr(cfg.b_t): + gate = gate_raw[row] + token_idx = chunk_idx * cutlass.Int32(cfg.b_t) + cutlass.Int32(row) + if token_idx < batch_seqlen: + gate = gate_scale( + cfg, + cg0_a_log_exp * (gate + cg0_dt_bias_value), + ) + else: + gate = cutlass.Float32(0.0) + g_prefix_regs[row] = gate else: for row in cutlass.range_constexpr(cfg.b_t): gate = gate_raw[row] token_idx = chunk_idx * cutlass.Int32(cfg.b_t) + cutlass.Int32(row) - if token_idx < seqlen_b: + if token_idx < batch_seqlen: gate = gate_scale( cfg, gate, @@ -1968,28 +1899,27 @@ def compute0_warp_group( g_prefix_regs[row] = cute.math.exp2(g_prefix_regs[row], fastmath=True) exp_g_last = g_prefix_regs[cfg.b_t - 1] - # ---- decay-slot guard: previous use fully consumed ------------------- + # ---- decay-slot guard: previous use fully consumed ----------------------- operand_done_phase = ((gc // cfg.smem_decay_stages) + 1) % 2 bars.mb_decay_done[decay_stage].wait(operand_done_phase) for row in cutlass.range_constexpr(cfg.b_t): - f32_segment = prefix_dim // 32 - f32_segment_dim = prefix_dim - f32_segment * 32 + f32_segment = channel_dim // 32 + f32_segment_dim = channel_dim - f32_segment * 32 prefix_idx = f32_segment * (cfg.b_t * 32) + row * 32 + swizzle_xor_128b(row, f32_segment_dim, elem_bytes=4) (sGate_ptr + prefix_idx).store(g_prefix_regs[row]) - # ---- state-scale diag: stage exp2(g_last) decay blocks --------------- - block = prefix_dim // cutlass.Int32(16) - coord = prefix_dim - block * cutlass.Int32(16) - linear_idx = block * cutlass.Int32(256) + coord * cutlass.Int32(16) + coord - diag_idx = swizzle_lin_S(linear_idx, bbits=1, mbase=3, sshift=3) + # ---- state-scale diag: stage exp2(g last) decay blocks ------------------- + block = channel_dim // cutlass.Int32(16) + coord = channel_dim - block * cutlass.Int32(16) + diag_idx = block * cutlass.Int32(256) + coord * cutlass.Int32(16) + swizzle_xor_32b(channel_dim, coord) sState_scale_diag_ptr[diag_idx] = exp_g_last.to(cfg.io_dtype) - # ---- raw Q/K: SMEM -> TMEM ring (channel-major, for WG2) ------------- + # ---- raw Q/K: SMEM -> TMEM ring (channel-major, for WG2) ----------------- qk_raw_stage = gc % cfg.tmem_qk_raw_stages bars.mb_qk_raw_done[qk_raw_stage].wait(((gc // cfg.tmem_qk_raw_stages) + 1) % 2) - raw_seg = prefix_dim // 64 - raw_dim = prefix_dim - raw_seg * 64 + raw_seg = channel_dim // 64 + raw_dim = channel_dim - raw_seg * 64 q_raw_words = cutlass.Array(cutlass.Int32, cfg.b_t // 2, alignment=16) k_raw_words = cutlass.Array(cutlass.Int32, cfg.b_t // 2, alignment=16) for t2 in cutlass.range_constexpr(cfg.b_t // 2): @@ -2004,12 +1934,12 @@ def compute0_warp_group( k_raw_words[t2] = fp32_to_fp16(k0, k1, dtype=cfg.io_dtype) nvvm.tcgen05_st( "32x32b", - nvvm.make_tmem_ptr(state_copy_addr + (tmem_col + cfg.tmem_qraw_inp_offset + qk_raw_stage * (cfg.b_t // 2)), cutlass.Int8), + nvvm.make_tmem_ptr(state_copy_addr + (tmem_col + cfg.tmem_qraw_input_offset + qk_raw_stage * (cfg.b_t // 2)), cutlass.Int8), q_raw_words[0 : (cfg.b_t // 2)], ) nvvm.tcgen05_st( "32x32b", - nvvm.make_tmem_ptr(state_copy_addr + (tmem_col + cfg.tmem_kraw_inp_offset + qk_raw_stage * (cfg.b_t // 2)), cutlass.Int8), + nvvm.make_tmem_ptr(state_copy_addr + (tmem_col + cfg.tmem_kraw_input_offset + qk_raw_stage * (cfg.b_t // 2)), cutlass.Int8), k_raw_words[0 : (cfg.b_t // 2)], ) nvvm.tcgen05_wait("store") @@ -2021,7 +1951,7 @@ def compute0_warp_group( raw_q_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) raw_k_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) raw_beta_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) - # ---- optional Q/K L2-norm -------------------------------------------- + # ---- optional Q/K L2-norm ------------------------------------------------ if cutlass.const_expr(cfg.l2norm): qk0_lo = opaque_f32_zero() qk0_hi = opaque_f32_zero() @@ -2046,8 +1976,7 @@ def compute0_warp_group( raw_k_regs[reg_base + dim_offset] = k_val beta_val = raw_beta_frag_f32[dim_offset] if cutlass.const_expr(cfg.beta_sigmoid): - half = cutlass.Float32(0.5) - beta_val = (cute.math.tanh(beta_val * half, approx=True) * half + half).to(cfg.io_dtype).to(cutlass.Float32) + beta_val = sigmoid(beta_val).to(cfg.io_dtype).to(cutlass.Float32) raw_beta_regs[reg_base + dim_offset] = beta_val if cutlass.const_expr(cfg.l2norm): if cutlass.const_expr(dim_offset % 2 == 0): @@ -2079,7 +2008,7 @@ def compute0_warp_group( sNorm_raw[(gc % cfg.tmem_qk_raw_stages) * (2 * cfg.b_t) + cfg.b_t + decay_row] = k_inv_norm q_stage_norm = q_inv_norm * scale - # ---- decay/restore operands: exp2(+-g) applied per key channel ------- + # ---- decay/restore operands: exp2(+-g) applied per key channel ----------- exp_g_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) exp_g_last_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) for dim_half in cutlass.range_constexpr(2): @@ -2108,7 +2037,7 @@ def compute0_warp_group( for dim_half in cutlass.range_constexpr(2): dim_base = dim_half * (cfg.d_k // 2) + lane_in_row_group * 8 reg_base = dim_half * 8 - # ---- K decay + K inv operands: exp2(g) * Beta * K / exp2(-g) * K - + # ---- K decay + K inv operands: exp2(g) * Beta * K / exp2(-g) * K ----- k_decay_pack = cutlass.Array(cutlass.Int32, 4, alignment=16) for pair_idx in cutlass.range_constexpr(4): dim0 = pair_idx * 2 @@ -2147,7 +2076,7 @@ def compute0_warp_group( nvvm.fence_proxy("async.shared", space="cta") bars.mb_k_decay_inv_ready[decay_stage].arrive() - # ---- Q decay + K_restore operands ------------------------------------ + # ---- Q decay + K restore operands ---------------------------------------- for dim_half in cutlass.range_constexpr(2): dim_base = dim_half * (cfg.d_k // 2) + lane_in_row_group * 8 reg_base = dim_half * 8 @@ -2181,39 +2110,38 @@ def compute0_warp_group( nvvm.fence_proxy("async.shared", space="cta") bars.mb_q_decay_k_restore_ready[decay_stage].arrive() - # ---- state copy: SMEM -> TMEM f16 ------------------------------------ - bars.mb_state_inp_done[gc % 2].wait(((gc // 2) + 1) % 2) - bars.mb_state_inp_cg2_done[gc % 2].wait(((gc // 2) + 1) % 2) + # ---- state copy: SMEM -> TMEM f16 ---------------------------------------- + bars.mb_state_input_done[gc % 2].wait(((gc // 2) + 1) % 2) + bars.mb_state_input_cg2_done[gc % 2].wait(((gc // 2) + 1) % 2) if chunk_idx >= FIRST_STATE_CHUNK: bars.mb_state_ready[state_index.idx].wait(state_index.phase) state_src = sState_raw.data_ptr() + state_index.idx * (cfg.d_k * cfg.d_v) - for pl in cutlass.range_constexpr(2): - for g8 in cutlass.range_constexpr(8): - state_frag = cutlass.Vector.from_elements( - tuple( - ( - state_src - + (value_dim // 64) * (cfg.d_v * 64) - + (pl * 64 + g8 * 8 + e) * 64 - + swizzle_xor_128b(pl * 64 + g8 * 8 + e, value_dim % 64, elem_bytes=2) - ).load() - for e in range(8) - ), - cfg.io_dtype, - ) - nvvm.tcgen05_st( - "32x32b", - nvvm.make_tmem_ptr( - state_copy_addr + (tmem_col + cfg.tmem_state_inp_offset + (gc % 2) * (cfg.d_v // 2) + pl * 32 + g8 * 4), cutlass.Int8 - ), - state_frag.bitcast(cutlass.Int32), - ) + ldm_row_coord = (lane_idx // 16) * 8 + (lane_idx & 7) + ldm_col_offset = ((lane_idx // 8) & 1) * 8 + k_base = tmem_sp * cfg.threads_per_warp + k_seg_off = (k_base // 64) * (cfg.d_v * 64) + state_copy_hi_addr = state_copy_addr + (16 << 16) + state_col = tmem_col + cfg.tmem_state_input_offset + (gc % 2) * (cfg.d_v // 2) + for dv_blk in cutlass.range_constexpr(cfg.d_v // 16): + dv_row = dv_blk * 16 + ldm_row_coord + frag_lo = nvvm.ldmatrix( + state_src + k_seg_off + dv_row * 64 + swizzle_xor_128b(dv_row, (k_base + ldm_col_offset) % 64, elem_bytes=2), + 4, + nvvm.MMALayout.COL, + ) + frag_hi = nvvm.ldmatrix( + state_src + k_seg_off + dv_row * 64 + swizzle_xor_128b(dv_row, (k_base + 16 + ldm_col_offset) % 64, elem_bytes=2), + 4, + nvvm.MMALayout.COL, + ) + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(state_copy_addr + (state_col + dv_blk * 8), cutlass.Int8), frag_lo) + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(state_copy_hi_addr + (state_col + dv_blk * 8), cutlass.Int8), frag_hi) nvvm.tcgen05_wait("store") bars.mb_state_cg0_done[state_index.idx].arrive() state_index = advance(state_index, cfg.smem_state_stages) - bars.mb_state_inp_ready[gc % 2].arrive() + bars.mb_state_input_ready[gc % 2].arrive() gbase += sk_nt - tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one) @cute.jit @@ -2224,8 +2152,8 @@ def compute1_warp_group( num_ctas, cu_seqlens, mWorkItems, - sSched, - lane, + sScheduler, + lane_idx, tmem_hold, warp_idx, mDstate0, @@ -2242,16 +2170,17 @@ def compute1_warp_group( ) -> None: """WG1 warp role (warps 4-7): the value-side TMEM staging.""" nvvm.setmaxregister(cfg.num_regs_compute_group_1, nvvm.SetMaxRegisterAction.INCREASE) + elect_one = nvvm.elect_sync() nvvm.barrier_cta_sync(cfg.tmem_lifecycle_barrier_id, thread_count=cfg.tmem_user_threads) tmem_base = tmem_hold.load() tmem_col = tmem_base & 0xFFFF tmem_row = tmem_base >> 16 tmem_sp = warp_idx % (cfg.d_v // cfg.threads_per_warp) - ov_tok = (lane // 16) * 8 + (lane & 7) - ov_col = ((lane // 8) & 1) * 8 - value_dim = tmem_sp * cfg.threads_per_warp + lane + ov_row_coord = (lane_idx // 16) * 8 + (lane_idx & 7) + ov_col_offset = ((lane_idx // 8) & 1) * 8 + value_dim = tmem_sp * cfg.threads_per_warp + lane_idx value_dim_base = tmem_sp * cfg.threads_per_warp - cg1_tidx = warp_idx % 4 * cfg.threads_per_warp + lane + cg1_tidx = warp_idx % 4 * cfg.threads_per_warp + lane_idx raw_index = PipelineState.start(phase=0) state_k_index = PipelineState.start(phase=0) @@ -2265,32 +2194,34 @@ def compute1_warp_group( dwo_done_index = PipelineState.start(phase=1) gbase = cutlass.Int32(0) - sched_state = PipelineState.start(phase=0) + scheduler_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) FIRST_STATE_CHUNK = 0 if cfg.use_initial_state else 1 while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) - sk_nt = cend - wstart + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) + sk_nt = compute_end - write_start - # ---- dht seeding: dH acc + dh_inp f16 + sdH ------------------------------ + # ---- dstate seed: GMEM -> TMEM + SMEM ---------------------------------------- if cutlass.const_expr(cfg.use_dstate_in): if sk_nt > 0: - seed_true = cend == num_chunks_b + seed_true = compute_end == batch_num_chunks bars.mb_dstate_smem_done.wait(dstate_smem_done_index.phase) bars.mb_dstate_smem_cg2_done.wait(dstate_smem_done_index.phase) dstate_smem_done_index = advance(dstate_smem_done_index, 1) - row_addr = (tmem_row + tmem_sp * cfg.threads_per_warp) << 16 + row_lo_addr = tmem_row << 16 dstate_src = (mDstate_in.iterator + mDstate_in.layout((batch_idx, head_idx, value_dim, 0))).raw_ptr() - for sub in cutlass.range_constexpr(cfg.d_k // 16): + for i in cutlass.range_constexpr(cfg.d_k // 16): seed_block = cutlass.Array(cutlass.Float32, 16, alignment=16) for g in cutlass.range_constexpr(4): - seed_chunk = (dstate_src + sub * 16 + g * 4).load(count=4, alignment=16) + seed_chunk = (dstate_src + i * 16 + g * 4).load(count=4, alignment=16) for t in cutlass.range_constexpr(4): dval = seed_chunk[t].to(cutlass.Float32) seed_block[g * 4 + t] = dval if seed_true else cutlass.Float32(0.0) nvvm.tcgen05_st( "32x32b", - nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_dstate_acc_offset + sub * 16), cutlass.Float32), + nvvm.make_tmem_ptr(row_lo_addr + (tmem_col + cfg.tmem_dstate_acc_offset + i * 16), cutlass.Float32), seed_block[0:16], ) seed_pack = cutlass.Array(cutlass.Int32, 8, alignment=16) @@ -2298,19 +2229,19 @@ def compute1_warp_group( seed_pack[pc] = fp32_to_fp16(seed_block[2 * pc], seed_block[2 * pc + 1], dtype=cfg.io_dtype) nvvm.tcgen05_st( "32x32b", - nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_dstate_inp_offset + sub * 8), cutlass.Int8), + nvvm.make_tmem_ptr(row_lo_addr + (tmem_col + cfg.tmem_dstate_input_offset + i * 8), cutlass.Int8), seed_pack[0:8], ) nvvm.tcgen05_wait("store") - bars.mb_dstate_inp_ready.arrive() + bars.mb_dstate_input_ready.arrive() - # ---- dht seed -> sdH: re-read dh_inp after the TMEM publish ------ - for sub in cutlass.range_constexpr(cfg.d_k // 16): + # ---- dstate seed -> sdH: re-read after the TMEM publish -------------- + for i in cutlass.range_constexpr(cfg.d_k // 16): dstate_words = nvvm.tcgen05_ld( - "32x32b", nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_dstate_inp_offset + sub * 8), cutlass.Float32), num=8 + "32x32b", nvvm.make_tmem_ptr(row_lo_addr + (tmem_col + cfg.tmem_dstate_input_offset + i * 8), cutlass.Float32), num=8 ) for half in cutlass.range_constexpr(2): - d_base = sub * 16 + half * 8 + d_base = i * 16 + half * 8 h_pack = cutlass.Vector.from_elements( (dstate_words[half * 4], dstate_words[half * 4 + 1], dstate_words[half * 4 + 2], dstate_words[half * 4 + 3]), cutlass.Float32, @@ -2321,7 +2252,7 @@ def compute1_warp_group( bars.mb_dstate_smem_ready.arrive() for rev_idx in cutlass.range(sk_nt, unroll=1): - chunk_idx = cend - cutlass.Int32(1) - rev_idx + chunk_idx = compute_end - cutlass.Int32(1) - rev_idx gc = gbase + rev_idx has_dstate = cutlass.Boolean(rev_idx > 0) if cutlass.const_expr(cfg.use_dstate_in): @@ -2329,175 +2260,173 @@ def compute1_warp_group( sV_ptr = sV_raw.data_ptr() + raw_index.idx * (cfg.d_v * cfg.b_t) sW_ptr = sW_raw.data_ptr() + raw_index.idx * (cfg.d_v * cfg.b_t) - row_addr_lo = tmem_row << 16 - row_addr_hi = (tmem_row + 16) << 16 - row_id0 = tmem_row + value_dim_base - row_id1 = row_id0 + 16 + row_lo_addr = tmem_row << 16 + row_hi_addr = (tmem_row + 16) << 16 - # ---- Y staging: Y = W*V - state_k -> TMEM f16 ------------------------ + # ---- Y stage: Y = W*V - k state -> TMEM f16 ------------------------------ bars.mb_v_ready[raw_index.idx].wait((gc // cfg.smem_raw_stages) % 2) bars.mb_w_ready[raw_index.idx].wait((gc // cfg.smem_raw_stages) % 2) projection_col_id = tmem_col + cfg.tmem_state_k_acc_offset - input_col_id = tmem_col + cfg.tmem_y_inp_offset - raw_v_frag0 = nvvm.ldmatrix( + input_col_id = tmem_col + cfg.tmem_y_input_offset + raw_v_frag_lo = nvvm.ldmatrix( sV_ptr - + (value_dim_base + ov_col) // 64 * (cfg.b_t * 64) - + ov_tok * 64 - + swizzle_xor_128b(ov_tok, (value_dim_base + ov_col) % 64, elem_bytes=2), + + (value_dim_base + ov_col_offset) // 64 * (cfg.b_t * 64) + + ov_row_coord * 64 + + swizzle_xor_128b(ov_row_coord, (value_dim_base + ov_col_offset) % 64, elem_bytes=2), 4, nvvm.MMALayout.COL, ) - raw_v_frag1 = nvvm.ldmatrix( + raw_v_frag_hi = nvvm.ldmatrix( sV_ptr - + (value_dim_base + 16 + ov_col) // 64 * (cfg.b_t * 64) - + ov_tok * 64 - + swizzle_xor_128b(ov_tok, (value_dim_base + 16 + ov_col) % 64, elem_bytes=2), + + (value_dim_base + 16 + ov_col_offset) // 64 * (cfg.b_t * 64) + + ov_row_coord * 64 + + swizzle_xor_128b(ov_row_coord, (value_dim_base + 16 + ov_col_offset) % 64, elem_bytes=2), 4, nvvm.MMALayout.COL, ) - raw_w_frag0 = nvvm.ldmatrix( + raw_w_frag_lo = nvvm.ldmatrix( sW_ptr - + (value_dim_base + ov_col) // 64 * (cfg.b_t * 64) - + ov_tok * 64 - + swizzle_xor_128b(ov_tok, (value_dim_base + ov_col) % 64, elem_bytes=2), + + (value_dim_base + ov_col_offset) // 64 * (cfg.b_t * 64) + + ov_row_coord * 64 + + swizzle_xor_128b(ov_row_coord, (value_dim_base + ov_col_offset) % 64, elem_bytes=2), 4, nvvm.MMALayout.COL, ) - raw_w_frag1 = nvvm.ldmatrix( + raw_w_frag_hi = nvvm.ldmatrix( sW_ptr - + (value_dim_base + 16 + ov_col) // 64 * (cfg.b_t * 64) - + ov_tok * 64 - + swizzle_xor_128b(ov_tok, (value_dim_base + 16 + ov_col) % 64, elem_bytes=2), + + (value_dim_base + 16 + ov_col_offset) // 64 * (cfg.b_t * 64) + + ov_row_coord * 64 + + swizzle_xor_128b(ov_row_coord, (value_dim_base + 16 + ov_col_offset) % 64, elem_bytes=2), 4, nvvm.MMALayout.COL, ) - y_inp_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) - y_inp_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + y_input_pack_lo = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + y_input_pack_hi = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) if chunk_idx >= FIRST_STATE_CHUNK: bars.mb_state_k_acc_ready.wait(state_k_index.phase) state_k_index = advance(state_k_index, 1) - state_k_vec0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id0 << 16) + projection_col_id, cutlass.Float32), num=2) - state_k_vec1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id1 << 16) + projection_col_id, cutlass.Float32), num=2) + state_k_vec_lo = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_lo_addr + projection_col_id, cutlass.Float32), num=2) + state_k_vec_hi = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_hi_addr + projection_col_id, cutlass.Float32), num=2) for reg_idx in cutlass.range_constexpr(4): - raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1) - frag_pair = (reg_idx ^ 2) * 2 - state_k_pair = fp32_to_fp16(state_k_vec0[frag_pair], state_k_vec0[frag_pair + 1], dtype=cfg.io_dtype) - wv_pair = mul_f16x2(raw_w_frag0[raw_matrix], raw_v_frag0[raw_matrix], cfg.io_dtype) - y_inp_pack0[reg_idx ^ 2] = sub_f16x2(wv_pair, state_k_pair, cfg.io_dtype) + frag_pair = reg_idx * 2 + state_k_pair = fp32_to_fp16(state_k_vec_lo[frag_pair], state_k_vec_lo[frag_pair + 1], dtype=cfg.io_dtype) + wv_pair = mul_f16x2(raw_w_frag_lo[reg_idx], raw_v_frag_lo[reg_idx], cfg.io_dtype) + y_input_pack_lo[reg_idx] = sub_f16x2(wv_pair, state_k_pair, cfg.io_dtype) for reg_idx in cutlass.range_constexpr(4): - raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1) - frag_pair = (reg_idx ^ 2) * 2 - state_k_pair = fp32_to_fp16(state_k_vec1[frag_pair], state_k_vec1[frag_pair + 1], dtype=cfg.io_dtype) - wv_pair = mul_f16x2(raw_w_frag1[raw_matrix], raw_v_frag1[raw_matrix], cfg.io_dtype) - y_inp_pack1[reg_idx ^ 2] = sub_f16x2(wv_pair, state_k_pair, cfg.io_dtype) + frag_pair = reg_idx * 2 + state_k_pair = fp32_to_fp16(state_k_vec_hi[frag_pair], state_k_vec_hi[frag_pair + 1], dtype=cfg.io_dtype) + wv_pair = mul_f16x2(raw_w_frag_hi[reg_idx], raw_v_frag_hi[reg_idx], cfg.io_dtype) + y_input_pack_hi[reg_idx] = sub_f16x2(wv_pair, state_k_pair, cfg.io_dtype) if chunk_idx < FIRST_STATE_CHUNK: for reg_idx in cutlass.range_constexpr(4): - raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1) - y_inp_pack0[reg_idx ^ 2] = mul_f16x2(raw_w_frag0[raw_matrix], raw_v_frag0[raw_matrix], cfg.io_dtype) + y_input_pack_lo[reg_idx] = mul_f16x2(raw_w_frag_lo[reg_idx], raw_v_frag_lo[reg_idx], cfg.io_dtype) for reg_idx in cutlass.range_constexpr(4): - raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1) - y_inp_pack1[reg_idx ^ 2] = mul_f16x2(raw_w_frag1[raw_matrix], raw_v_frag1[raw_matrix], cfg.io_dtype) - nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_addr_lo + input_col_id, cutlass.Int8), y_inp_pack0[0:4]) - nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_addr_hi + input_col_id, cutlass.Int8), y_inp_pack1[0:4]) + y_input_pack_hi[reg_idx] = mul_f16x2(raw_w_frag_hi[reg_idx], raw_v_frag_hi[reg_idx], cfg.io_dtype) + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_lo_addr + input_col_id, cutlass.Int8), y_input_pack_lo[0:4]) + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_hi_addr + input_col_id, cutlass.Int8), y_input_pack_hi[0:4]) nvvm.tcgen05_wait("store") - bars.mb_y_inp_ready.arrive() + bars.mb_y_input_ready.arrive() - # ---- dU restage: dU acc -> TMEM f16 A operand ------------------------ + # ---- dU stage: dU acc -> TMEM f16 A operand ------------------------------ bars.mb_du_acc_ready.wait(du_acc_index.phase) du_acc_index = advance(du_acc_index, 1) du_col_id = tmem_col + cfg.tmem_du_acc_offset - du_vec0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id0 << 16) + du_col_id, cutlass.Float32), num=2) - du_vec1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id1 << 16) + du_col_id, cutlass.Float32), num=2) + du_vec_lo = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_lo_addr + du_col_id, cutlass.Float32), num=2) + du_vec_hi = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_hi_addr + du_col_id, cutlass.Float32), num=2) - du_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) - du_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + du_pack_lo = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + du_pack_hi = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) for reg_idx in cutlass.range_constexpr(4): frag_pair = reg_idx * 2 - du_pack0[reg_idx] = fp32_to_fp16(du_vec0[frag_pair], du_vec0[frag_pair + 1], dtype=cfg.io_dtype) - du_pack1[reg_idx] = fp32_to_fp16(du_vec1[frag_pair], du_vec1[frag_pair + 1], dtype=cfg.io_dtype) - nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_addr_lo + (tmem_col + cfg.tmem_du_inp_offset), cutlass.Int8), du_pack0[0:4]) - nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_addr_hi + (tmem_col + cfg.tmem_du_inp_offset), cutlass.Int8), du_pack1[0:4]) + du_pack_lo[reg_idx] = fp32_to_fp16(du_vec_lo[frag_pair], du_vec_lo[frag_pair + 1], dtype=cfg.io_dtype) + du_pack_hi[reg_idx] = fp32_to_fp16(du_vec_hi[frag_pair], du_vec_hi[frag_pair + 1], dtype=cfg.io_dtype) + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_lo_addr + (tmem_col + cfg.tmem_du_input_offset), cutlass.Int8), du_pack_lo[0:4]) + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_hi_addr + (tmem_col + cfg.tmem_du_input_offset), cutlass.Int8), du_pack_hi[0:4]) nvvm.tcgen05_wait("store") - bars.mb_du_inp_ready.arrive() + bars.mb_du_input_ready.arrive() - # ---- U readback -> sU ------------------------------------------------ + # ---- U read: TMEM -> sU -------------------------------------------------- bars.mb_u_acc_ready.wait(u_acc_index.phase) u_acc_index = advance(u_acc_index, 1) u_col_id = tmem_col + cfg.tmem_u_acc_offset - u_vec0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id0 << 16) + u_col_id, cutlass.Float32), num=2) - u_vec1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id1 << 16) + u_col_id, cutlass.Float32), num=2) + u_vec_lo = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_lo_addr + u_col_id, cutlass.Float32), num=2) + u_vec_hi = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_hi_addr + u_col_id, cutlass.Float32), num=2) - u_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) - u_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + u_pack_lo = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + u_pack_hi = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) for reg_idx in cutlass.range_constexpr(4): - u_pack0[reg_idx] = fp32_to_fp16(u_vec0[2 * reg_idx], u_vec0[2 * reg_idx + 1], dtype=cfg.io_dtype) - u_pack1[reg_idx] = fp32_to_fp16(u_vec1[2 * reg_idx], u_vec1[2 * reg_idx + 1], dtype=cfg.io_dtype) + u_pack_lo[reg_idx] = fp32_to_fp16(u_vec_lo[2 * reg_idx], u_vec_lo[2 * reg_idx + 1], dtype=cfg.io_dtype) + u_pack_hi[reg_idx] = fp32_to_fp16(u_vec_hi[2 * reg_idx], u_vec_hi[2 * reg_idx + 1], dtype=cfg.io_dtype) nvvm.stmatrix( sU_raw.data_ptr() - + (value_dim_base + ov_col) // 64 * (cfg.b_t * 64) - + ov_tok * 64 - + swizzle_xor_128b(ov_tok, (value_dim_base + ov_col) % 64, elem_bytes=2), - u_pack0.data_ptr().load(count=4, alignment=4), + + (value_dim_base + ov_col_offset) // 64 * (cfg.b_t * 64) + + ov_row_coord * 64 + + swizzle_xor_128b(ov_row_coord, (value_dim_base + ov_col_offset) % 64, elem_bytes=2), + u_pack_lo.data_ptr().load(count=4, alignment=4), nvvm.MMALayout.COL, shape=nvvm.StoreShape.M8N8, ) nvvm.stmatrix( sU_raw.data_ptr() - + (value_dim_base + 16 + ov_col) // 64 * (cfg.b_t * 64) - + ov_tok * 64 - + swizzle_xor_128b(ov_tok, (value_dim_base + 16 + ov_col) % 64, elem_bytes=2), - u_pack1.data_ptr().load(count=4, alignment=4), + + (value_dim_base + 16 + ov_col_offset) // 64 * (cfg.b_t * 64) + + ov_row_coord * 64 + + swizzle_xor_128b(ov_row_coord, (value_dim_base + 16 + ov_col_offset) % 64, elem_bytes=2), + u_pack_hi.data_ptr().load(count=4, alignment=4), nvvm.MMALayout.COL, shape=nvvm.StoreShape.M8N8, ) nvvm.fence_proxy("async.shared", space="cta") bars.mb_u_smem_ready.arrive() - # ---- dY readback ------------------------------------------------------- + # ---- dY read ------------------------------------------------------------- bars.mb_dy_acc_ready.wait(dy_acc_index.phase) dy_acc_index = advance(dy_acc_index, 1) dy_col_id = tmem_col + cfg.tmem_dy_acc_offset - dy_vec0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id0 << 16) + dy_col_id, cutlass.Float32), num=2) - dy_vec1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id1 << 16) + dy_col_id, cutlass.Float32), num=2) + dy_vec_lo = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_lo_addr + dy_col_id, cutlass.Float32), num=2) + dy_vec_hi = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_hi_addr + dy_col_id, cutlass.Float32), num=2) - # ---- -dY -> TMEM: A operand of the dstate K-term ----------------------- - neg_dy_regs0 = cutlass.Array(cutlass.Float32, 8, alignment=16) - neg_dy_regs1 = cutlass.Array(cutlass.Float32, 8, alignment=16) + # ---- -dY -> TMEM: A operand of the dstate K-term ------------------------- + neg_dy_regs_lo = cutlass.Array(cutlass.Float32, 8, alignment=16) + neg_dy_regs_hi = cutlass.Array(cutlass.Float32, 8, alignment=16) for e in cutlass.range_constexpr(8): - neg_dy_regs0[e] = -dy_vec0[e] - neg_dy_regs1[e] = -dy_vec1[e] - neg_dy_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) - neg_dy_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + neg_dy_regs_lo[e] = -dy_vec_lo[e] + neg_dy_regs_hi[e] = -dy_vec_hi[e] + neg_dy_pack_lo = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + neg_dy_pack_hi = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) for reg_idx in cutlass.range_constexpr(4): frag_pair = reg_idx * 2 - neg_dy_pack0[reg_idx] = fp32_to_fp16(neg_dy_regs0[frag_pair], neg_dy_regs0[frag_pair + 1], dtype=cfg.io_dtype) - neg_dy_pack1[reg_idx] = fp32_to_fp16(neg_dy_regs1[frag_pair], neg_dy_regs1[frag_pair + 1], dtype=cfg.io_dtype) - nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_addr_lo + (tmem_col + cfg.tmem_neg_dy_inp_offset), cutlass.Int8), neg_dy_pack0[0:4]) - nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_addr_hi + (tmem_col + cfg.tmem_neg_dy_inp_offset), cutlass.Int8), neg_dy_pack1[0:4]) + neg_dy_pack_lo[reg_idx] = fp32_to_fp16(neg_dy_regs_lo[frag_pair], neg_dy_regs_lo[frag_pair + 1], dtype=cfg.io_dtype) + neg_dy_pack_hi[reg_idx] = fp32_to_fp16(neg_dy_regs_hi[frag_pair], neg_dy_regs_hi[frag_pair + 1], dtype=cfg.io_dtype) + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_lo_addr + (tmem_col + cfg.tmem_neg_dy_input_offset), cutlass.Int8), neg_dy_pack_lo[0:4]) + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_hi_addr + (tmem_col + cfg.tmem_neg_dy_input_offset), cutlass.Int8), neg_dy_pack_hi[0:4]) nvvm.tcgen05_wait("store") - bars.mb_neg_dy_inp_ready.arrive() - - # ---- dY -> sdY: pack + store + publish (super dM + dV scalar operand) -- - addr_lo0 = (value_dim_base + ov_col) // 64 * (cfg.b_t * 64) + ov_tok * 64 + swizzle_xor_128b(ov_tok, (value_dim_base + ov_col) % 64, elem_bytes=2) - addr_lo1 = ( - (value_dim_base + 16 + ov_col) // 64 * (cfg.b_t * 64) - + ov_tok * 64 - + swizzle_xor_128b(ov_tok, (value_dim_base + 16 + ov_col) % 64, elem_bytes=2) + bars.mb_neg_dy_input_ready.arrive() + + # ---- dY -> sdY: pack + store + publish (super dM + dV scalar operand) ---- + dy_addr_lo = ( + (value_dim_base + ov_col_offset) // 64 * (cfg.b_t * 64) + + ov_row_coord * 64 + + swizzle_xor_128b(ov_row_coord, (value_dim_base + ov_col_offset) % 64, elem_bytes=2) ) - dy_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) - dy_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + dy_addr_hi = ( + (value_dim_base + 16 + ov_col_offset) // 64 * (cfg.b_t * 64) + + ov_row_coord * 64 + + swizzle_xor_128b(ov_row_coord, (value_dim_base + 16 + ov_col_offset) % 64, elem_bytes=2) + ) + dy_pack_lo = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + dy_pack_hi = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) for reg_idx in cutlass.range_constexpr(4): - dy_pack0[reg_idx] = fp32_to_fp16(dy_vec0[2 * reg_idx], dy_vec0[2 * reg_idx + 1], dtype=cfg.io_dtype) - dy_pack1[reg_idx] = fp32_to_fp16(dy_vec1[2 * reg_idx], dy_vec1[2 * reg_idx + 1], dtype=cfg.io_dtype) + dy_pack_lo[reg_idx] = fp32_to_fp16(dy_vec_lo[2 * reg_idx], dy_vec_lo[2 * reg_idx + 1], dtype=cfg.io_dtype) + dy_pack_hi[reg_idx] = fp32_to_fp16(dy_vec_hi[2 * reg_idx], dy_vec_hi[2 * reg_idx + 1], dtype=cfg.io_dtype) bars.mb_dy_smem_done.wait(sdy_done_index.phase) sdy_done_index = advance(sdy_done_index, 1) - nvvm.stmatrix(sDy_raw.data_ptr() + addr_lo0, dy_pack0.data_ptr().load(count=4, alignment=4), nvvm.MMALayout.COL, shape=nvvm.StoreShape.M8N8) - nvvm.stmatrix(sDy_raw.data_ptr() + addr_lo1, dy_pack1.data_ptr().load(count=4, alignment=4), nvvm.MMALayout.COL, shape=nvvm.StoreShape.M8N8) + nvvm.stmatrix(sDy_raw.data_ptr() + dy_addr_lo, dy_pack_lo.data_ptr().load(count=4, alignment=4), nvvm.MMALayout.COL, shape=nvvm.StoreShape.M8N8) + nvvm.stmatrix(sDy_raw.data_ptr() + dy_addr_hi, dy_pack_hi.data_ptr().load(count=4, alignment=4), nvvm.MMALayout.COL, shape=nvvm.StoreShape.M8N8) nvvm.fence_proxy("async.shared", space="cta") bars.mb_dy_smem_ready.arrive() - # ---- scalar pass over own sdY: dV staging ---------------------------- + # ---- scalar pass over own sdY: dV stage ---------------------------------- dv_stage = gc % cfg.smem_dv_stages bars.mb_dv_tmastg_done[dv_stage].wait(dv_done_index.phase) dv_done_index = advance(dv_done_index, cfg.smem_dv_stages) @@ -2522,36 +2451,36 @@ def compute1_warp_group( bars.mb_v_done[raw_index.idx].arrive() bars.mb_w_done[raw_index.idx].arrive() - # ---- dH capture for the next ----------------------------------------- + # ---- dH capture for the next --------------------------------------------- bars.mb_dstate_acc_ready.wait(dstate_ready_index.phase) dstate_ready_index = advance(dstate_ready_index, 1) if rev_idx + cutlass.Int32(1) < sk_nt: bars.mb_dstate_smem_done.wait(dstate_smem_done_index.phase) bars.mb_dstate_smem_cg2_done.wait(dstate_smem_done_index.phase) dstate_smem_done_index = advance(dstate_smem_done_index, 1) - row_addr = (tmem_row + tmem_sp * cfg.threads_per_warp) << 16 - for sub in cutlass.range_constexpr(cfg.d_k // 32): + row_lo_addr = tmem_row << 16 + for i in cutlass.range_constexpr(cfg.d_k // 32): dstate_vec = nvvm.tcgen05_ld( - "32x32b", nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_dstate_acc_offset + sub * 32), cutlass.Float32), num=32 + "32x32b", nvvm.make_tmem_ptr(row_lo_addr + (tmem_col + cfg.tmem_dstate_acc_offset + i * 32), cutlass.Float32), num=32 ) dstate_pack = cutlass.Array(cutlass.Int32, 16, alignment=16) for pc in cutlass.range_constexpr(16): dstate_pack[pc] = fp32_to_fp16(dstate_vec[2 * pc], dstate_vec[2 * pc + 1], dtype=cfg.io_dtype) nvvm.tcgen05_st( "32x32b", - nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_dstate_inp_offset + sub * 16), cutlass.Int8), + nvvm.make_tmem_ptr(row_lo_addr + (tmem_col + cfg.tmem_dstate_input_offset + i * 16), cutlass.Int8), dstate_pack[0:16], ) nvvm.tcgen05_wait("store") - bars.mb_dstate_inp_ready.arrive() + bars.mb_dstate_input_ready.arrive() - # ---- dh_inp -> sdH: re-read after the TMEM publish --------------- - for sub in cutlass.range_constexpr(cfg.d_k // 32): + # ---- dstate input -> sdH: re-read after the TMEM publish ------------- + for i in cutlass.range_constexpr(cfg.d_k // 32): dstate_words = nvvm.tcgen05_ld( - "32x32b", nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_dstate_inp_offset + sub * 16), cutlass.Float32), num=16 + "32x32b", nvvm.make_tmem_ptr(row_lo_addr + (tmem_col + cfg.tmem_dstate_input_offset + i * 16), cutlass.Float32), num=16 ) for half in cutlass.range_constexpr(4): - d_base = sub * 32 + half * 8 + d_base = i * 32 + half * 8 h_pack = cutlass.Vector.from_elements( (dstate_words[half * 4], dstate_words[half * 4 + 1], dstate_words[half * 4 + 2], dstate_words[half * 4 + 3]), cutlass.Float32, @@ -2562,18 +2491,18 @@ def compute1_warp_group( bars.mb_dstate_smem_ready.arrive() raw_index = advance(raw_index, cfg.smem_raw_stages) - # ---- tile end: dS0 drain / zero-length pass-through ---------------------- + # ---- tile end: dstate0 store / zero-length pass-through ---------------------- if cutlass.const_expr(mDstate0 is not None): if sk_nt > 0: - if wstart == 0: - row_addr = (tmem_row + tmem_sp * cfg.threads_per_warp) << 16 + if write_start == 0: + row_lo_addr = tmem_row << 16 dstate0_dst = (mDstate0.iterator + mDstate0.layout((batch_idx, head_idx, value_dim, 0))).raw_ptr() - for sub in cutlass.range_constexpr(cfg.d_k // 32): + for i in cutlass.range_constexpr(cfg.d_k // 32): dstate0_vec = nvvm.tcgen05_ld( - "32x32b", nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_dstate_acc_offset + sub * 32), cutlass.Float32), num=32 + "32x32b", nvvm.make_tmem_ptr(row_lo_addr + (tmem_col + cfg.tmem_dstate_acc_offset + i * 32), cutlass.Float32), num=32 ) for g in cutlass.range_constexpr(8): - (dstate0_dst + sub * 32 + g * 4).store( + (dstate0_dst + i * 32 + g * 4).store( cutlass.Vector.from_elements(tuple(dstate0_vec[g * 4 + t] for t in range(4)), cutlass.Float32), alignment=16, ) @@ -2585,9 +2514,10 @@ def compute1_warp_group( mDstate0[batch_idx, head_idx, value_dim, kd] = mDstate_in[batch_idx, head_idx, value_dim, kd] else: mDstate0[batch_idx, head_idx, value_dim, kd] = cutlass.Float32(0.0) - bars.mb_dstate0_acc_stored.arrive() + if sk_nt > 0: + bars.mb_dstate0_acc_stored.arrive() gbase += sk_nt - tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one) bars.mb_tmem_done[0].arrive() @@ -2600,8 +2530,8 @@ def compute2_warp_group( num_ctas, cu_seqlens, mWorkItems, - sSched, - lane, + sScheduler, + lane_idx, tmem_hold, warp_idx, sBeta_raw, @@ -2625,13 +2555,14 @@ def compute2_warp_group( (db) for the epilogue's TMA stores, and stages dQ/dK for the epilogue's TMA stores.""" nvvm.setmaxregister(cfg.num_regs_compute_group_2, nvvm.SetMaxRegisterAction.INCREASE) + elect_one = nvvm.elect_sync() nvvm.barrier_cta_sync(cfg.tmem_lifecycle_barrier_id, thread_count=cfg.tmem_user_threads) tmem_base = tmem_hold.load() tmem_col = tmem_base & 0xFFFF tmem_row = tmem_base >> 16 wg1_sp = warp_idx % 4 - channel = wg1_sp * cfg.threads_per_warp + lane - row_addr = (tmem_row + wg1_sp * cfg.threads_per_warp) << 16 + channel = wg1_sp * cfg.threads_per_warp + lane_idx + row_lo_addr = tmem_row << 16 cg2_tidx = channel raw_index = PipelineState.start(phase=0) @@ -2641,14 +2572,16 @@ def compute2_warp_group( dk_restore_part_index = PipelineState.start(phase=0) dgate_last_dstate_smem_index = PipelineState.start(phase=0) gbase = cutlass.Int32(0) - sched_state = PipelineState.start(phase=0) + scheduler_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) FIRST_STATE_CHUNK = 0 if cfg.use_initial_state else 1 while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) - sk_nt = cend - wstart + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) + sk_nt = compute_end - write_start for rev_idx in cutlass.range(sk_nt, unroll=1): - chunk_idx = cend - cutlass.Int32(1) - rev_idx + chunk_idx = compute_end - cutlass.Int32(1) - rev_idx gc = gbase + rev_idx chunk_start = chunk_idx * cfg.b_t raw_stage = gc % cfg.smem_raw_stages @@ -2658,13 +2591,13 @@ def compute2_warp_group( has_dstate = cutlass.Boolean(True) sBetaP_ptr = sBeta_raw.data_ptr() + raw_stage * (cfg.d_k * cfg.b_t) sGate_ptr = sGate_raw.data_ptr() + raw_stage * (cfg.d_k * cfg.b_t) - writes = chunk_idx < wend + writes = chunk_idx < write_end # ---- raw q/k/beta/gate landed: CG0 publishes the decay ring only after - # consuming them, so this wait is CG2's visibility guard --------------- + # consuming them, so this wait is CG2's visibility guard ------------------- bars.mb_k_decay_inv_ready[decay_stage].wait((gc // cfg.smem_decay_stages) % 2) - # ---- per-channel gate factors ---------------------------------------- + # ---- per-channel gate factors -------------------------------------------- f32_seg = channel // 32 f32_dim = channel - f32_seg * 32 f16_seg = channel // 64 @@ -2676,16 +2609,16 @@ def compute2_warp_group( nvvm.fence_proxy("async.shared", space="cta") bars.mb_gate_done[raw_stage].arrive() - # ---- staged raw Q/K: TMEM ring cols for this chunk ------------------- + # ---- staged raw Q/K: TMEM ring cols for this chunk ----------------------- qk_raw_stage = gc % cfg.tmem_qk_raw_stages - qraw_col = tmem_col + cfg.tmem_qraw_inp_offset + qk_raw_stage * (cfg.b_t // 2) - kraw_col = tmem_col + cfg.tmem_kraw_inp_offset + qk_raw_stage * (cfg.b_t // 2) + qraw_col = tmem_col + cfg.tmem_qraw_input_offset + qk_raw_stage * (cfg.b_t // 2) + kraw_col = tmem_col + cfg.tmem_kraw_input_offset + qk_raw_stage * (cfg.b_t // 2) norm_base = qk_raw_stage * (2 * cfg.b_t) bars.mb_qk_raw_ready[qk_raw_stage].wait((gc // cfg.tmem_qk_raw_stages) % 2) - # ---- dGate_last hdot: sum_v sdH[v, c] * S0[c, v] --------------------- + # ---- dGate last hdot: sum over v of sdH[v, c] * S0[c, v] ----------------- dgate_last_val = cutlass.Float32(0.0) - bars.mb_state_inp_ready[gc % 2].wait((gc // 2) % 2) + bars.mb_state_input_ready[gc % 2].wait((gc // 2) % 2) if has_dstate: bars.mb_dstate_smem_ready.wait(dgate_last_dstate_smem_index.phase) dgate_last_dstate_smem_index = advance(dgate_last_dstate_smem_index, 1) @@ -2694,7 +2627,7 @@ def compute2_warp_group( state_vec = nvvm.tcgen05_ld( "32x32b", nvvm.make_tmem_ptr( - row_addr + (tmem_col + cfg.tmem_state_inp_offset + (gc % 2) * (cfg.d_v // 2) + pl * 32 + row_half * 16), cutlass.Float32 + row_lo_addr + (tmem_col + cfg.tmem_state_input_offset + (gc % 2) * (cfg.d_v // 2) + pl * 32 + row_half * 16), cutlass.Float32 ), num=16, ) @@ -2715,9 +2648,9 @@ def compute2_warp_group( part_a, part_b = fadd2(pa0, pb0, pa1, pb1) dgate_last_val = dgate_last_val + (part_a + part_b) bars.mb_dstate_smem_cg2_done.arrive() - bars.mb_state_inp_cg2_done[gc % 2].arrive() + bars.mb_state_input_cg2_done[gc % 2].arrive() - # ---- part-drain accumulators ------------------------------------------- + # ---- part-store accumulators --------------------------------------------- dq_n = cutlass.Array(cutlass.Float32, cfg.b_t, alignment=16) dk_n = cutlass.Array(cutlass.Float32, cfg.b_t, alignment=16) db_regs = cutlass.Array(cutlass.Float32, cfg.b_t, alignment=16) @@ -2728,14 +2661,14 @@ def compute2_warp_group( for t in cutlass.range_constexpr(cfg.b_t): dk_n[t] = cutlass.Float32(0.0) - # ---- dK_restore part drain: (eGl/eG) scale + dGate_last k-dot ---------- + # ---- dK restore part store: (eGl/eG) scale + dGate last K-dot ------------ if has_dstate: bars.mb_dk_restore_part_acc_ready.wait(dk_restore_part_index.phase) dk_restore_part_index = advance(dk_restore_part_index, 1) dk_restore_part_vec = nvvm.tcgen05_ld( - "32x32b", nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_dk_restore_acc_offset), cutlass.Float32), num=cfg.b_t + "32x32b", nvvm.make_tmem_ptr(row_lo_addr + (tmem_col + cfg.tmem_dk_restore_acc_offset), cutlass.Float32), num=cfg.b_t ) - kr_words = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + kraw_col, cutlass.Float32), num=cfg.b_t // 2) + kr_words = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_lo_addr + kraw_col, cutlass.Float32), num=cfg.b_t // 2) for t in cutlass.range_constexpr(cfg.b_t): dk_hat = egl * cute.math.rcp(eg[t], approx=True, ftz=True) * dk_restore_part_vec[t] dk_n[t] = dk_hat @@ -2745,29 +2678,29 @@ def compute2_warp_group( k_v = k_v * sNorm_raw[norm_base + cfg.b_t + t] dgate_last_acc[t % 4] = dgate_last_acc[t % 4] + k_v * dk_hat - # ---- dQ acc drain: eG.scale --------------------------------------------- + # ---- dQ acc store: eG.scale ---------------------------------------------- bars.mb_dq_acc_ready.wait(dq_acc_index.phase) dq_acc_index = advance(dq_acc_index, 1) - dq_vec = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_dq_acc_offset), cutlass.Float32), num=cfg.b_t) + dq_vec = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_lo_addr + (tmem_col + cfg.tmem_dq_acc_offset), cutlass.Float32), num=cfg.b_t) for t2 in cutlass.range_constexpr(cfg.b_t // 2): t = 2 * t2 - es_lo, es_hi = fmul2(eg[t], eg[t + 1], scale, scale) - dq_n[t], dq_n[t + 1] = fmul2(es_lo, es_hi, dq_vec[t], dq_vec[t + 1]) + lo, hi = fmul2(eg[t], eg[t + 1], scale, scale) + dq_n[t], dq_n[t + 1] = fmul2(lo, hi, dq_vec[t], dq_vec[t + 1]) - # ---- dK_inv part drain: (dA - dM) term, 1/eG scale ---------------------- + # ---- dK inv part store: (dA - dM) term, 1/eG scale ----------------------- bars.mb_dk_inv_part_acc_ready.wait(dk_inv_part_index.phase) dk_inv_part_index = advance(dk_inv_part_index, 1) - dk_inv_part_vec = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_dk_inv_acc_offset), cutlass.Float32), num=cfg.b_t) + dk_inv_part_vec = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_lo_addr + (tmem_col + cfg.tmem_dk_inv_acc_offset), cutlass.Float32), num=cfg.b_t) for t in cutlass.range_constexpr(cfg.b_t): dk_n[t] = dk_n[t] + dk_inv_part_vec[t] * cute.math.rcp(eg[t], approx=True, ftz=True) - # ---- dK_decay part drain: -eG scale, seeds dBeta and dGate -------------- + # ---- dK decay part store: -eG scale, seeds dBeta and dGate --------------- bars.mb_dk_decay_part_acc_ready.wait(dk_decay_part_index.phase) dk_decay_part_index = advance(dk_decay_part_index, 1) dk_decay_part_vec = nvvm.tcgen05_ld( - "32x32b", nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_dk_decay_acc_offset), cutlass.Float32), num=cfg.b_t + "32x32b", nvvm.make_tmem_ptr(row_lo_addr + (tmem_col + cfg.tmem_dk_decay_acc_offset), cutlass.Float32), num=cfg.b_t ) - kd_words = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + kraw_col, cutlass.Float32), num=cfg.b_t // 2) + kd_words = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_lo_addr + kraw_col, cutlass.Float32), num=cfg.b_t // 2) for t in cutlass.range_constexpr(cfg.b_t): dk_decay = -eg[t] * dk_decay_part_vec[t] k_pair = cutlass.Vector.from_elements((kd_words[t // 2],), cutlass.Float32).bitcast(cfg.io_dtype) @@ -2778,16 +2711,16 @@ def compute2_warp_group( beta_v = (sBetaP_ptr + f16_seg * (cfg.b_t * 64) + t * 64 + swizzle_xor_128b(t, f16_dim, elem_bytes=2)).load().to(cutlass.Float32) if cutlass.const_expr(cfg.beta_sigmoid): half = cutlass.Float32(0.5) - beta_v = (cute.math.tanh(beta_v * half, approx=True) * half + half).to(cfg.io_dtype).to(cutlass.Float32) + beta_v = sigmoid(beta_v).to(cfg.io_dtype).to(cutlass.Float32) dgate_regs[t] = beta_v * dk_decay dk_n[t] = dk_n[t] + dgate_regs[t] nvvm.tcgen05_wait("load") bars.mb_dqk_acc_done.arrive() - # ---- dGate finalize -------------------------------------------------- - qf_words = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + qraw_col, cutlass.Float32), num=cfg.b_t // 2) - kf_words = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + kraw_col, cutlass.Float32), num=cfg.b_t // 2) + # ---- dGate finalize ------------------------------------------------------ + qf_words = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_lo_addr + qraw_col, cutlass.Float32), num=cfg.b_t // 2) + kf_words = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_lo_addr + kraw_col, cutlass.Float32), num=cfg.b_t // 2) for t in cutlass.range_constexpr(cfg.b_t): q_pair = cutlass.Vector.from_elements((qf_words[t // 2],), cutlass.Float32).bitcast(cfg.io_dtype) k_pair = cutlass.Vector.from_elements((kf_words[t // 2],), cutlass.Float32).bitcast(cfg.io_dtype) @@ -2799,21 +2732,20 @@ def compute2_warp_group( beta_v = (sBetaP_ptr + f16_seg * (cfg.b_t * 64) + t * 64 + swizzle_xor_128b(t, f16_dim, elem_bytes=2)).load().to(cutlass.Float32) if cutlass.const_expr(cfg.beta_sigmoid): half = cutlass.Float32(0.5) - beta_v = (cute.math.tanh(beta_v * half, approx=True) * half + half).to(cfg.io_dtype).to(cutlass.Float32) + beta_v = sigmoid(beta_v).to(cfg.io_dtype).to(cutlass.Float32) dgate_regs[t] = q_v * dq_n[t] + beta_v * db_regs[t] - k_v * (dk_n[t] - dgate_regs[t]) if cutlass.const_expr(cfg.beta_sigmoid): - # after dgate, which consumes db_regs pre-chain-rule db_regs[t] = db_regs[t] * (beta_v - beta_v * beta_v) dgate_regs[cfg.b_t - 1] = dgate_regs[cfg.b_t - 1] + ((dgate_last_acc[0] + dgate_last_acc[1]) + (dgate_last_acc[2] + dgate_last_acc[3])) nvvm.fence_proxy("async.shared", space="cta") bars.mb_beta_done[raw_stage].arrive() - # ---- L2-norm backward row projection --------------------------------- + # ---- L2-norm backward row projection ------------------------------------- if cutlass.const_expr(cfg.l2norm): for grad, qk_col, inv_off in ((dq_n, qraw_col, 0), (dk_n, kraw_col, cfg.b_t)): dots = cutlass.Array(cutlass.Float32, cfg.b_t, alignment=16) for half in cutlass.range_constexpr(2): - p_words = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + qk_col + half * (cfg.b_t // 4), cutlass.Float32), num=cfg.b_t // 4) + p_words = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_lo_addr + qk_col + half * (cfg.b_t // 4), cutlass.Float32), num=cfg.b_t // 4) for tt2 in cutlass.range_constexpr(cfg.b_t // 4): tt = 2 * tt2 t = half * (cfg.b_t // 2) + tt @@ -2827,12 +2759,12 @@ def compute2_warp_group( bfly_lo = cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, dots[t], step, 31, kind=nvvm.Shfl.BFLY)) bfly_hi = cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, dots[t + 1], step, 31, kind=nvvm.Shfl.BFLY)) dots[t], dots[t + 1] = fadd2(dots[t], dots[t + 1], bfly_lo, bfly_hi) - if lane == 0: + if lane_idx == 0: for t in cutlass.range_constexpr(cfg.b_t): sRed1_raw[wg1_sp * cfg.b_t + t] = dots[t] nvvm.barrier_cta_sync(cfg.cg2_sync_barrier_id, thread_count=cfg.cg2_threads) for half in cutlass.range_constexpr(2): - a_words = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + qk_col + half * (cfg.b_t // 4), cutlass.Float32), num=cfg.b_t // 4) + a_words = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_lo_addr + qk_col + half * (cfg.b_t // 4), cutlass.Float32), num=cfg.b_t // 4) for tt2 in cutlass.range_constexpr(cfg.b_t // 4): t = half * (cfg.b_t // 2) + 2 * tt2 a_pair = cutlass.Vector.from_elements((a_words[tt2],), cutlass.Float32).bitcast(cfg.io_dtype) @@ -2848,7 +2780,7 @@ def compute2_warp_group( nvvm.tcgen05_wait("load") bars.mb_qk_raw_done[qk_raw_stage].arrive() - # ---- stage dQ/dK for the epilogue TMA stores ------------------------- + # ---- stage dQ/dK for the epilogue TMA stores ----------------------------- dq_stage = gc % cfg.smem_dq_stages dk_stage = gc % cfg.smem_dk_stages bars.mb_dq_tmastg_done[dq_stage].wait(((gc // cfg.smem_dq_stages) + 1) % 2) @@ -2863,19 +2795,19 @@ def compute2_warp_group( bars.mb_dq_tmastg_ready[dq_stage].arrive() bars.mb_dk_tmastg_ready[dk_stage].arrive() - # ---- dGate_last add -------------------------------------------------- + # ---- dGate last add ------------------------------------------------------ if has_dstate: if chunk_idx >= FIRST_STATE_CHUNK: dgate_regs[cfg.b_t - 1] = dgate_regs[cfg.b_t - 1] + egl * dgate_last_val - # ---- dGate reverse cumsum -------------------------------------------- + # ---- dGate reverse cumsum ------------------------------------------------ suffix = cutlass.Float32(0.0) for rt in cutlass.range_constexpr(cfg.b_t): t = cfg.b_t - 1 - rt suffix = suffix + dgate_regs[t] dgate_regs[t] = suffix - # ---- stage dGate + dBeta for the epilogue TMA stores ----------------- + # ---- stage dGate + dBeta for the epilogue TMA stores --------------------- dgate_stage = gc % cfg.smem_dgate_stages bars.mb_dgate_tmastg_done[dgate_stage].wait(((gc // cfg.smem_dgate_stages) + 1) % 2) db_stage = gc % cfg.smem_db_stages @@ -2890,7 +2822,7 @@ def compute2_warp_group( bars.mb_db_tmastg_ready[db_stage].arrive() raw_index = advance(raw_index, cfg.smem_raw_stages) gbase += sk_nt - tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one) bars.mb_tmem_done[0].arrive() @@ -2917,7 +2849,7 @@ def build_descs_body( base_dwo, base_dbo, base_checkpoint, - desc_ws: cute.Tensor, + desc_workspace: cute.Tensor, cu_seqlens: cute.Tensor, q: cute.Tensor, k: cute.Tensor, @@ -2954,20 +2886,20 @@ def build_descs_body( prologue kernel after its order pass; warps past the array count fall through the widx guards.""" arr_words = n_batch * cutlass.Int32(TENSOR_MAP_QWORDS) - sub0 = cute.make_tensor(desc_ws.iterator, cute.make_layout((arr_words,), stride=(1,))) - sub1 = cute.make_tensor(desc_ws.iterator + arr_words, cute.make_layout((arr_words,), stride=(1,))) - sub2 = cute.make_tensor(desc_ws.iterator + 2 * arr_words, cute.make_layout((arr_words,), stride=(1,))) - sub3 = cute.make_tensor(desc_ws.iterator + 3 * arr_words, cute.make_layout((arr_words,), stride=(1,))) - sub4 = cute.make_tensor(desc_ws.iterator + 4 * arr_words, cute.make_layout((arr_words,), stride=(1,))) - sub5 = cute.make_tensor(desc_ws.iterator + 5 * arr_words, cute.make_layout((arr_words,), stride=(1,))) - sub6 = cute.make_tensor(desc_ws.iterator + 6 * arr_words, cute.make_layout((arr_words,), stride=(1,))) - sub7 = cute.make_tensor(desc_ws.iterator + 7 * arr_words, cute.make_layout((arr_words,), stride=(1,))) - sub8 = cute.make_tensor(desc_ws.iterator + 8 * arr_words, cute.make_layout((arr_words,), stride=(1,))) - sub9 = cute.make_tensor(desc_ws.iterator + 9 * arr_words, cute.make_layout((arr_words,), stride=(1,))) - sub10 = cute.make_tensor(desc_ws.iterator + 10 * arr_words, cute.make_layout((arr_words,), stride=(1,))) - sub11 = cute.make_tensor(desc_ws.iterator + 11 * arr_words, cute.make_layout((arr_words,), stride=(1,))) - sub12 = cute.make_tensor(desc_ws.iterator + 12 * arr_words, cute.make_layout((arr_words,), stride=(1,))) - sub13 = cute.make_tensor(desc_ws.iterator + 13 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub0 = cute.make_tensor(desc_workspace.iterator, cute.make_layout((arr_words,), stride=(1,))) + sub1 = cute.make_tensor(desc_workspace.iterator + arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub2 = cute.make_tensor(desc_workspace.iterator + 2 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub3 = cute.make_tensor(desc_workspace.iterator + 3 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub4 = cute.make_tensor(desc_workspace.iterator + 4 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub5 = cute.make_tensor(desc_workspace.iterator + 5 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub6 = cute.make_tensor(desc_workspace.iterator + 6 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub7 = cute.make_tensor(desc_workspace.iterator + 7 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub8 = cute.make_tensor(desc_workspace.iterator + 8 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub9 = cute.make_tensor(desc_workspace.iterator + 9 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub10 = cute.make_tensor(desc_workspace.iterator + 10 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub11 = cute.make_tensor(desc_workspace.iterator + 11 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub12 = cute.make_tensor(desc_workspace.iterator + 12 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub13 = cute.make_tensor(desc_workspace.iterator + 13 * arr_words, cute.make_layout((arr_words,), stride=(1,))) if widx == 0: if nvvm.elect_sync(): @@ -3031,7 +2963,7 @@ def build_descs_body( def prologue_kernel( run_order: cutlass.Constexpr[bool], order_gen: cutlass.Constexpr[bool], - has_sched: cutlass.Constexpr[bool], + has_scheduler: cutlass.Constexpr[bool], b_t: cutlass.Constexpr[int], base_q: cutlass.GridConstant[cuda.tensor_map.TensorMap], base_k: cutlass.GridConstant[cuda.tensor_map.TensorMap], @@ -3047,7 +2979,7 @@ def prologue_kernel( base_dwo: cutlass.GridConstant[cuda.tensor_map.TensorMap], base_dbo: cutlass.GridConstant[cuda.tensor_map.TensorMap], base_checkpoint: cutlass.GridConstant[cuda.tensor_map.TensorMap], - desc_ws: cute.Tensor, + desc_workspace: cute.Tensor, cu_seqlens: cute.Tensor, q: cute.Tensor, k: cute.Tensor, @@ -3066,7 +2998,7 @@ def prologue_kernel( mStaging: cute.Tensor | None, mCount: cute.Tensor, mWorkItems: cute.Tensor, - mSched: cute.Tensor | None, + mScheduler: cute.Tensor | None, n_batch: cutlass.Int32, q_rs: cutlass.Int32, k_rs: cutlass.Int32, @@ -3086,7 +3018,7 @@ def prologue_kernel( ) -> None: """Single-CTA prologue. Under ``run_order`` this kernel is the first work-item-table consumer, so it LPT-orders the table and zeroes both - consumers' sched rings via :func:`order_body`; it then builds the + consumers' scheduler rings via :func:`order_body`; it then builds the per-batch TMA-descriptor arrays via :func:`build_descs_body`, one warp per array (the extra warps only take part in the order phase).""" tidx, _, _ = cute.arch.thread_idx() @@ -3099,10 +3031,10 @@ def prologue_kernel( n_heads_out = cutlass.Int32(gate.shape[1]) order_body( order_gen, - has_sched, + has_scheduler, b_t, ORDER_THREADS, - ORDER_ELEMS, + ORDER_ELEMENTS, tidx, n_heads_out, n_heads_out * n_batch, @@ -3110,7 +3042,7 @@ def prologue_kernel( mStaging, mCount, mWorkItems, - mSched, + mScheduler, sKey, sIdx, sSpread, @@ -3131,7 +3063,7 @@ def prologue_kernel( base_dwo, base_dbo, base_checkpoint, - desc_ws, + desc_workspace, cu_seqlens, q, k, @@ -3172,7 +3104,7 @@ def prologue( b_t: cutlass.Constexpr[int], run_order: cutlass.Constexpr[bool], order_gen: cutlass.Constexpr[bool], - has_sched: cutlass.Constexpr[bool], + has_scheduler: cutlass.Constexpr[bool], q: cute.Tensor, k: cute.Tensor, v: cute.Tensor, @@ -3191,7 +3123,7 @@ def prologue( work_item_staging: cute.Tensor | None, work_count: cute.Tensor, work_items: cute.Tensor, - sched_all: cute.Tensor | None, + scheduler_all: cute.Tensor | None, tensormap_workspace: cute.Tensor, stream: cuda_driver.CUstream, ): @@ -3206,7 +3138,7 @@ def prologue( d_k = q.shape[2] d_v = v.shape[2] bpe = io_dtype.width // 8 - granu = 128 // bpe + granule = 128 // bpe seqlen = q.shape[0] q_headed = cute.make_tensor(q.iterator, cute.make_layout((d_k, h_q, seqlen), stride=(1, q.stride[1], q.stride[0]))) @@ -3224,19 +3156,19 @@ def prologue( dbo_headed = cute.make_tensor(dbo.iterator, cute.make_layout((d_k, ho, seqlen), stride=(1, dbo.stride[1], dbo.stride[0]))) swz = cuda.TensorMapSwizzle.s128b - base_q = cuda.create_tensor_map_tiled_from_view(q_headed, box_dims=(granu, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) - base_k = cuda.create_tensor_map_tiled_from_view(k_headed, box_dims=(granu, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) - base_v = cuda.create_tensor_map_tiled_from_view(v_headed, box_dims=(granu, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_q = cuda.create_tensor_map_tiled_from_view(q_headed, box_dims=(granule, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_k = cuda.create_tensor_map_tiled_from_view(k_headed, box_dims=(granule, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_v = cuda.create_tensor_map_tiled_from_view(v_headed, box_dims=(granule, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) base_gate = cuda.create_tensor_map_tiled_from_view(gate_headed, box_dims=(32, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) - base_do = cuda.create_tensor_map_tiled_from_view(do_headed, box_dims=(granu, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) - base_beta = cuda.create_tensor_map_tiled_from_view(beta_headed, box_dims=(granu, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) - base_w = cuda.create_tensor_map_tiled_from_view(w_headed, box_dims=(granu, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) - base_dq = cuda.create_tensor_map_tiled_from_view(dq_headed, box_dims=(granu, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) - base_dk = cuda.create_tensor_map_tiled_from_view(dk_headed, box_dims=(granu, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) - base_dv = cuda.create_tensor_map_tiled_from_view(dv_headed, box_dims=(granu, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_do = cuda.create_tensor_map_tiled_from_view(do_headed, box_dims=(granule, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_beta = cuda.create_tensor_map_tiled_from_view(beta_headed, box_dims=(granule, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_w = cuda.create_tensor_map_tiled_from_view(w_headed, box_dims=(granule, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_dq = cuda.create_tensor_map_tiled_from_view(dq_headed, box_dims=(granule, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_dk = cuda.create_tensor_map_tiled_from_view(dk_headed, box_dims=(granule, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_dv = cuda.create_tensor_map_tiled_from_view(dv_headed, box_dims=(granule, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) base_dgate = cuda.create_tensor_map_tiled_from_view(dgate_headed, box_dims=(32, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) - base_dwo = cuda.create_tensor_map_tiled_from_view(dwo_headed, box_dims=(granu, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) - base_dbo = cuda.create_tensor_map_tiled_from_view(dbo_headed, box_dims=(granu, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_dwo = cuda.create_tensor_map_tiled_from_view(dwo_headed, box_dims=(granule, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_dbo = cuda.create_tensor_map_tiled_from_view(dbo_headed, box_dims=(granule, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) checkpoint_view = cute.make_tensor( state_checkpoints.iterator, @@ -3250,7 +3182,7 @@ def prologue( prologue_kernel( run_order, order_gen, - has_sched, + has_scheduler, b_t, base_q, base_k, @@ -3285,7 +3217,7 @@ def prologue( work_item_staging, work_count, work_items, - sched_all, + scheduler_all, cutlass.Int32(batch_size), cutlass.Int32(q.stride[0]), cutlass.Int32(k.stride[0]), @@ -3319,14 +3251,14 @@ def host( d_final_state: cute.Tensor | None, work_items: cute.Tensor | None, work_count: cute.Tensor | None, - sched_ctr: cute.Tensor | None, + scheduler_counter: cute.Tensor | None, tensormap_workspace: cute.Tensor, scale: cutlass.Float32, stream, ) -> None: num_sequences = cu_seqlens.shape[0] - 1 - # ---- launch ------------------------------------------------------------------ + # ---- launch ---------------------------------------------------------------------- n_desc = num_sequences grid_shape = (cfg.max_active_clusters, 1, 1) kernel( @@ -3343,7 +3275,7 @@ def host( d_final_state, work_items, work_count, - sched_ctr, + scheduler_counter, scale, ).launch( grid=grid_shape, @@ -3368,7 +3300,7 @@ def kernel( mDstate_in: cute.Tensor | None, mWorkItems: cute.Tensor, mCount: cute.Tensor, - mSched: cute.Tensor | None, + mScheduler: cute.Tensor | None, scale: cutlass.Float32, ) -> None: """BT=16 GDN-2 backward kernel (persistent, 16 warps).""" @@ -3376,7 +3308,7 @@ def kernel( bidx = cute.arch.block_idx()[0] num_ctas = cute.arch.grid_dim()[0] warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - lane = tidx % cfg.threads_per_warp + lane_idx = tidx % cfg.threads_per_warp total_tiles = mCount[0] assert cu_seqlens.element_type in (cutlass.Int32, cutlass.Int64) @@ -3403,7 +3335,7 @@ def kernel( SMEM = cutlass.AddressSpace.smem bars = make_gdn2_bwd_bars(cfg) tmem_hold = cutlass.Array(cutlass.Int32, 1, space=SMEM, alignment=4) - sSched = cutlass.Array(cutlass.Int32, cfg.sched_stages, space=SMEM, alignment=16) + sScheduler = cutlass.Array(cutlass.Int32, cfg.scheduler_stages, space=SMEM, alignment=16) bpe = cfg.io_dtype.width // 8 SWZ = 2 LEAD = 16 @@ -3439,7 +3371,7 @@ def kernel( sDwOut_raw = cutlass.Array(cfg.io_dtype, cfg.dwo_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) sDb_raw = cutlass.Array(cfg.io_dtype, cfg.db_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) - sState_alt = SmemTile( + sState_trans = SmemTile( base=sState_raw.data_ptr().toint(), elems_per_stage=((cfg.state_cosize) // (cfg.smem_state_stages)) * bpe, stages=cfg.smem_state_stages, @@ -3447,7 +3379,7 @@ def kernel( stride_byte_offset=STRIDE, layout=SWZ, ) - sState_direct = SmemTile( + sState = SmemTile( base=sState_raw.data_ptr().toint(), elems_per_stage=((cfg.state_cosize) // (cfg.smem_state_stages)) * bpe, stages=cfg.smem_state_stages, @@ -3455,7 +3387,7 @@ def kernel( stride_byte_offset=STRIDE, layout=SWZ, ) - sK_decay_lead16 = SmemTile( + sK_decay = SmemTile( base=sK_decay_raw.data_ptr().toint(), elems_per_stage=((cfg.operand_cosize) // (cfg.smem_decay_stages)) * bpe, stages=cfg.smem_decay_stages, @@ -3463,7 +3395,7 @@ def kernel( stride_byte_offset=STRIDE, layout=SWZ, ) - sK_restore_lead16 = SmemTile( + sK_restore = SmemTile( base=sK_restore_raw.data_ptr().toint(), elems_per_stage=((cfg.operand_cosize) // (cfg.smem_decay_stages)) * bpe, stages=cfg.smem_decay_stages, @@ -3471,7 +3403,7 @@ def kernel( stride_byte_offset=STRIDE, layout=SWZ, ) - sDo_lead16 = SmemTile( + sDo = SmemTile( base=sDo_raw.data_ptr().toint(), elems_per_stage=((cfg.raw_v_cosize) // (cfg.smem_raw_stages)) * bpe, stages=cfg.smem_raw_stages, @@ -3479,7 +3411,7 @@ def kernel( stride_byte_offset=STRIDE, layout=SWZ, ) - sDo_amaj = SmemTile( + sDo_trans = SmemTile( base=sDo_raw.data_ptr().toint(), elems_per_stage=((cfg.raw_v_cosize) // (cfg.smem_raw_stages)) * bpe, stages=cfg.smem_raw_stages, @@ -3487,7 +3419,7 @@ def kernel( stride_byte_offset=STRIDE, layout=SWZ, ) - sU_lead16 = SmemTile( + sU = SmemTile( base=sU_raw.data_ptr().toint(), elems_per_stage=((cfg.b_t * cfg.d_v) // (1)) * bpe, stages=1, @@ -3495,7 +3427,7 @@ def kernel( stride_byte_offset=STRIDE, layout=SWZ, ) - sDy_lead16 = SmemTile( + sDy = SmemTile( base=sDy_raw.data_ptr().toint(), elems_per_stage=((cfg.b_t * cfg.d_v) // (1)) * bpe, stages=1, @@ -3503,7 +3435,7 @@ def kernel( stride_byte_offset=STRIDE, layout=SWZ, ) - sDstate_alt = SmemTile( + sDstate_trans = SmemTile( base=sDstate_raw.data_ptr().toint(), elems_per_stage=((cfg.d_k * cfg.d_v) // (1)) * bpe, stages=1, @@ -3520,7 +3452,7 @@ def kernel( stride_byte_offset=STRIDE, layout=SWZ, ) - sK_inv_amaj = SmemTile( + sK_inv_trans = SmemTile( base=sK_inv_raw.data_ptr().toint(), elems_per_stage=((cfg.operand_cosize) // (cfg.smem_decay_stages)) * bpe, stages=cfg.smem_decay_stages, @@ -3584,23 +3516,23 @@ def kernel( bars.mb_state_done[stage].init() bars.mb_state_cg0_done[stage].init() for stage in cutlass.range_constexpr(2): - bars.mb_state_inp_ready[stage].init() - bars.mb_state_inp_done[stage].init() - bars.mb_state_inp_cg2_done[stage].init() + bars.mb_state_input_ready[stage].init() + bars.mb_state_input_done[stage].init() + bars.mb_state_input_cg2_done[stage].init() elif warp_idx == cfg.tcgen05_mma_warp_id: if elect_one: bars.mb_state_k_acc_ready.init() - bars.mb_y_inp_ready.init() + bars.mb_y_input_ready.init() bars.mb_u_acc_ready.init() bars.mb_u_smem_ready.init() bars.mb_du_acc_ready.init() - bars.mb_du_inp_ready.init() + bars.mb_du_input_ready.init() bars.mb_dy_acc_ready.init() - bars.mb_neg_dy_inp_ready.init() + bars.mb_neg_dy_input_ready.init() bars.mb_dy_smem_ready.init() bars.mb_dy_smem_done.init() bars.mb_dstate_acc_ready.init() - bars.mb_dstate_inp_ready.init() + bars.mb_dstate_input_ready.init() bars.mb_dstate_smem_ready.init() bars.mb_dstate_smem_done.init() bars.mb_dstate_smem_cg2_done.init() @@ -3649,9 +3581,9 @@ def kernel( for stage in cutlass.range_constexpr(cfg.smem_dwo_stages): bars.mb_dwo_tmastg_ready[stage].init() bars.mb_dwo_tmastg_done[stage].init() - for stage in cutlass.range_constexpr(cfg.sched_stages): - bars.mb_sched_ready[stage].init() - bars.mb_sched_done[stage].init() + for stage in cutlass.range_constexpr(cfg.scheduler_stages): + bars.mb_scheduler_ready[stage].init() + bars.mb_scheduler_done[stage].init() diag_zero = cfg.io_dtype(0.0) for diag_idx in cutlass.range(tidx, cfg.diag_cosize, cfg.threads_per_cta, unroll=1): sState_scale_diag_raw[diag_idx] = diag_zero @@ -3665,8 +3597,8 @@ def kernel( num_ctas, cu_seqlens, mWorkItems, - mSched, - sSched, + mScheduler, + sScheduler, q_tx_bytes, k_tx_bytes, gate_tx_bytes, @@ -3700,8 +3632,8 @@ def kernel( num_ctas, cu_seqlens, mWorkItems, - sSched, - lane, + sScheduler, + lane_idx, sK_decay_raw, sK_inv_raw, sU_raw, @@ -3717,20 +3649,20 @@ def kernel( num_ctas, cu_seqlens, mWorkItems, - sSched, + sScheduler, tmem_hold, - sState_alt, - sState_direct, - sK_decay_lead16, - sK_inv_amaj, - sK_restore_lead16, - sDo_lead16, - sDo_amaj, + sState_trans, + sState, + sK_decay, + sK_inv_trans, + sK_restore, + sDo, + sDo_trans, sQ_decay_trans, sK_decay_trans, - sU_lead16, - sDy_lead16, - sDstate_alt, + sU, + sDy, + sDstate_trans, sIntermediate, sState_scale_diag, bars, @@ -3743,8 +3675,8 @@ def kernel( num_ctas, cu_seqlens, mWorkItems, - sSched, - lane, + sScheduler, + lane_idx, sK_inv_raw, sQ_decay_raw, sDo_raw, @@ -3772,8 +3704,8 @@ def kernel( num_ctas, cu_seqlens, mWorkItems, - sSched, - lane, + sScheduler, + lane_idx, tmem_hold, warp_idx, scale, @@ -3803,8 +3735,8 @@ def kernel( num_ctas, cu_seqlens, mWorkItems, - sSched, - lane, + sScheduler, + lane_idx, tmem_hold, warp_idx, sBeta_raw, @@ -3827,8 +3759,8 @@ def kernel( num_ctas, cu_seqlens, mWorkItems, - sSched, - lane, + sScheduler, + lane_idx, tmem_hold, warp_idx, mDstate0, @@ -3863,10 +3795,10 @@ class Gdn2BwdCfg: v_ratio: int n_heads_out: int max_active_clusters: int - dyn_sched: bool = False - sched_stages: int = 8 + dynamic_scheduling: bool = False + scheduler_stages: int = 8 - # ---- fixed constants stamped from CFG by build_cfg --------------------------- + # ---- fixed constants stamped from CFG at build time ------------------------------ compute_group_0_warp_ids: tuple = CFG.COMPUTE_GROUP_0_WARP_IDS compute_group_2_warp_ids: tuple = CFG.COMPUTE_GROUP_2_WARP_IDS compute_group_1_warp_ids: tuple = CFG.COMPUTE_GROUP_1_WARP_IDS @@ -3884,7 +3816,7 @@ class Gdn2BwdCfg: num_regs_compute_group_2: int = CFG.NUM_REGS_COMPUTE_GROUP_2 num_regs_other: int = CFG.NUM_REGS_OTHER - # ---- named barrier slots (ids 1-4; 0 is the CTA-wide sync) ------------------- + # ---- named barrier slots (ids 1-4; 0 is the CTA-wide sync) ----------------------- cg0_sync_barrier_id: int = 1 cg0_threads: int = 0 cg2_sync_barrier_id: int = 2 @@ -3894,7 +3826,7 @@ class Gdn2BwdCfg: cg1_sync_barrier_id: int = 4 cg1_threads: int = 0 - # ---- SMEM / TMEM stage counts + TMEM column offsets -------------------------- + # ---- SMEM / TMEM stage counts + TMEM column offsets ------------------------------ smem_raw_stages: int = CFG.SMEM_RAW_STAGES smem_state_stages: int = CFG.SMEM_STATE_STAGES smem_decay_stages: int = CFG.SMEM_DECAY_STAGES @@ -3907,7 +3839,7 @@ class Gdn2BwdCfg: smem_dwo_stages: int = CFG.SMEM_DWO_STAGES intermediate_tiles: int = 5 tmem_dstate_acc_offset: int = 0 - tmem_dstate_inp_offset: int = 0 + tmem_dstate_input_offset: int = 0 tmem_state_k_acc_offset: int = 0 tmem_u_acc_offset: int = 0 tmem_du_acc_offset: int = 0 @@ -3917,15 +3849,15 @@ class Gdn2BwdCfg: tmem_dk_inv_acc_offset: int = 0 tmem_dk_restore_acc_offset: int = 0 tmem_qk_raw_stages: int = 4 - tmem_qraw_inp_offset: int = 0 - tmem_kraw_inp_offset: int = 0 - tmem_y_inp_offset: int = 0 - tmem_du_inp_offset: int = 0 - tmem_neg_dy_inp_offset: int = 0 - tmem_state_inp_offset: int = 0 + tmem_qraw_input_offset: int = 0 + tmem_kraw_input_offset: int = 0 + tmem_y_input_offset: int = 0 + tmem_du_input_offset: int = 0 + tmem_neg_dy_input_offset: int = 0 + tmem_state_input_offset: int = 0 buffer_align_bytes: int = CFG.BUFFER_ALIGN_BYTES - # ---- buffer cosizes / TMA bytes stamped by build_cfg ------------------------- + # ---- buffer cosizes / TMA bytes stamped at build time ---------------------------- raw_qk_cosize: int = 0 raw_v_cosize: int = 0 raw_gate_cosize: int = 0 @@ -3957,7 +3889,7 @@ def build_cfg( v_ratio: int, n_heads_out: int, max_active_clusters: int, - dyn_sched: bool = False, + dynamic_scheduling: bool = False, ) -> Gdn2BwdCfg: if io_dtype not in (cutlass.Float16, cutlass.BFloat16): raise ValueError(f"io_dtype={io_dtype} not supported; only Float16 and BFloat16 are supported") @@ -3975,7 +3907,7 @@ def build_cfg( v_ratio=v_ratio, n_heads_out=n_heads_out, max_active_clusters=max_active_clusters, - dyn_sched=dyn_sched, + dynamic_scheduling=dynamic_scheduling, ) cfg.threads_per_cta = 16 * cfg.threads_per_warp cfg.cg0_threads = len(cfg.compute_group_0_warp_ids) * cfg.threads_per_warp @@ -3986,29 +3918,22 @@ def build_cfg( ) * cfg.threads_per_warp cfg.tmem_dstate_acc_offset = 0 - cfg.tmem_dstate_inp_offset = cfg.d_k - cfg.tmem_state_inp_offset = cfg.tmem_dstate_inp_offset + cfg.d_k // 2 - cfg.tmem_state_k_acc_offset = cfg.tmem_state_inp_offset + cfg.d_v + cfg.tmem_dstate_input_offset = cfg.d_k + cfg.tmem_state_input_offset = cfg.tmem_dstate_input_offset + cfg.d_k // 2 + cfg.tmem_state_k_acc_offset = cfg.tmem_state_input_offset + cfg.d_v cfg.tmem_u_acc_offset = cfg.tmem_state_k_acc_offset + cfg.b_t cfg.tmem_du_acc_offset = cfg.tmem_u_acc_offset + cfg.b_t - # dY overwrites the state_k slot: WG1's Y staging consumes state_k - # before the dY = dU @ T_inv MMA writes (du_inp chain), and the dY - # readback precedes state_k(c+1) = state @ K_decay^T via - # neg_dy_ready -> the -dY @ K_decay dstate MMA -> in-order MMA cfg.tmem_dy_acc_offset = cfg.tmem_state_k_acc_offset cfg.tmem_dq_acc_offset = cfg.tmem_du_acc_offset + cfg.b_t cfg.tmem_dk_decay_acc_offset = cfg.tmem_dq_acc_offset + cfg.b_t cfg.tmem_dk_inv_acc_offset = cfg.tmem_dk_decay_acc_offset + cfg.b_t cfg.tmem_dk_restore_acc_offset = cfg.tmem_dk_inv_acc_offset + cfg.b_t - cfg.tmem_y_inp_offset = cfg.tmem_dk_restore_acc_offset + cfg.b_t - # -dY overwrites the y_inp slot: U = Y @ T_inv consumed Y before the dY - # block runs (u_acc_ready wait), and y_inp(c+1) is gated by - # state_k_acc_ready(c+1) whose commit covers the -dY @ K_decay MMA (c) - cfg.tmem_neg_dy_inp_offset = cfg.tmem_y_inp_offset - cfg.tmem_du_inp_offset = cfg.tmem_y_inp_offset + cfg.b_t // 2 - cfg.tmem_qraw_inp_offset = cfg.tmem_du_inp_offset + cfg.b_t // 2 - cfg.tmem_kraw_inp_offset = cfg.tmem_qraw_inp_offset + cfg.tmem_qk_raw_stages * (cfg.b_t // 2) - assert cfg.tmem_kraw_inp_offset + cfg.tmem_qk_raw_stages * (cfg.b_t // 2) <= 512 + cfg.tmem_y_input_offset = cfg.tmem_dk_restore_acc_offset + cfg.b_t + cfg.tmem_neg_dy_input_offset = cfg.tmem_y_input_offset + cfg.tmem_du_input_offset = cfg.tmem_y_input_offset + cfg.b_t // 2 + cfg.tmem_qraw_input_offset = cfg.tmem_du_input_offset + cfg.b_t // 2 + cfg.tmem_kraw_input_offset = cfg.tmem_qraw_input_offset + cfg.tmem_qk_raw_stages * (cfg.b_t // 2) + assert cfg.tmem_kraw_input_offset + cfg.tmem_qk_raw_stages * (cfg.b_t // 2) <= 512 cfg.raw_qk_cosize = cfg.smem_raw_stages * cfg.d_k * cfg.b_t cfg.raw_v_cosize = cfg.smem_raw_stages * cfg.d_v * cfg.b_t @@ -4031,7 +3956,7 @@ def build_cfg( TENSORMAP_STATIC_SLOTS = 0 -# ---- Torch adapter / host-side compilation --------------------------------------- +# ---- Torch adapter / host-side compilation ------------------------------------------- @lru_cache(maxsize=None) @@ -4048,10 +3973,10 @@ def get_compiled_cache( gate_lower_bound: float, beta_sigmoid: bool, use_initial_state: bool, - dyn_sched: bool, + dynamic_scheduling: bool, order_in_prologue: bool, order_gen: bool, - has_sched: bool, + has_scheduler: bool, ): return {} @@ -4085,8 +4010,8 @@ def chunk_gdn2_bwd_sm100( use_beta_sigmoid: bool = False, work_items=None, work_count=None, - sched_ctr=None, - sched_all=None, + scheduler_counter=None, + scheduler_all=None, work_item_scratch=None, order_in_prologue: bool = False, tensormap_workspace, @@ -4131,9 +4056,9 @@ def chunk_gdn2_bwd_sm100( use_beta_sigmoid: ``beta`` holds logits; sigmoid in-kernel work_items/work_count: split-K table (``common/split_k.py``, REQUIRED; an uncut table row is the whole (b, h) sequence); each item - computes chunks ``[wstart, cend)`` backward and writes - gradients only for ``[wstart, wend)`` - sched_ctr: ``(2,)`` int32 zeroed scratch enabling the dynamic + computes chunks ``[write_start, compute_end)`` backward and writes + gradients only for ``[write_start, write_end)`` + scheduler_counter: ``(2,)`` int32 zeroed scratch enabling the dynamic (work-stealing) tile scheduler tensormap_workspace: ``tensormap_workspace_bytes(module, B)`` bytes, 128-byte aligned, for the per-(batch, head) TMA-descriptor @@ -4146,11 +4071,13 @@ def chunk_gdn2_bwd_sm100( use_dstate_in = d_final_state is not None use_dstate0 = d_initial_state is not None if work_items is None or work_count is None: - raise ValueError("work_items/work_count are required (the split-table stage builds them for every launch)") - dyn_sched = sched_ctr is not None + raise ValueError( + "work_items/work_count are required (built by the split-table stage, or by an order-generating prologue when work_item_scratch is None)" + ) + dynamic_scheduling = scheduler_counter is not None order_gen = work_item_scratch is None - if order_in_prologue and sched_all is None: - raise ValueError("order_in_prologue requires sched_all (the prologue zeroes both consumers' sched rings)") + if order_in_prologue and scheduler_all is None: + raise ValueError("order_in_prologue requires scheduler_all (the prologue zeroes both consumers' scheduler rings)") for name, t in (("state_checkpoints", state_checkpoints), ("beta", beta), ("w", w), ("dbeta", dbeta), ("dw", dw)): if str(t.dtype).split(".")[-1] != str(q.dtype).split(".")[-1]: raise ValueError(f"{name} dtype must match the io dtype: got {t.dtype} with io {q.dtype}") @@ -4180,10 +4107,10 @@ def chunk_gdn2_bwd_sm100( gate_lower_bound, use_beta_sigmoid, use_initial_state, - dyn_sched, + dynamic_scheduling, order_in_prologue, order_gen, - sched_all is not None, + scheduler_all is not None, ) if "compiled" not in cache: @@ -4202,7 +4129,7 @@ def chunk_gdn2_bwd_sm100( v_ratio=HO // HV, n_heads_out=HO, max_active_clusters=multiprocessor_count(current_device()), - dyn_sched=dyn_sched, + dynamic_scheduling=dynamic_scheduling, ) dstate0_cute = None @@ -4215,10 +4142,10 @@ def chunk_gdn2_bwd_sm100( wi_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1), divisibility=1) wc_cute = from_dlpack(work_count, assumed_align=4).mark_layout_dynamic() sc_cute = None - if dyn_sched: - sc_cute = from_dlpack(sched_ctr, assumed_align=4).mark_layout_dynamic() + if dynamic_scheduling: + sc_cute = from_dlpack(scheduler_counter, assumed_align=4).mark_layout_dynamic() - tensormap_ws_cute = from_dlpack(tensormap_workspace, assumed_align=128).mark_layout_dynamic() + tensormap_workspace_cute = from_dlpack(tensormap_workspace, assumed_align=128).mark_layout_dynamic() state_checkpoints_cute = from_dlpack(state_checkpoints, assumed_align=16).mark_layout_dynamic(leading_dim=len(state_checkpoints.shape) - 1) a_log_cute = from_dlpack(a_log, assumed_align=4) if a_log is not None else None @@ -4241,7 +4168,7 @@ def chunk_gdn2_bwd_sm100( wi_cute, wc_cute, sc_cute, - tensormap_ws_cute, + tensormap_workspace_cute, scale, cu_stream, options="--enable-tvm-ffi --opt-level 2", @@ -4264,7 +4191,7 @@ def chunk_gdn2_bwd_sm100( dbo_pl = from_dlpack(dbeta, assumed_align=16).mark_layout_dynamic(leading_dim=2) state_checkpoints_pl = from_dlpack(state_checkpoints, assumed_align=16).mark_layout_dynamic(leading_dim=3) cu_pl = from_dlpack(cu_seqlens, assumed_align=8).mark_layout_dynamic() - ws_pl = from_dlpack(tensormap_workspace, assumed_align=128).mark_layout_dynamic() + workspace_pl = from_dlpack(tensormap_workspace, assumed_align=128).mark_layout_dynamic() staging_pl = None if not order_gen: staging_pl = from_dlpack(work_item_scratch, assumed_align=16) @@ -4272,16 +4199,16 @@ def chunk_gdn2_bwd_sm100( work_items_pl = from_dlpack(work_items, assumed_align=16) work_items_pl.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1), divisibility=1) work_count_pl = from_dlpack(work_count, assumed_align=4).mark_layout_dynamic() - sched_all_pl = None - if sched_all is not None: - sched_all_pl = from_dlpack(sched_all, assumed_align=4).mark_layout_dynamic() + scheduler_all_pl = None + if scheduler_all is not None: + scheduler_all_pl = from_dlpack(scheduler_all, assumed_align=4).mark_layout_dynamic() cache["prologue"] = cute.compile( prologue, io_dtype, CFG.B_T, order_in_prologue, order_gen, - sched_all is not None, + scheduler_all is not None, q_pl, k_pl, v_pl, @@ -4300,8 +4227,8 @@ def chunk_gdn2_bwd_sm100( staging_pl, work_count_pl, work_items_pl, - sched_all_pl, - ws_pl, + scheduler_all_pl, + workspace_pl, cu_stream, options="--enable-tvm-ffi", ) @@ -4324,7 +4251,7 @@ def chunk_gdn2_bwd_sm100( work_item_scratch if not order_gen else None, work_count, work_items, - sched_all, + scheduler_all, tensormap_workspace, cu_stream, ) @@ -4340,7 +4267,7 @@ def chunk_gdn2_bwd_sm100( d_final_state, work_items, work_count, - sched_ctr, + scheduler_counter, tensormap_workspace, scale, cu_stream, @@ -4369,8 +4296,8 @@ def run_bwd( d_final_state, work_items, work_count, - sched_ctr, - sched_all, + scheduler_counter, + scheduler_all, work_item_scratch, tensormap_workspace, scale, @@ -4401,7 +4328,7 @@ def run_bwd( work_item_scratch, work_count, work_items, - sched_all, + scheduler_all, tensormap_workspace, cu_stream, ) @@ -4417,7 +4344,7 @@ def run_bwd( d_final_state, work_items, work_count, - sched_ctr, + scheduler_counter, tensormap_workspace, scale, cu_stream, diff --git a/python/cudnn/linear_attention/frost/kernel/gdn2_prefill_config.py b/python/cudnn/linear_attention/frost/kernel/gdn2_prefill_config.py index 495f6cf82..e9d3cf1bf 100644 --- a/python/cudnn/linear_attention/frost/kernel/gdn2_prefill_config.py +++ b/python/cudnn/linear_attention/frost/kernel/gdn2_prefill_config.py @@ -54,7 +54,7 @@ class Cfg: # --- SMEM / TMEM ring stage counts --- SMEM_RAW_STAGES: int = 5 - SMEM_SCHED_STAGES: int = 8 + SMEM_SCHEDULER_STAGES: int = 8 SMEM_O_STAGES: int = 2 SMEM_DECAY_STAGES: int = 2 SMEM_INTERMEDIATE_STAGES: int = 2 diff --git a/python/cudnn/linear_attention/frost/kernel/gdn2_prefill_f16.py b/python/cudnn/linear_attention/frost/kernel/gdn2_prefill_f16.py index 98f9ff763..b8646ad53 100644 --- a/python/cudnn/linear_attention/frost/kernel/gdn2_prefill_f16.py +++ b/python/cudnn/linear_attention/frost/kernel/gdn2_prefill_f16.py @@ -15,73 +15,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Chunked Gated DeltaNet v2 (GDN-2) prefill kernel for Blackwell SM100/SM103 -(Cutlass DSL), BT=16 tiling with per-key-channel decay + per-key erase gate -(beta) + per-value write gate (W), using direct CUTLASS primitives. -Framework-neutral entry ``chunk_gdn2_sm100``. - -Extends the KDA BT=16 schedule to the channel-wise gated delta rule: - - S_t = S_{t-1} * diag(exp(g_t)) - U_t = W_t * V_t - (beta_t * K_t)^T S_t - S_t += K_t (x) U_t - O_t = scale * Q_t^T S_t - -vs KDA: the erase gate beta and write gate W are per-channel tensors. beta -is folded into the k_decay operand (feeds KK^T and state*K), the strict-lower -tile loses its per-row beta scale, and Y becomes `W*V - state*K`; beta/W -arrive by TMA alongside Q/K/V. - -ABI: q `[T, HQ, DK]`, k `[T, HK, DK]`, v `[T, HV, DV]`, gate -`[T, HO, DK]` fp32 (natural-log decay unless SAFE_GATE), beta `[T, HO, DK]` -and w `[T, HO, DV]` in the io dtype, cu_seqlens int32, states/checkpoints -`[N, HO, DV, DK]` (VK, k contiguous). GQA/GVA head broadcast follows repeat_interleave: -source head = head_idx // (HO // H_x). State presence, L2NORM, SAFE_GATE, -checkpoints, and the head ratios are compile-time specializations. - -Warp assignments (16 warps = 512 threads): - warps 0-7 : compute group 0 - Gate prefix scan + decay/restore operands - warps 8-11 : compute group 1 - TMEM value side, O drain, state stores - warp 12 : super-MMA - register-MMA KK^T + Neumann inverse - warp 13 : tcgen05-MMA - the six state GEMMs + the TMEM lifecycle - warp 14 : TMA load - per-chunk input G->S loads - warp 15 : epilogue - register-MMA A + the O TMA store - -SMEM layout (~205 KB total): - Buffer Bytes Stages - Q / K / V raw 20480 5 <-- SW128 TMA ring (io dtype) - Beta / W raw 2x 20480 5 <-- per-channel gates, same ring - Gate raw 40960 5 <-- fp32 prefix-scan source - dt_bias (+a_log slot) 516 1 <-- SAFE_GATE only - K_inv 8192 2 <-- token-major ldmatrix/tcgen05 B operand - K decay / Q decay 2x 8192 2 <-- tcgen05 SW128 K-box-major A/B operands - K restore 8192 2 <-- tcgen05 B operand for the state update - state-scale diag 12288 3 <-- per-k-atom decay diagonal blocks - Intermediate (T_inv / A) 2048 2 <-- SW32 16x16 register-MMA tiles - O staging 8192 2 <-- W128 output drain - -TMEM layout (272 of 512 columns): - Buffer Cols Purpose - state 0-127 state[DK,DV] fp32 recurrent state - state inp 128-191 packed b16 A operand view of the state - q_state_acc 192-223 2-stage state*Q -> O accumulator - state_k_acc 224-239 state*K fp32 accumulator - u_acc 240-255 U fp32 accumulator - y_inp 256-263 packed b16 A operand: Y = W*V - state*(Beta*K) - u_inp 264-271 packed b16 A operand: the U readback - -GEMM schedule (tcgen05-MMA warp, in issue order per chunk): - state*K -> state_k_acc - state*Q -> q_state_acc (the O acc) - state decay (diag blocks) - U = Y(T) @ T_inv -> u_acc - final_state += U @ K_restore - O += A @ U -> q_state_acc - -Requires a cutlass DSL build providing `cutlass.experimental.*`; not available -in the pip nvidia-cutlass-dsl releases. -""" - from dataclasses import dataclass from functools import lru_cache from typing import Callable, NamedTuple, Optional, Type @@ -93,7 +26,7 @@ import cutlass.cute as cute from cutlass.cute.runtime import from_dlpack -from ..common.split_k import ORDER_CAPACITY, ORDER_ELEMS, ORDER_THREADS, decode_work_item, order_body +from ..common.split_k import ORDER_CAPACITY, ORDER_ELEMENTS, ORDER_THREADS, decode_work_item, order_body from ..common.host import get_dtype from cudnn.frost.buffers import data_ptr from cudnn.frost.device import current_device, multiprocessor_count @@ -108,9 +41,10 @@ ) from cudnn.frost.tile_dsl.handles import GmemTileTma, MmaDesc, SmemTile, tma_slice_runtime_desc from cudnn.frost.tile_dsl.mma import mma_step, mma_ts_step -from cudnn.frost.tile_dsl.swizzle import swizzle_lin_128b, swizzle_lin_S, swizzle_xor_128b +from cudnn.frost.tile_dsl.swizzle import swizzle_xor_128b, swizzle_xor_32b from cudnn.frost.tile_dsl.tma import tma_load_tile, tma_store_commit, tma_store_tile, tma_store_wait, tma_tensormap_acquire from cudnn.frost.tile_dsl.pointwise import ( + sigmoid, f16x2_to_f32, fadd2, fmul2, @@ -123,14 +57,7 @@ ) LOG2_E: float = 1.4426950408889634 - - DEFAULT_GATE_LOWER_BOUND: float = -5.0 - - -# Host-side API defaults. - - L2_NORM_EPS: float = 1.0e-12 @@ -156,9 +83,9 @@ class Gdn2Bars(NamedTuple): mb_state_k_acc_ready: MBarrier mb_u_acc_ready: MBarrier - mb_state_inp_ready: MBarrier - mb_y_inp_ready: MBarrier - mb_u_inp_ready: MBarrier + mb_state_input_ready: MBarrier + mb_y_input_ready: MBarrier + mb_u_input_ready: MBarrier mb_t_inv_ready: MBarrier mb_intermediate_done: MBarrier @@ -179,13 +106,12 @@ class Gdn2Bars(NamedTuple): mb_checkpoint_tmastg_ready: MBarrier mb_checkpoint_tmastg_done: MBarrier - mb_sched_ready: MBarrier - mb_sched_done: MBarrier + mb_scheduler_ready: MBarrier + mb_scheduler_done: MBarrier def make_gdn2_bars(cfg) -> Gdn2Bars: - """Bars factory. MUST be called from inside ``kernel`` (allocates the - mbarrier rings in SMEM ahead of the data buffers).""" + """Gdn2Bars factory.""" def alloc(n): return cutlass.Array(cutlass.Int64, n, space=cutlass.AddressSpace.smem, alignment=8) @@ -211,9 +137,9 @@ def alloc(n): mb_o_acc_done=MBarrier(alloc(cfg.tmem_q_state_acc_stages), stages=cfg.tmem_q_state_acc_stages, init_count=CG1_THREADS, producer=Producer.THREAD), mb_state_k_acc_ready=MBarrier(alloc(1), stages=1, init_count=1, producer=Producer.MMA_COMMIT), mb_u_acc_ready=MBarrier(alloc(1), stages=1, init_count=1, producer=Producer.MMA_COMMIT), - mb_state_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), - mb_y_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), - mb_u_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_state_input_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_y_input_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_u_input_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), mb_t_inv_ready=MBarrier(alloc(cfg.smem_intermediate_stages), stages=cfg.smem_intermediate_stages, init_count=WARP, producer=Producer.THREAD), mb_intermediate_done=MBarrier(alloc(cfg.smem_intermediate_stages), stages=cfg.smem_intermediate_stages, init_count=1, producer=Producer.MMA_COMMIT), mb_a_ready=MBarrier(alloc(cfg.smem_intermediate_stages), stages=cfg.smem_intermediate_stages, init_count=WARP, producer=Producer.THREAD), @@ -241,40 +167,40 @@ def alloc(n): alloc(cfg.smem_checkpoint_stages), stages=cfg.smem_checkpoint_stages, init_count=CG1_THREADS, producer=Producer.THREAD ), mb_checkpoint_tmastg_done=MBarrier(alloc(cfg.smem_checkpoint_stages), stages=cfg.smem_checkpoint_stages, init_count=WARP, producer=Producer.THREAD), - mb_sched_ready=MBarrier(alloc(cfg.sched_stages), stages=cfg.sched_stages, init_count=1, producer=Producer.THREAD), - mb_sched_done=MBarrier(alloc(cfg.sched_stages), stages=cfg.sched_stages, init_count=15, producer=Producer.THREAD), + mb_scheduler_ready=MBarrier(alloc(cfg.scheduler_stages), stages=cfg.scheduler_stages, init_count=1, producer=Producer.THREAD), + mb_scheduler_done=MBarrier(alloc(cfg.scheduler_stages), stages=cfg.scheduler_stages, init_count=15, producer=Producer.THREAD), ) -# ---- Dynamic tile scheduler ------------------------------------------------------ +# ---- Dynamic tile scheduler ---------------------------------------------------------- @cute.jit -def sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ctas): +def scheduler_publish_next(cfg, bars, sScheduler, mScheduler, scheduler_state, tile_idx, num_ctas, elect_one): """TMA-warp side: pull the next tile off the global ticket, publish it.""" - if cutlass.const_expr(cfg.dyn_sched): - bars.mb_sched_done[sched_state.idx].wait(sched_state.phase) - if nvvm.elect_sync(): - fetched = cutlass.Int32(nvvm.atomicrmw("add", mSched.iterator, cutlass.Int32(1), mem_order="relaxed", syncscope="gpu")) - sSched[sched_state.idx] = num_ctas + fetched + if cutlass.const_expr(cfg.dynamic_scheduling): + bars.mb_scheduler_done[scheduler_state.idx].wait(scheduler_state.phase) + if elect_one: + fetched = cutlass.Int32(nvvm.atomicrmw("add", mScheduler.iterator, cutlass.Int32(1), mem_order="relaxed", syncscope="gpu")) + sScheduler[scheduler_state.idx] = num_ctas + fetched nvvm.bar_warp_sync(cute.arch.FULL_MASK) - next_tile = sSched[sched_state.idx] - if nvvm.elect_sync(): - bars.mb_sched_ready[sched_state.idx].arrive() - return next_tile, advance(sched_state, cfg.sched_stages) - return tile_idx + num_ctas, sched_state + next_tile = sScheduler[scheduler_state.idx] + if elect_one: + bars.mb_scheduler_ready[scheduler_state.idx].arrive() + return next_tile, advance(scheduler_state, cfg.scheduler_stages) + return tile_idx + num_ctas, scheduler_state @cute.jit -def sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas): +def scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one): """Consumer side: read the TMA warp's published next tile.""" - if cutlass.const_expr(cfg.dyn_sched): - bars.mb_sched_ready[sched_state.idx].wait(sched_state.phase) - next_tile = sSched[sched_state.idx] - if nvvm.elect_sync(): - bars.mb_sched_done[sched_state.idx].arrive() - return next_tile, advance(sched_state, cfg.sched_stages) - return tile_idx + num_ctas, sched_state + if cutlass.const_expr(cfg.dynamic_scheduling): + bars.mb_scheduler_ready[scheduler_state.idx].wait(scheduler_state.phase) + next_tile = sScheduler[scheduler_state.idx] + if elect_one: + bars.mb_scheduler_done[scheduler_state.idx].arrive() + return next_tile, advance(scheduler_state, cfg.scheduler_stages) + return tile_idx + num_ctas, scheduler_state @cute.jit @@ -285,9 +211,9 @@ def tmaldg_warp( num_ctas, cu_seqlens, mWorkItems, - mSched, - sSched, - lane, + mScheduler, + sScheduler, + lane_idx, sBeta_raw, sGate_raw, sK_raw, @@ -304,8 +230,13 @@ def tmaldg_warp( ) -> None: """TMA-LDG warp role (warp 14): persistent scheduler loop issuing the per-chunk Q/K/V/Beta/W/Gate G->S loads.""" - elect_one = nvvm.elect_sync() nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) + + raw_index = PipelineState.start(phase=1) + raw_bar_index = PipelineState.start(phase=0) + scheduler_state = PipelineState.start(phase=1) + + elect_one = nvvm.elect_sync() sQ_tma = SmemTile( base=sQ_raw, elems_per_stage=(cfg.d_k * cfg.b_t), @@ -372,12 +303,11 @@ def tmaldg_warp( tma_granu_elems=32, tma_subtile_stride_elems=(cfg.b_t * 32), ) - raw_index = PipelineState.start(phase=1) - raw_bar_index = PipelineState.start(phase=0) - sched_state = PipelineState.start(phase=1) tile_idx = cutlass.Int32(bidx) while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) head_o = head_idx head_q = head_idx if cfg.q_ratio == 1 else head_idx // cutlass.Int32(cfg.q_ratio) head_k = head_idx if cfg.k_ratio == 1 else head_idx // cutlass.Int32(cfg.k_ratio) @@ -396,45 +326,45 @@ def tmaldg_warp( tma_tensormap_acquire(desc_gate_slot) tma_tensormap_acquire(desc_beta_slot) tma_tensormap_acquire(desc_w_slot) - for chunk_idx in cutlass.range(cstart, wend, 1, unroll=1): + for chunk_idx in cutlass.range(compute_start, write_end, 1, unroll=1): chunk_start = chunk_idx * cfg.b_t - # ---- Q load ---------------------------------------------------------- + # ---- Q load -------------------------------------------------------------- bars.mb_q_done[raw_index.idx].wait(raw_index.phase) if elect_one: bars.mb_q_ready[raw_bar_index.idx].arrive(n_bytes=cfg.tma_q_bytes) q_slice = tma_slice_runtime_desc(desc_q_slot, cutlass.Int32(0), head_q, chunk_start) tma_load_tile(sQ_tma[raw_index.idx], q_slice, bars.mb_q_ready[raw_bar_index.idx].smem_ptr, acquire=False) - # ---- K load ---------------------------------------------------------- + # ---- K load -------------------------------------------------------------- bars.mb_k_done[raw_index.idx].wait(raw_index.phase) if elect_one: bars.mb_k_ready[raw_bar_index.idx].arrive(n_bytes=cfg.tma_k_bytes) k_slice = tma_slice_runtime_desc(desc_k_slot, cutlass.Int32(0), head_k, chunk_start) tma_load_tile(sK_tma[raw_index.idx], k_slice, bars.mb_k_ready[raw_bar_index.idx].smem_ptr, acquire=False) - # ---- V load ---------------------------------------------------------- + # ---- V load -------------------------------------------------------------- bars.mb_v_done[raw_index.idx].wait(raw_index.phase) if elect_one: bars.mb_v_ready[raw_bar_index.idx].arrive(n_bytes=cfg.tma_v_bytes) v_slice = tma_slice_runtime_desc(desc_v_slot, cutlass.Int32(0), head_v, chunk_start) tma_load_tile(sV_tma[raw_index.idx], v_slice, bars.mb_v_ready[raw_bar_index.idx].smem_ptr, acquire=False) - # ---- Beta load ------------------------------------------------------- + # ---- Beta load: GMEM -> SMEM --------------------------------------------- bars.mb_beta_done[raw_index.idx].wait(raw_index.phase) if elect_one: bars.mb_beta_ready[raw_bar_index.idx].arrive(n_bytes=cfg.tma_beta_bytes) beta_slice = tma_slice_runtime_desc(desc_beta_slot, cutlass.Int32(0), head_o, chunk_start) tma_load_tile(sBeta_tma[raw_index.idx], beta_slice, bars.mb_beta_ready[raw_bar_index.idx].smem_ptr, acquire=False) - # ---- W load ---------------------------------------------------------- + # ---- W load -------------------------------------------------------------- bars.mb_w_done[raw_index.idx].wait(raw_index.phase) if elect_one: bars.mb_w_ready[raw_bar_index.idx].arrive(n_bytes=cfg.tma_w_bytes) w_slice = tma_slice_runtime_desc(desc_w_slot, cutlass.Int32(0), head_o, chunk_start) tma_load_tile(sW_tma[raw_index.idx], w_slice, bars.mb_w_ready[raw_bar_index.idx].smem_ptr, acquire=False) - # ---- Gate load ------------------------------------------------------- + # ---- Gate load: GMEM -> SMEM --------------------------------------------- bars.mb_gate_done[raw_index.idx].wait(raw_index.phase) if elect_one: bars.mb_gate_ready[raw_bar_index.idx].arrive(n_bytes=cfg.tma_gate_bytes) @@ -442,7 +372,7 @@ def tmaldg_warp( tma_load_tile(sGate_tma[raw_index.idx], gate_slice, bars.mb_gate_ready[raw_bar_index.idx].smem_ptr, acquire=False) raw_index = advance(raw_index, cfg.smem_raw_stages) raw_bar_index = advance(raw_bar_index, cfg.smem_raw_bar_stages) - tile_idx, sched_state = sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_publish_next(cfg, bars, sScheduler, mScheduler, scheduler_state, tile_idx, num_ctas, elect_one) @cute.jit @@ -453,8 +383,8 @@ def super_mma_warp( num_ctas, cu_seqlens, mWorkItems, - sSched, - lane, + sScheduler, + lane_idx, sK_inv_raw, sIntermediate_raw, sK_decay_raw, @@ -463,26 +393,28 @@ def super_mma_warp( """Super-MMA warp role (warp 12): persistent scheduler loop computing the Neumann-series T_inv by register MMA.""" nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) + elect_one = nvvm.elect_sync() - # ---- ldmatrix/stmatrix lane decode ------------------------------------------- - rhs_row_coord = lane % 8 + (cutlass.Int32(8) if (lane // 16) else cutlass.Int32(0)) - rhs_col_offset = cutlass.Int32(8) if ((lane // 8) % 2) else cutlass.Int32(0) - lhs_row_coord = lane % 8 + (cutlass.Int32(8) if ((lane // 8) % 2) else cutlass.Int32(0)) - lhs_col_offset = cutlass.Int32(8) if ((lane // 8) // 2) else cutlass.Int32(0) - decay_key_mask = cutlass.Int32(8) - stsm_row_coord = lane & 7 - stsm_col_coord = cutlass.Int32(0) - if (lane // 8) & 1: - stsm_row_coord = stsm_row_coord + cutlass.Int32(8) - if lane // 8 >= 2: - stsm_col_coord = cutlass.Int32(8) - stsm_idx = swizzle_lin_S(stsm_row_coord * cfg.b_t + (stsm_col_coord ^ (cfg.b_t // 2)), bbits=1, mbase=3, sshift=3) + # ---- ldmatrix/stmatrix lane decode ----------------------------------------------- + k_inv_row_coord = lane_idx % 8 + (cutlass.Int32(8) if (lane_idx // 16) else cutlass.Int32(0)) + k_inv_col_offset = cutlass.Int32(8) if ((lane_idx // 8) % 2) else cutlass.Int32(0) + k_decay_row_coord = lane_idx % 8 + (cutlass.Int32(8) if ((lane_idx // 8) % 2) else cutlass.Int32(0)) + k_decay_col_offset = cutlass.Int32(8) if ((lane_idx // 8) // 2) else cutlass.Int32(0) + t_inv_row_coord = lane_idx & 7 + t_inv_col_coord = cutlass.Int32(0) + if (lane_idx // 8) & 1: + t_inv_row_coord = t_inv_row_coord + cutlass.Int32(8) + if lane_idx // 8 >= 2: + t_inv_col_coord = cutlass.Int32(8) + t_inv_idx = t_inv_row_coord * cfg.b_t + swizzle_xor_32b(t_inv_row_coord, t_inv_col_coord) global_chunk_base = cutlass.Int32(0) - sched_state = PipelineState.start(phase=0) + scheduler_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) - num_tile_chunks = wend - cstart # processed chunks; ring bookkeeping runs on global_chunk_base + local_chunk + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) + num_tile_chunks = write_end - compute_start for local_chunk in cutlass.range(num_tile_chunks, unroll=1): global_chunk = global_chunk_base + local_chunk decay_stage = global_chunk % cfg.smem_decay_stages @@ -493,53 +425,51 @@ def super_mma_warp( bars.mb_k_decay_inv_cg0_ready[decay_stage].wait((global_chunk // cfg.smem_decay_stages) % 2) - # ---- KK = K_decay @ K_inv^T ------------------------------------------ + # ---- KK = K decay @ K inv^T ---------------------------------------------- kk_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): kk_acc[accum_idx] = cutlass.Float32(0.0) - for k_block in cutlass.range_constexpr((cfg.d_k // 16)): - # Load B operand - k_inv_col = k_block * 16 + rhs_col_offset + for i in cutlass.range_constexpr((cfg.d_k // 16)): + k_inv_col = i * 16 + k_inv_col_offset k_inv_segment = k_inv_col // 64 - rhs_frag = nvvm.ldmatrix( + k_inv_frag = nvvm.ldmatrix( sK_inv_ptr + k_inv_segment * (cfg.b_t * 64) - + rhs_row_coord * 64 - + swizzle_xor_128b(rhs_row_coord, k_inv_col - k_inv_segment * 64, elem_bytes=2), + + k_inv_row_coord * 64 + + swizzle_xor_128b(k_inv_row_coord, k_inv_col - k_inv_segment * 64, elem_bytes=2), 4, nvvm.MMALayout.ROW, ) - # Load A operand - storage_key = (k_block * 16 + lhs_col_offset) ^ decay_key_mask - storage_slice = storage_key // 64 - kk_lhs_frag = nvvm.ldmatrix( + k_decay_col = i * 16 + k_decay_col_offset + k_decay_segment = k_decay_col // 64 + k_decay_frag = nvvm.ldmatrix( sK_decay_ptr - + storage_slice * (cfg.b_t * 64) - + swizzle_xor_128b(lhs_row_coord, lhs_row_coord * 64 + storage_key - storage_slice * 64, elem_bytes=2), + + k_decay_segment * (cfg.b_t * 64) + + swizzle_xor_128b(k_decay_row_coord, k_decay_row_coord * 64 + k_decay_col - k_decay_segment * 64, elem_bytes=2), 4, nvvm.MMALayout.ROW, ) mma_step( kk_acc, - (kk_lhs_frag[0], kk_lhs_frag[1], kk_lhs_frag[2], kk_lhs_frag[3]), - (rhs_frag[0], rhs_frag[1], rhs_frag[2], rhs_frag[3]), + (k_decay_frag[0], k_decay_frag[1], k_decay_frag[2], k_decay_frag[3]), + (k_inv_frag[0], k_inv_frag[1], k_inv_frag[2], k_inv_frag[3]), k_step=0, M=16, N=16, ab_dtype=cfg.io_dtype, ) - # ---- L = tril(KK, -1) fragment --------------------------------------- - row_lo = lane // 4 + # ---- L = tril(KK, -1) ---------------------------------------------------- + row_lo = lane_idx // 4 row_hi = row_lo + cutlass.Int32(8) l_regs = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): row_coord = row_lo if cutlass.const_expr(accum_idx % 4 >= 2): row_coord = row_hi - col_coord = (accum_idx // 4) * 8 + 2 * (lane % 4) + col_coord = (accum_idx // 4) * 8 + 2 * (lane_idx % 4) if cutlass.const_expr(accum_idx % 2 == 1): col_coord = col_coord + cutlass.Int32(1) l_regs[accum_idx] = kk_acc[accum_idx] if row_coord > col_coord else cutlass.Float32(0.0) @@ -549,13 +479,13 @@ def super_mma_warp( l_a3 = fp32_to_fp16(l_regs[6], l_regs[7], dtype=cfg.io_dtype) l_values = cutlass.Vector.from_elements((l_a0, l_a1, l_a2, l_a3), cutlass.Int32).bitcast(cfg.io_dtype).to(cutlass.Float32) - # ---- T_inv = I - L, then three Neumann doubling rounds --------------- + # ---- T^-1 = I - L, then three Neumann doubling rounds -------------------- tinv_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): row_coord = row_lo if cutlass.const_expr(accum_idx % 4 >= 2): row_coord = row_hi - col_coord = (accum_idx // 4) * 8 + 2 * (lane % 4) + col_coord = (accum_idx // 4) * 8 + 2 * (lane_idx % 4) if cutlass.const_expr(accum_idx % 2 == 1): col_coord = col_coord + cutlass.Int32(1) eye = cutlass.Float32(1.0) if row_coord == col_coord else cutlass.Float32(0.0) @@ -563,8 +493,8 @@ def super_mma_warp( lpow_a0, lpow_a1, lpow_a2, lpow_a3 = l_a0, l_a1, l_a2, l_a3 mov_lpow0, mov_lpow1, mov_lpow2, mov_lpow3 = movmatrix_16b(l_a0), movmatrix_16b(l_a1), movmatrix_16b(l_a2), movmatrix_16b(l_a3) - for _round in cutlass.range_constexpr(3): - # ---- Lpow = Lpow @ Lpow ------------------------------------------ + for neumann_round in cutlass.range_constexpr(3): + # ---- Lpow = Lpow @ Lpow ---------------------------------------------- sq_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): sq_acc[accum_idx] = cutlass.Float32(0.0) @@ -582,7 +512,7 @@ def super_mma_warp( lpow_a2 = fp32_to_fp16(sq_acc[4], sq_acc[5], dtype=cfg.io_dtype) lpow_a3 = fp32_to_fp16(sq_acc[6], sq_acc[7], dtype=cfg.io_dtype) mov_lpow0, mov_lpow1, mov_lpow2, mov_lpow3 = movmatrix_16b(lpow_a0), movmatrix_16b(lpow_a1), movmatrix_16b(lpow_a2), movmatrix_16b(lpow_a3) - # ---- T_inv += T_inv @ Lpow --------------------------------------- + # ---- T^-1 += T^-1 @ Lpow --------------------------------------------- upd_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): upd_acc[accum_idx] = cutlass.Float32(0.0) @@ -610,7 +540,7 @@ def super_mma_warp( bars.mb_intermediate_done[intermediate_stage].wait(((global_chunk // cfg.smem_intermediate_stages) + 1) % 2) nvvm.stmatrix( - sIntermediate_ptr + (cfg.b_t * cfg.b_t) + stsm_idx, + sIntermediate_ptr + (cfg.b_t * cfg.b_t) + t_inv_idx, [ fp32_to_fp16(tinv_acc[0], tinv_acc[1], dtype=cfg.io_dtype), fp32_to_fp16(tinv_acc[2], tinv_acc[3], dtype=cfg.io_dtype), @@ -624,7 +554,7 @@ def super_mma_warp( bars.mb_t_inv_ready[intermediate_stage].arrive() bars.mb_decay_super_done[decay_stage].arrive() global_chunk_base += num_tile_chunks - tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one) @cute.jit @@ -635,11 +565,11 @@ def tcgen05_mma_warp( num_ctas, cu_seqlens, mWorkItems, - sSched, + sScheduler, sTmem_base, sIntermediate, sK_decay, - sK_restore, + sK_restore_trans, sQ_decay, sState_scale_diag, bars, @@ -651,20 +581,20 @@ def tcgen05_mma_warp( nvvm.tcgen05_alloc(sTmem_base, cutlass.Int32(512), group=nvvm.CTAGroup.CTA_1) nvvm.barrier_cta_sync(cfg.tmem_lifecycle_barrier_id, thread_count=cfg.tmem_user_threads) tmem_base = sTmem_base.load() - state_inp_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_state_inp_offset, cutlass.Int8) + state_input_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_state_input_offset, cutlass.Int8) state_dsts = tuple(nvvm.make_tmem_ptr(tmem_base + cfg.tmem_state_acc_offset + k * 16, cutlass.Float32) for k in range(cfg.d_k // 16)) state_k_acc_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_state_k_acc_offset, cutlass.Float32) u_acc_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_u_acc_offset, cutlass.Float32) - y_inp_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_y_inp_offset, cutlass.Int8) - u_inp_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_u_inp_offset, cutlass.Int8) + y_input_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_y_input_offset, cutlass.Int8) + u_input_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_u_input_offset, cutlass.Int8) state_dst_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_state_acc_offset, cutlass.Float32) - state_inp_index = PipelineState.start(phase=0) + state_input_index = PipelineState.start(phase=0) state_read_index = PipelineState.start(phase=0) - y_inp_index = PipelineState.start(phase=0) - u_inp_index = PipelineState.start(phase=0) + y_input_index = PipelineState.start(phase=0) + u_input_index = PipelineState.start(phase=0) qk_scale_index = PipelineState.start(phase=0) - # ---- chunk-invariant GEMM descriptors ---------------------------------------- + # ---- chunk-invariant GEMM descriptors -------------------------------------------- bpe = cfg.io_dtype.width // 8 idesc_acc = nvvm.Tcgen05InstrDesc.build( c_dtype=cutlass.Float32, @@ -690,7 +620,7 @@ def tcgen05_mma_warp( m_dim=cfg.d_v, b_major=1, ) - bmm_state_desc = MmaDesc( + bmm_state_k_decay_desc = MmaDesc( M=cfg.d_v, N=cfg.b_t, K=cfg.d_k, @@ -702,7 +632,8 @@ def tcgen05_mma_warp( idesc=idesc_acc, kind=nvvm.Tcgen05MMAKind.F16, ) - bmm_diag_desc = MmaDesc( + bmm_state_q_decay_desc = bmm_state_k_decay_desc + bmm_state_diag_desc = MmaDesc( M=cfg.d_v, N=16, K=16, @@ -714,7 +645,7 @@ def tcgen05_mma_warp( idesc=idesc_diag, kind=nvvm.Tcgen05MMAKind.F16, ) - bmm_qk_desc = MmaDesc( + bmm_u_a_desc = MmaDesc( M=cfg.d_v, N=cfg.b_t, K=cfg.b_t, @@ -726,7 +657,8 @@ def tcgen05_mma_warp( idesc=idesc_acc, kind=nvvm.Tcgen05MMAKind.F16, ) - bmm_final_state_desc = MmaDesc( + bmm_y_t_inv_desc = bmm_u_a_desc + bmm_u_k_restore_desc = MmaDesc( M=cfg.d_v, N=cfg.d_k, K=cfg.b_t, @@ -738,39 +670,46 @@ def tcgen05_mma_warp( idesc=idesc_final_state, kind=nvvm.Tcgen05MMAKind.F16, ) - STATE_A_SEG = bmm_state_desc.sps_B * bmm_state_desc.tmem_advance_A - STATE_B_SEG = bmm_state_desc.smem_subtile_B >> 4 + STATE_A_SEG = bmm_state_k_decay_desc.sps_B * bmm_state_k_decay_desc.tmem_advance_A + STATE_B_SEG = bmm_state_k_decay_desc.smem_subtile_B >> 4 global_chunk_base = cutlass.Int32(0) - sched_state = PipelineState.start(phase=0) + scheduler_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) - num_tile_chunks = wend - cstart + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) + num_tile_chunks = write_end - compute_start + if cutlass.const_expr(cfg.use_initial_state): + seed_state = compute_start == 0 for local_chunk in cutlass.range(num_tile_chunks, unroll=1): global_chunk = global_chunk_base + local_chunk - have_state = cutlass.Boolean(True) if cutlass.const_expr(cfg.use_initial_state) else local_chunk > 0 + if cutlass.const_expr(cfg.use_initial_state): + have_state = local_chunk > 0 or seed_state + else: + have_state = local_chunk > 0 q_state_acc_stage = global_chunk % cfg.tmem_q_state_acc_stages decay_stage = global_chunk % cfg.smem_decay_stages state_scale_diag_stage = qk_scale_index.idx intermediate_stage = global_chunk % cfg.smem_intermediate_stages sK_decay_stage = sK_decay[decay_stage] sQ_decay_stage = sQ_decay[decay_stage] - sK_restore_stage = sK_restore[decay_stage] + sK_restore_stage = sK_restore_trans[decay_stage] sState_scale_diag_stage = sState_scale_diag[state_scale_diag_stage] sIntermediate_stage = sIntermediate[intermediate_stage] - # ---- state_k = state(T) @ K_decay^T ---------------------------------- + # ---- k state = state(T) @ K decay^T -------------------------------------- bars.mb_k_decay_inv_cg0_ready[decay_stage].wait((global_chunk // cfg.smem_decay_stages) % 2) if have_state: - bars.mb_state_inp_ready.wait(state_inp_index.phase) - state_inp_index = advance(state_inp_index, 1) + bars.mb_state_input_ready.wait(state_input_index.phase) + state_input_index = advance(state_input_index, 1) desc_k_decay = sK_decay_stage.desc() - for s in cutlass.range_constexpr(bmm_state_desc.num_subtiles_B): - for k in cutlass.range_constexpr(bmm_state_desc.sps_B): + for s in cutlass.range_constexpr(bmm_state_k_decay_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_state_k_decay_desc.sps_B): mma_ts_step( - bmm_state_desc, - state_inp_ptr.subview(s * STATE_A_SEG), + bmm_state_k_decay_desc, + state_input_ptr.subview(s * STATE_A_SEG), desc_k_decay + s * STATE_B_SEG, state_k_acc_ptr, k, @@ -780,17 +719,17 @@ def tcgen05_mma_warp( if elect_one: bars.mb_state_k_acc_ready.arrive(cta_group=1) - # ---- q_state = state(T) @ Q_decay^T ---------------------------------- + # ---- q state = state(T) @ Q decay^T -------------------------------------- bars.mb_qk_scale_ready[qk_scale_index.idx].wait(qk_scale_index.phase) bars.mb_o_acc_done[q_state_acc_stage].wait(((global_chunk // cfg.tmem_q_state_acc_stages + cutlass.Int32(1)) % cutlass.Int32(2))) q_state_acc_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_q_state_acc_offset + q_state_acc_stage * cfg.b_t, cutlass.Float32) if have_state: desc_q_decay = sQ_decay_stage.desc() - for s in cutlass.range_constexpr(bmm_state_desc.num_subtiles_B): - for k in cutlass.range_constexpr(bmm_state_desc.sps_B): + for s in cutlass.range_constexpr(bmm_state_q_decay_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_state_q_decay_desc.sps_B): mma_ts_step( - bmm_state_desc, - state_inp_ptr.subview(s * STATE_A_SEG), + bmm_state_q_decay_desc, + state_input_ptr.subview(s * STATE_A_SEG), desc_q_decay + s * STATE_B_SEG, q_state_acc_ptr, k, @@ -805,15 +744,15 @@ def tcgen05_mma_warp( bars.mb_state_acc_read_done.wait(state_read_index.phase) state_read_index = advance(state_read_index, 1) - # ---- state decay = state(T) @ diag(exp2(g_last)) (per-k-atom blocks) ---- + # ---- state decay = state(T) @ diag(exp2(g last)) (per-k-atom blocks) ----- if have_state: desc_diag = sState_scale_diag_stage.desc() - for k_block in cutlass.range_constexpr(cfg.d_k // 16): + for i in cutlass.range_constexpr(cfg.d_k // 16): mma_ts_step( - bmm_diag_desc, - state_inp_ptr.subview(k_block * bmm_diag_desc.tmem_advance_A), - desc_diag.advance_start_address(k_block * 256 * 2), - state_dsts[k_block], + bmm_state_diag_desc, + state_input_ptr.subview(i * bmm_state_diag_desc.tmem_advance_A), + desc_diag.advance_start_address(i * 256 * 2), + state_dsts[i], 0, cutlass.Boolean(False), ) @@ -821,30 +760,30 @@ def tcgen05_mma_warp( if elect_one: bars.mb_state_scale_diag_done[state_scale_diag_stage].arrive(cta_group=1) - # ---- u_acc = Y(T) @ T_inv -------------------------------------------- + # ---- u acc = Y(T) @ T^-1 ------------------------------------------------- bars.mb_t_inv_ready[intermediate_stage].wait((global_chunk // cfg.smem_intermediate_stages) % 2) - bars.mb_y_inp_ready.wait(y_inp_index.phase) - y_inp_index = advance(y_inp_index, 1) + bars.mb_y_input_ready.wait(y_input_index.phase) + y_input_index = advance(y_input_index, 1) desc_qk = sIntermediate_stage.shifted((cfg.b_t * cfg.b_t)).desc() - mma_ts_step(bmm_qk_desc, y_inp_ptr, desc_qk, u_acc_ptr, 0, cutlass.Boolean(False)) + mma_ts_step(bmm_y_t_inv_desc, y_input_ptr, desc_qk, u_acc_ptr, 0, cutlass.Boolean(False)) if elect_one: bars.mb_u_acc_ready.arrive(cta_group=1) - # ---- final_state += U(T) @ K_restore --------------------------------- - bars.mb_u_inp_ready.wait(u_inp_index.phase) - u_inp_index = advance(u_inp_index, 1) + # ---- final state += U(T) @ K restore ------------------------------------- + bars.mb_u_input_ready.wait(u_input_index.phase) + u_input_index = advance(u_input_index, 1) desc_k_restore = sK_restore_stage.desc() - mma_ts_step(bmm_final_state_desc, u_inp_ptr, desc_k_restore, state_dst_ptr, 0, have_state) + mma_ts_step(bmm_u_k_restore_desc, u_input_ptr, desc_k_restore, state_dst_ptr, 0, have_state) if elect_one: bars.mb_k_restore_acc_done[decay_stage].arrive(cta_group=1) - # ---- O += U(T) @ A --------------------------------------------------- + # ---- O += U(T) @ A ------------------------------------------------------- bars.mb_a_ready[intermediate_stage].wait((global_chunk // cfg.smem_intermediate_stages) % 2) desc_qk = sIntermediate_stage.desc() mma_ts_step( - bmm_qk_desc, - u_inp_ptr, + bmm_u_a_desc, + u_input_ptr, desc_qk, nvvm.make_tmem_ptr(tmem_base + cfg.tmem_q_state_acc_offset + q_state_acc_stage * cfg.b_t, cutlass.Float32), 0, @@ -856,7 +795,7 @@ def tcgen05_mma_warp( qk_scale_index = advance(qk_scale_index, cfg.smem_state_scale_diag_stages) global_chunk_base += num_tile_chunks - tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one) bars.mb_tmem_done[0].wait(0) nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_1) nvvm.tcgen05_dealloc( @@ -874,8 +813,8 @@ def epilogue_warp( num_ctas, cu_seqlens, mWorkItems, - sSched, - lane, + sScheduler, + lane_idx, mO, sK_inv_raw, sO_raw, @@ -917,39 +856,40 @@ def epilogue_warp( ) qk_scale_index = PipelineState.start(phase=0) - # ---- ldmatrix/stmatrix lane decode ------------------------------------------- - rhs_row_coord = lane % 8 + (cutlass.Int32(8) if (lane // 16) else cutlass.Int32(0)) - rhs_col_offset = cutlass.Int32(8) if ((lane // 8) % 2) else cutlass.Int32(0) - lhs_row_coord = lane % 8 + (cutlass.Int32(8) if ((lane // 8) % 2) else cutlass.Int32(0)) - lhs_col_offset = cutlass.Int32(8) if ((lane // 8) // 2) else cutlass.Int32(0) - decay_key_mask = cutlass.Int32(8) - stsm_row_coord = lane & 7 - stsm_col_coord = cutlass.Int32(0) - if (lane // 8) & 1: - stsm_row_coord = stsm_row_coord + cutlass.Int32(8) - if lane // 8 >= 2: - stsm_col_coord = cutlass.Int32(8) - stsm_idx = swizzle_lin_S(stsm_row_coord * cfg.b_t + (stsm_col_coord ^ (cfg.b_t // 2)), bbits=1, mbase=3, sshift=3) + # ---- ldmatrix/stmatrix lane decode ----------------------------------------------- + k_inv_row_coord = lane_idx % 8 + (cutlass.Int32(8) if (lane_idx // 16) else cutlass.Int32(0)) + k_inv_col_offset = cutlass.Int32(8) if ((lane_idx // 8) % 2) else cutlass.Int32(0) + q_decay_row_coord = lane_idx % 8 + (cutlass.Int32(8) if ((lane_idx // 8) % 2) else cutlass.Int32(0)) + q_decay_col_offset = cutlass.Int32(8) if ((lane_idx // 8) // 2) else cutlass.Int32(0) + a_row_coord = lane_idx & 7 + a_col_coord = cutlass.Int32(0) + if (lane_idx // 8) & 1: + a_row_coord = a_row_coord + cutlass.Int32(8) + if lane_idx // 8 >= 2: + a_col_coord = cutlass.Int32(8) + a_idx = a_row_coord * cfg.b_t + swizzle_xor_32b(a_row_coord, a_col_coord) global_chunk_base = cutlass.Int32(0) - sched_state = PipelineState.start(phase=0) + scheduler_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) head_o = head_idx o_slot = batch_idx * cutlass.Int32(TENSOR_MAP_QWORDS) desc_o_slot = (desc_o_base + o_slot).tospace(cutlass.AddressSpace.generic) if cutlass.const_expr(cfg.enable_checkpoints): desc_checkpoint_slot = (desc_checkpoint_base + o_slot).tospace(cutlass.AddressSpace.generic) checkpoint_chunks = checkpoint_every_n_tokens // cutlass.Int32(cfg.b_t) - checkpoint_quot = (cstart + cutlass.Int32(1)) // checkpoint_chunks - checkpoint_mod = (cstart + cutlass.Int32(1)) % checkpoint_chunks + checkpoint_quot = (compute_start + cutlass.Int32(1)) // checkpoint_chunks + checkpoint_mod = (compute_start + cutlass.Int32(1)) % checkpoint_chunks if elect_one: tma_tensormap_acquire(desc_checkpoint_slot) if elect_one: tma_tensormap_acquire(desc_o_slot) - num_tile_chunks = wend - cstart + num_tile_chunks = write_end - compute_start if cutlass.const_expr(cfg.enable_checkpoints): - if num_tile_chunks > 0 and wstart == 0: + if num_tile_chunks > 0 and write_start == 0: checkpoint_stage = checkpoint_ready_index.idx bars.mb_checkpoint_tmastg_ready[checkpoint_stage].wait(checkpoint_ready_index.phase) checkpoint_ready_index = advance(checkpoint_ready_index, cfg.smem_checkpoint_stages) @@ -959,7 +899,7 @@ def epilogue_warp( tma_store_wait(0) bars.mb_checkpoint_tmastg_done[checkpoint_stage].arrive() for local_chunk in cutlass.range(num_tile_chunks, unroll=1): - chunk_idx = cstart + local_chunk + chunk_idx = compute_start + local_chunk global_chunk = global_chunk_base + local_chunk decay_stage = global_chunk % cfg.smem_decay_stages intermediate_stage = global_chunk % cfg.smem_intermediate_stages @@ -969,38 +909,36 @@ def epilogue_warp( bars.mb_qk_scale_ready[qk_scale_index.idx].wait(qk_scale_index.phase) - # ---- A = Q_decay @ K_inv^T ------------------------------------------ + # ---- A = Q decay @ K inv^T ----------------------------------------------- a_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): a_acc[accum_idx] = cutlass.Float32(0.0) - for k_block in cutlass.range_constexpr((cfg.d_k // 16)): - # Load B operand - k_inv_col = k_block * 16 + rhs_col_offset + for i in cutlass.range_constexpr((cfg.d_k // 16)): + k_inv_col = i * 16 + k_inv_col_offset k_inv_segment = k_inv_col // 64 - rhs_frag = nvvm.ldmatrix( + k_inv_frag = nvvm.ldmatrix( sK_inv_ptr + k_inv_segment * (cfg.b_t * 64) - + rhs_row_coord * 64 - + swizzle_xor_128b(rhs_row_coord, k_inv_col - k_inv_segment * 64, elem_bytes=2), + + k_inv_row_coord * 64 + + swizzle_xor_128b(k_inv_row_coord, k_inv_col - k_inv_segment * 64, elem_bytes=2), 4, nvvm.MMALayout.ROW, ) - # Load A operand - storage_key = (k_block * 16 + lhs_col_offset) ^ decay_key_mask - storage_slice = storage_key // 64 - a_lhs_frag = nvvm.ldmatrix( + q_decay_col = i * 16 + q_decay_col_offset + q_decay_segment = q_decay_col // 64 + q_decay_frag = nvvm.ldmatrix( sQ_decay_ptr - + storage_slice * (cfg.b_t * 64) - + swizzle_xor_128b(lhs_row_coord, lhs_row_coord * 64 + storage_key - storage_slice * 64, elem_bytes=2), + + q_decay_segment * (cfg.b_t * 64) + + swizzle_xor_128b(q_decay_row_coord, q_decay_row_coord * 64 + q_decay_col - q_decay_segment * 64, elem_bytes=2), 4, nvvm.MMALayout.ROW, ) mma_step( a_acc, - (a_lhs_frag[0], a_lhs_frag[1], a_lhs_frag[2], a_lhs_frag[3]), - (rhs_frag[0], rhs_frag[1], rhs_frag[2], rhs_frag[3]), + (q_decay_frag[0], q_decay_frag[1], q_decay_frag[2], q_decay_frag[3]), + (k_inv_frag[0], k_inv_frag[1], k_inv_frag[2], k_inv_frag[3]), k_step=0, M=16, N=16, @@ -1008,17 +946,17 @@ def epilogue_warp( ) for accum_idx in cutlass.range_constexpr(8): - row_coord = lane // 4 + row_coord = lane_idx // 4 if cutlass.const_expr(accum_idx % 4 >= 2): row_coord = row_coord + cutlass.Int32(8) - col_coord = (accum_idx // 4) * 8 + 2 * (lane % 4) + col_coord = (accum_idx // 4) * 8 + 2 * (lane_idx % 4) if cutlass.const_expr(accum_idx % 2 == 1): col_coord = col_coord + cutlass.Int32(1) a_acc[accum_idx] = a_acc[accum_idx] if row_coord >= col_coord else cutlass.Float32(0.0) bars.mb_intermediate_done[intermediate_stage].wait(((global_chunk // cfg.smem_intermediate_stages) + 1) % 2) nvvm.stmatrix( - sIntermediate_ptr + stsm_idx, + sIntermediate_ptr + a_idx, [ fp32_to_fp16(a_acc[0], a_acc[1], dtype=cfg.io_dtype), fp32_to_fp16(a_acc[2], a_acc[3], dtype=cfg.io_dtype), @@ -1033,7 +971,7 @@ def epilogue_warp( bars.mb_decay_super_done[decay_stage].arrive() qk_scale_index = advance(qk_scale_index, cfg.qk_scale_ready_stages) - # ---- checkpoint + O drain: checkpoint stores first (CG1 stages checkpoint before O) -- + # ---- checkpoint + O store: checkpoint first ------------------------------ if local_chunk > 0: output_chunk = chunk_idx - cutlass.Int32(1) output_chunk_start = output_chunk * cfg.b_t @@ -1041,9 +979,9 @@ def epilogue_warp( did_checkpoint = cutlass.Int32(0) checkpoint_stage = cutlass.Int32(0) if cutlass.const_expr(cfg.enable_checkpoints): - # ---- checkpoint store ---------------------------------------- + # ---- checkpoint store -------------------------------------------- do_checkpoint = checkpoint_mod == 0 - do_checkpoint = do_checkpoint and chunk_idx >= wstart + do_checkpoint = do_checkpoint and chunk_idx >= write_start checkpoint_stage = checkpoint_ready_index.idx if do_checkpoint: bars.mb_checkpoint_tmastg_ready[checkpoint_ready_index.idx].wait(checkpoint_ready_index.phase) @@ -1060,7 +998,7 @@ def epilogue_warp( bars.mb_o_tmastg_ready[o_stage].wait(((global_chunk - cutlass.Int32(1)) // cfg.smem_o_stages) % 2) o_slice = tma_slice_runtime_desc(desc_o_slot, cutlass.Int32(0), head_o, output_chunk_start) did_o = cutlass.Int32(0) - if output_chunk >= wstart: + if output_chunk >= write_start: tma_store_tile(sO_tma[o_stage], o_slice, acquire=False) tma_store_commit() did_o = cutlass.Int32(1) @@ -1082,9 +1020,9 @@ def epilogue_warp( tma_store_wait(0) bars.mb_o_tmastg_done[o_stage].arrive() - # ---- last computed chunk drain (always owned: it is wend - 1) ------------ + # ---- last computed chunk store (always owned: it is wend - 1) ---------------- if num_tile_chunks > 0: - output_chunk = wend - cutlass.Int32(1) + output_chunk = write_end - cutlass.Int32(1) last_global_chunk = global_chunk_base + num_tile_chunks - cutlass.Int32(1) output_chunk_start = output_chunk * cfg.b_t o_stage = last_global_chunk % cfg.smem_o_stages @@ -1095,7 +1033,7 @@ def epilogue_warp( tma_store_wait(0) bars.mb_o_tmastg_done[o_stage].arrive() global_chunk_base += num_tile_chunks - tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one) @cute.jit @@ -1103,9 +1041,7 @@ def gate_scale(cfg, raw_gate: cutlass.Float32) -> cutlass.Float32: """Map raw gate to the log2-domain decay increment.""" if cutlass.const_expr(cfg.safe_gate): - half = cutlass.Float32(0.5) - sigmoid = cute.math.tanh(raw_gate * half, approx=True) * half + half - return cfg.gate_scale_log2 * sigmoid + return cfg.gate_scale_log2 * sigmoid(raw_gate) # Default ABI: gate arrives in natural-log space return raw_gate * cutlass.Float32(LOG2_E) @@ -1118,8 +1054,8 @@ def compute0_warp_group( num_ctas, cu_seqlens, mWorkItems, - sSched, - lane, + sScheduler, + lane_idx, warp_idx, mQ, mA_log, @@ -1140,25 +1076,30 @@ def compute0_warp_group( """CG0 warp role (warps 0-7, two ping-pong groups): Gate prefix scan and the decay/restore operand materialization.""" nvvm.setmaxregister(cfg.num_regs_compute_group_0, nvvm.SetMaxRegisterAction.INCREASE) + elect_one = nvvm.elect_sync() + + scheduler_state = PipelineState.start(phase=0) + cg0_warp = warp_idx - cfg.compute_group_0_warp_ids[0] - cg0_group_id = cg0_warp // cfg.cg0_warps_per_group cg0_local_warp = cg0_warp % cfg.cg0_warps_per_group - prefix_dim = cg0_local_warp * cfg.threads_per_warp + lane + + cg0_group_id = cg0_warp // cfg.cg0_warps_per_group + channel_dim = cg0_local_warp * cfg.threads_per_warp + lane_idx cg0_a_log_exp = cutlass.Float32(1.0) cg0_dt_bias_value = cutlass.Float32(0.0) global_chunk_base = cutlass.Int32(0) - sched_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) opaque_one = opaque_f32_zero() + cutlass.Float32(1.0) while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) head_o = head_idx - num_tile_chunks = wend - cstart + num_tile_chunks = write_end - compute_start if cutlass.const_expr(cfg.safe_gate): if num_tile_chunks > 0: cg0_a_log_exp = cute.math.exp2(mA_log[head_o].to(cutlass.Float32) * LOG2_E, fastmath=True) - cg0_dt_bias_value = mDt_bias[head_o, prefix_dim].to(cutlass.Float32) - # tile entry: both ping-pong groups inherit each other's delivery proofs (parity-swap guard) + cg0_dt_bias_value = mDt_bias[head_o, channel_dim].to(cutlass.Float32) nvvm.barrier_cta_sync(cfg.cg0_tile_entry_barrier_id, thread_count=cfg.cg0_group_count * cfg.cg0_threads_per_group) cg0_first_global_chunk = global_chunk_base + cutlass.Int32(cg0_group_id) diag_ring_stage = cg0_first_global_chunk % cutlass.Int32(cfg.smem_state_scale_diag_stages) @@ -1167,7 +1108,7 @@ def compute0_warp_group( raw_bar_stage = cg0_first_global_chunk % cutlass.Int32(cfg.smem_raw_bar_stages) raw_bar_phase = (cg0_first_global_chunk // cutlass.Int32(cfg.smem_raw_bar_stages)) % cutlass.Int32(2) for local_chunk in cutlass.range(cg0_group_id, num_tile_chunks, cfg.cg0_group_count, unroll=1): - chunk_idx = cstart + local_chunk + chunk_idx = compute_start + local_chunk global_chunk = global_chunk_base + local_chunk chunk_start = chunk_idx * cfg.b_t decay_stage = global_chunk % cfg.smem_decay_stages @@ -1189,48 +1130,36 @@ def compute0_warp_group( bars.mb_gate_ready[raw_bar_stage].wait(raw_bar_phase) row_group_start = cg0_local_warp * (cfg.b_t // cfg.cg0_warps_per_group) - lane_row_group = lane // 8 - lane_in_row_group = lane - lane_row_group * 8 + lane_row_group = lane_idx // 8 + lane_in_row_group = lane_idx - lane_row_group * 8 decay_row = row_group_start + lane_row_group - decay_key_mask = cutlass.Int32(8) - prefix_dim = cg0_local_warp * cfg.threads_per_warp + lane + channel_dim = cg0_local_warp * cfg.threads_per_warp + lane_idx - # ---- Gate prefix scan ----------------------------------------------- - f32_segment = prefix_dim // 32 + # ---- Gate prefix scan ---------------------------------------------------- + f32_segment = channel_dim // 32 prefix_seg_base = f32_segment * (cfg.b_t * 32) - prefix_col = prefix_dim - f32_segment * 32 + prefix_col = channel_dim - f32_segment * 32 g_prefix_regs = cutlass.Array(cutlass.Float32, cfg.b_t, alignment=16) if cutlass.const_expr(cfg.safe_gate): - valid_rows = seqlen_b - chunk_idx * cutlass.Int32(cfg.b_t) - valid_mask = cutlass.vector.create_mask([cfg.b_t], [valid_rows]) - for row_pair in cutlass.range_constexpr(cfg.b_t // 2): - row0 = row_pair * 2 - row1 = row0 + 1 - prefix_idx0 = prefix_seg_base + swizzle_xor_128b(row0, row0 * 32 + prefix_col, elem_bytes=4) - prefix_idx1 = prefix_seg_base + swizzle_xor_128b(row1, row1 * 32 + prefix_col, elem_bytes=4) - gate0 = (sGate_ptr + prefix_idx0).load() - gate1 = (sGate_ptr + prefix_idx1).load() - gate0 = cg0_a_log_exp * (gate0 + cg0_dt_bias_value) - gate1 = cg0_a_log_exp * (gate1 + cg0_dt_bias_value) - gate0 = gate_scale( - cfg, - gate0, - ) - gate1 = gate_scale( - cfg, - gate1, - ) - gate_pair = cutlass.Vector.from_elements((gate0, gate1), cutlass.Float32) - gate_pair = cutlass.vector.where(valid_mask[row0 : row1 + 1], gate_pair, 0.0) - g_prefix_regs[row0] = gate_pair[0] - g_prefix_regs[row1] = gate_pair[1] + for row in cutlass.range_constexpr(cfg.b_t): + prefix_idx = prefix_seg_base + swizzle_xor_128b(row, row * 32 + prefix_col, elem_bytes=4) + gate = (sGate_ptr + prefix_idx).load() + token_idx = chunk_idx * cutlass.Int32(cfg.b_t) + cutlass.Int32(row) + if token_idx < batch_seqlen: + gate = gate_scale( + cfg, + cg0_a_log_exp * (gate + cg0_dt_bias_value), + ) + else: + gate = cutlass.Float32(0.0) + g_prefix_regs[row] = gate else: for row in cutlass.range_constexpr(cfg.b_t): prefix_idx = prefix_seg_base + swizzle_xor_128b(row, row * 32 + prefix_col, elem_bytes=4) gate = (sGate_ptr + prefix_idx).load() token_idx = chunk_idx * cutlass.Int32(cfg.b_t) + cutlass.Int32(row) - if token_idx < seqlen_b: + if token_idx < batch_seqlen: gate = gate_scale( cfg, gate, @@ -1251,7 +1180,7 @@ def compute0_warp_group( g_prefix_regs[row1] = prefix1 prefix_acc = prefix1 - # ---- exp2(g): stage prefixes + final-token decay --------------------- + # ---- exp2(g): stage prefixes + final-token decay ------------------------- for row in cutlass.range_constexpr(cfg.b_t): g_prefix_regs[row] = cute.math.exp2(g_prefix_regs[row], fastmath=True) @@ -1260,13 +1189,11 @@ def compute0_warp_group( prefix_idx = prefix_seg_base + swizzle_xor_128b(row, row * 32 + prefix_col, elem_bytes=4) (sGate_ptr + prefix_idx).store(g_prefix_regs[row]) - # ---- state-scale diag: stage exp2(g_last) decay blocks --------------- + # ---- state-scale diag: stage exp2(g last) decay blocks ------------------- bars.mb_state_scale_diag_done[state_scale_diag_stage].wait(diag_ring_phase ^ cutlass.Int32(1)) - block = prefix_dim // cutlass.Int32(16) - coord = prefix_dim - block * cutlass.Int32(16) - storage_col = coord ^ cutlass.Int32((cfg.b_t // 2)) - linear_idx = block * cutlass.Int32(256) + coord * cutlass.Int32(16) + storage_col - diag_idx = swizzle_lin_S(linear_idx, bbits=1, mbase=3, sshift=3) + block = channel_dim // cutlass.Int32(16) + coord = channel_dim - block * cutlass.Int32(16) + diag_idx = block * cutlass.Int32(256) + coord * cutlass.Int32(16) + swizzle_xor_32b(channel_dim, coord) sState_scale_diag_ptr[diag_idx] = exp_g_last.to(cfg.io_dtype) nvvm.barrier_cta_sync(cfg.cg0_group_sync_barrier_base_id + cg0_group_id, thread_count=cfg.cg0_threads_per_group) @@ -1279,7 +1206,7 @@ def compute0_warp_group( raw_k_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) raw_beta_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) - # ---- optional Q/K L2-norm ------------------------------------------- + # ---- optional Q/K L2-norm ------------------------------------------------ if cutlass.const_expr(cfg.l2norm): qk0_lo = opaque_f32_zero() qk0_hi = opaque_f32_zero() @@ -1304,8 +1231,7 @@ def compute0_warp_group( raw_k_regs[reg_base + dim_offset] = k_val beta_val = raw_beta_vec_f32[dim_offset] if cutlass.const_expr(cfg.beta_sigmoid): - half = cutlass.Float32(0.5) - beta_val = (cute.math.tanh(beta_val * half, approx=True) * half + half).to(cfg.io_dtype).to(cutlass.Float32) + beta_val = sigmoid(beta_val).to(cfg.io_dtype).to(cutlass.Float32) raw_beta_regs[reg_base + dim_offset] = beta_val if cutlass.const_expr(cfg.l2norm): if cutlass.const_expr(dim_offset % 2 == 0): @@ -1328,7 +1254,7 @@ def compute0_warp_group( q_inv_norm = cute.math.rsqrt(cute.math.max(q_sum_sq, norm_floor_sq), fastmath=True) k_inv_norm = cute.math.rsqrt(cute.math.max(k_sum_sq, norm_floor_sq), fastmath=True) - # ---- decay/restore operands: exp2(+-g) applied per key channel ------- + # ---- decay/restore operands: exp2(+-g) applied per key channel ----------- exp_g_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) exp_g_last_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) for dim_half in cutlass.range_constexpr(2): @@ -1360,7 +1286,7 @@ def compute0_warp_group( exp_g_last_regs[f32_reg_base + 2] = exp_g_last_frag[2] exp_g_last_regs[f32_reg_base + 3] = exp_g_last_frag[3] - # ---- K decay + K_inv operands: K * exp2(+g) and K * exp2(-g) ----- + # ---- K decay + K inv operands: K * exp2(+g) and K * exp2(-g) --------- k_decay_pack = cutlass.Array(cutlass.Int32, 4, alignment=16) for pair_idx in cutlass.range_constexpr(4): dim0 = pair_idx * 2 @@ -1402,11 +1328,9 @@ def compute0_warp_group( f16_segment_dim = dim_base - f16_segment * 64 k_inv_swizzled_idx = f16_segment * (cfg.b_t * 64) + decay_row * 64 + swizzle_xor_128b(decay_row, f16_segment_dim, elem_bytes=2) (sK_inv_ptr + k_inv_swizzled_idx).store(k_inv_vec, alignment=16) - storage_key = dim_base ^ decay_key_mask - storage_slice = storage_key // 64 - decay_swizzled_idx = storage_slice * (cfg.b_t * 64) + swizzle_xor_128b( - decay_row, decay_row * 64 + storage_key - storage_slice * 64, elem_bytes=2 - ) + decay_col = dim_base + decay_segment = decay_col // 64 + decay_swizzled_idx = decay_segment * (cfg.b_t * 64) + swizzle_xor_128b(decay_row, decay_row * 64 + decay_col - decay_segment * 64, elem_bytes=2) (sK_decay_ptr + decay_swizzled_idx).store(k_decay_vec, alignment=16) nvvm.fence_proxy("async.shared", space="cta") bars.mb_k_decay_inv_cg0_ready[decay_stage].arrive() @@ -1415,7 +1339,7 @@ def compute0_warp_group( bars.mb_gate_done[raw_stage].arrive() bars.mb_beta_done[raw_stage].arrive() - # ---- Q_decay operand: Q * q_inv_norm -------------------------------- + # ---- Q decay operand: Q * q inv norm ------------------------------------- for dim_half in cutlass.range_constexpr(2): dim_base = dim_half * (cfg.d_k // 2) + lane_in_row_group * 8 reg_base = dim_half * 8 @@ -1439,14 +1363,12 @@ def compute0_warp_group( ), cutlass.Int32, ).bitcast(cfg.io_dtype) - storage_key = dim_base ^ decay_key_mask - storage_slice = storage_key // 64 - decay_swizzled_idx = storage_slice * (cfg.b_t * 64) + swizzle_xor_128b( - decay_row, decay_row * 64 + storage_key - storage_slice * 64, elem_bytes=2 - ) + decay_col = dim_base + decay_segment = decay_col // 64 + decay_swizzled_idx = decay_segment * (cfg.b_t * 64) + swizzle_xor_128b(decay_row, decay_row * 64 + decay_col - decay_segment * 64, elem_bytes=2) (sQ_decay_ptr + decay_swizzled_idx).store(q_decay_vec, alignment=16) - # ---- K_restore operand: K_inv * exp_g_last -------------------------- + # ---- K restore operand: K inv * exp2(g last) ----------------------------- bars.mb_k_restore_acc_done[decay_stage].wait(((global_chunk // cfg.smem_decay_stages + 1) % 2)) for dim_half in cutlass.range_constexpr(2): dim_base = dim_half * (cfg.d_k // 2) + lane_in_row_group * 8 @@ -1457,10 +1379,9 @@ def compute0_warp_group( dim1 = dim0 + 1 exp_g_last_pair = fp32_to_fp16(exp_g_last_regs[reg_base + dim0], exp_g_last_regs[reg_base + dim1], dtype=cfg.io_dtype) k_restore_pack[pair_idx] = mul_f16x2(k_inv_pack[dim_half * 4 + pair_idx], exp_g_last_pair, cfg.io_dtype) - storage_row = decay_row ^ (cfg.b_t // 2) f16_segment = dim_base // 64 f16_segment_dim = dim_base - f16_segment * 64 - k_restore_idx = f16_segment * (cfg.b_t * 64) + storage_row * 64 + swizzle_xor_128b(storage_row, f16_segment_dim, elem_bytes=2) + k_restore_idx = f16_segment * (cfg.b_t * 64) + decay_row * 64 + swizzle_xor_128b(decay_row, f16_segment_dim, elem_bytes=2) k_restore_vec = cutlass.Vector.from_elements( ( k_restore_pack[0], @@ -1485,7 +1406,7 @@ def compute0_warp_group( raw_bar_stage = raw_bar_stage - cutlass.Int32(cfg.smem_raw_bar_stages) if raw_bar_wrapped else raw_bar_stage raw_bar_phase = raw_bar_phase ^ (cutlass.Int32(1) if raw_bar_wrapped else cutlass.Int32(0)) global_chunk_base += num_tile_chunks - tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one) @cute.jit @@ -1496,8 +1417,8 @@ def compute1_warp_group( num_ctas, cu_seqlens, mWorkItems, - sSched, - lane, + sScheduler, + lane_idx, sTmem_base, warp_idx, mState_out, @@ -1514,126 +1435,107 @@ def compute1_warp_group( """CG1 warp role (warps 8-11): persistent scheduler loop running the value-side TMEM staging, O drain, and checkpoint/final-state stores.""" nvvm.setmaxregister(cfg.num_regs_compute_group_1, nvvm.SetMaxRegisterAction.INCREASE) + elect_one = nvvm.elect_sync() + + checkpoint_done_index = PipelineState.start(phase=1) + sO_ptr = sO_raw.data_ptr() sCheckpoint_ptr = sCheckpoint_raw.data_ptr() if cutlass.const_expr(cfg.enable_checkpoints) else sO_raw.data_ptr() - checkpoint_done_index = PipelineState.start(phase=1) nvvm.barrier_cta_sync(cfg.tmem_lifecycle_barrier_id, thread_count=cfg.tmem_user_threads) tmem_base = sTmem_base.load() tmem_col = tmem_base & 0xFFFF tmem_row = tmem_base >> 16 + row_lo_addr = tmem_row << 16 + row_hi_addr = (tmem_row + 16) << 16 tmem_subpartition = warp_idx % (cfg.d_v // cfg.threads_per_warp) - frag_row_coord = (lane // 16) * 8 + (lane & 7) - frag_col_offset = ((lane // 8) & 1) * 8 - row_id = tmem_row + tmem_subpartition * cfg.threads_per_warp - value_dim = tmem_subpartition * cfg.threads_per_warp + lane + ov_row_coord = (lane_idx // 16) * 8 + (lane_idx & 7) + ov_col_offset = ((lane_idx // 8) & 1) * 8 + value_dim = tmem_subpartition * cfg.threads_per_warp + lane_idx state_k_acc_index = PipelineState.start(phase=0) u_acc_index = PipelineState.start(phase=0) o_acc_index = PipelineState.start(phase=0) k_restore_index = PipelineState.start(phase=0) # CG1's per-chunk mb_k_restore_acc_done wait slot raw_index = PipelineState.start(phase=0) # raw-ring slot for the sV/sW reads + inputs_done arrives - raw_bar_index = PipelineState.start(phase=0) # even-depth ready-ring slot (decoupled from the data ring) + raw_bar_index = PipelineState.start(phase=0) # even-depth ready-ring slot global_chunk_base = cutlass.Int32(0) - sched_state = PipelineState.start(phase=0) + scheduler_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) head_o = head_idx - num_tile_chunks = wend - cstart + num_tile_chunks = write_end - compute_start if num_tile_chunks > 0: - # ---- first chunk: seed state TMEM from mState_init ---------- - seed_from_initial_state = cstart == 0 + # ---- first chunk: seed state TMEM from mState init ----------------------- + seed_from_initial_state = compute_start == 0 + sV_ptr = sV_raw.data_ptr() + raw_index.idx * (cfg.d_v * cfg.b_t) + sW_ptr = sW_raw.data_ptr() + raw_index.idx * (cfg.d_v * cfg.b_t) + + # ---- state seed: initial state GMEM -> state TMEM ------------------------ + state_col_id = tmem_col + cfg.tmem_state_acc_offset + packed_col_id = tmem_col + cfg.tmem_state_input_offset if cutlass.const_expr(mState_init is not None): if seed_from_initial_state: seed_vw = 16 // (mState_init.element_type.width // 8) seed_src = (mState_init.iterator + mState_init.layout((batch_idx, head_o, value_dim, 0))).raw_ptr() - for key_block_start in cutlass.range_constexpr(0, cfg.d_k, 32): - state_block = cutlass.Array(cutlass.Float32, 32, alignment=16) - for g in cutlass.range_constexpr(32 // seed_vw): - seed_chunk = (seed_src + key_block_start + g * seed_vw).load(count=seed_vw, alignment=16) + state_vecs = [] + for i in cutlass.range_constexpr(cfg.d_k // 16): + state_block = [] + for g in cutlass.range_constexpr(16 // seed_vw): + seed_chunk = (seed_src + i * 16 + g * seed_vw).load(count=seed_vw, alignment=16) for t in cutlass.range_constexpr(seed_vw): - state_block[g * seed_vw + t] = seed_chunk[t].to(cutlass.Float32) + state_block.append(seed_chunk[t].to(cutlass.Float32)) + state_vecs.append(state_block) + for i in cutlass.range_constexpr(cfg.d_k // 16): + state_pack = cutlass.Array(cutlass.Int32, 8, alignment=16) + for packed_col in cutlass.range_constexpr(8): + state_pack[packed_col] = fp32_to_fp16(state_vecs[i][2 * packed_col], state_vecs[i][2 * packed_col + 1], dtype=cfg.io_dtype) nvvm.tcgen05_st( "32x32b", - nvvm.make_tmem_ptr((row_id << 16) + (tmem_col + cfg.tmem_state_acc_offset + key_block_start), cutlass.Float32), - state_block[0:32], + nvvm.make_tmem_ptr(row_lo_addr + packed_col_id + i * 8, cutlass.Int8), + state_pack[0:8], ) - else: - for key_block_start in cutlass.range_constexpr(0, cfg.d_k, 32): - state_block = cutlass.Array(cutlass.Float32, 32, alignment=16) - for col in cutlass.range_constexpr(32): - state_block[col] = cutlass.Float32(0.0) - nvvm.tcgen05_st( - "32x32b", - nvvm.make_tmem_ptr((row_id << 16) + (tmem_col + cfg.tmem_state_acc_offset + key_block_start), cutlass.Float32), - state_block[0:32], - ) - if cutlass.const_expr(mState_init is not None): - nvvm.tcgen05_wait("store") - sV_ptr = sV_raw.data_ptr() + raw_index.idx * (cfg.d_v * cfg.b_t) - sW_ptr = sW_raw.data_ptr() + raw_index.idx * (cfg.d_v * cfg.b_t) - - row_addr = (tmem_row + tmem_subpartition * cfg.threads_per_warp) << 16 - state_col_id = tmem_col + cfg.tmem_state_acc_offset - - # ---- state repack: acc TMEM -> packed b16 TMEM ---------------------- - packed_col_id = tmem_col + cfg.tmem_state_inp_offset - if cutlass.const_expr(mState_init is not None): - state_vecs = [] - for k_block in cutlass.range_constexpr(cfg.d_k // 16): - state_vecs.append(nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + state_col_id + k_block * 16, cutlass.Float32), num=16)) - - for k_block in cutlass.range_constexpr(cfg.d_k // 16): - state_pack = cutlass.Array(cutlass.Int32, 8, alignment=16) - for packed_col in cutlass.range_constexpr(8): - source_pair = packed_col ^ 4 - state_pack[packed_col] = fp32_to_fp16( - state_vecs[k_block][2 * source_pair], state_vecs[k_block][2 * source_pair + 1], dtype=cfg.io_dtype - ) - nvvm.tcgen05_st( - "32x32b", - nvvm.make_tmem_ptr((tmem_row << 16) + packed_col_id + k_block * 8, cutlass.Int8), - state_pack[0:8], - ) - - nvvm.tcgen05_wait("store") - bars.mb_state_inp_ready.arrive() - if cutlass.const_expr(cfg.enable_checkpoints): - if wstart == 0: - checkpoint_stage = checkpoint_done_index.idx - bars.mb_checkpoint_tmastg_done[checkpoint_stage].wait(checkpoint_done_index.phase) - checkpoint_done_index = advance(checkpoint_done_index, cfg.smem_checkpoint_stages) - checkpoint_stage_base = checkpoint_stage * (cfg.d_k * cfg.d_v) - for k_block in cutlass.range_constexpr(cfg.d_k // 16): - for g in cutlass.range_constexpr(2): - packs = tuple( - fp32_to_fp16(state_vecs[k_block][g * 8 + 2 * t], state_vecs[k_block][g * 8 + 2 * t + 1], dtype=cfg.io_dtype) - for t in range(4) - ) - dk = k_block * 16 + g * 8 - checkpoint_addr = ( - checkpoint_stage_base + (dk // 64) * (cfg.d_v * 64) + value_dim * 64 + swizzle_xor_128b(value_dim, dk % 64, elem_bytes=2) - ) - (sCheckpoint_ptr + checkpoint_addr).store( - cutlass.Vector.from_elements(packs, cutlass.Int32).bitcast(cfg.io_dtype), alignment=16 - ) - nvvm.tcgen05_wait("load") - nvvm.fence_proxy("async.shared", space="cta") - bars.mb_checkpoint_tmastg_ready[checkpoint_stage].arrive() - bars.mb_state_acc_read_done.arrive() + nvvm.tcgen05_wait("store") + bars.mb_state_input_ready.arrive() + if cutlass.const_expr(cfg.enable_checkpoints): + if write_start == 0: + checkpoint_stage = checkpoint_done_index.idx + bars.mb_checkpoint_tmastg_done[checkpoint_stage].wait(checkpoint_done_index.phase) + checkpoint_done_index = advance(checkpoint_done_index, cfg.smem_checkpoint_stages) + checkpoint_stage_base = checkpoint_stage * (cfg.d_k * cfg.d_v) + for i in cutlass.range_constexpr(cfg.d_k // 16): + for g in cutlass.range_constexpr(2): + packs = tuple( + fp32_to_fp16(state_vecs[i][g * 8 + 2 * t], state_vecs[i][g * 8 + 2 * t + 1], dtype=cfg.io_dtype) for t in range(4) + ) + dk = i * 16 + g * 8 + checkpoint_addr = ( + checkpoint_stage_base + + (dk // 64) * (cfg.d_v * 64) + + value_dim * 64 + + swizzle_xor_128b(value_dim, dk % 64, elem_bytes=2) + ) + (sCheckpoint_ptr + checkpoint_addr).store( + cutlass.Vector.from_elements(packs, cutlass.Int32).bitcast(cfg.io_dtype), alignment=16 + ) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_checkpoint_tmastg_ready[checkpoint_stage].arrive() + bars.mb_state_acc_read_done.arrive() if cutlass.const_expr(cfg.enable_checkpoints and mState_init is None): - if wstart == 0: + if write_start == 0: checkpoint_stage = checkpoint_done_index.idx bars.mb_checkpoint_tmastg_done[checkpoint_stage].wait(checkpoint_done_index.phase) checkpoint_done_index = advance(checkpoint_done_index, cfg.smem_checkpoint_stages) checkpoint_stage_base = checkpoint_stage * (cfg.d_k * cfg.d_v) zero_packs = tuple(cutlass.Int32(0) for _ in range(4)) - for k_block in cutlass.range_constexpr(cfg.d_k // 16): + for i in cutlass.range_constexpr(cfg.d_k // 16): for g in cutlass.range_constexpr(2): - dk = k_block * 16 + g * 8 + dk = i * 16 + g * 8 checkpoint_addr = ( checkpoint_stage_base + (dk // 64) * (cfg.d_v * 64) + value_dim * 64 + swizzle_xor_128b(value_dim, dk % 64, elem_bytes=2) ) @@ -1643,129 +1545,98 @@ def compute1_warp_group( nvvm.fence_proxy("async.shared", space="cta") bars.mb_checkpoint_tmastg_ready[checkpoint_stage].arrive() - # ---- Y staging: Y = W*V - state*(Beta*K) ----------------------------- + # ---- Y stage: Y = W*V - state*(Beta*K) ----------------------------------- bars.mb_v_ready[raw_bar_index.idx].wait(raw_bar_index.phase) projection_col_id = tmem_col + cfg.tmem_state_k_acc_offset - input_col_id = tmem_col + cfg.tmem_y_inp_offset + input_col_id = tmem_col + cfg.tmem_y_input_offset value_dim_base = tmem_subpartition * cfg.threads_per_warp - # ---- raw V fragments, then W, then the state*K acc readback ---------- - row_id0 = tmem_row + value_dim_base - row_id1 = row_id0 + 16 - raw_v_frag0 = nvvm.ldmatrix( + # ---- raw V fragments, then W, then the k state acc read ------------------ + raw_v_frag_lo = nvvm.ldmatrix( sV_ptr - + (value_dim_base + frag_col_offset) // 64 * (cfg.b_t * 64) - + frag_row_coord * 64 - + swizzle_xor_128b(frag_row_coord, (value_dim_base + frag_col_offset) % 64, elem_bytes=2), + + (value_dim_base + ov_col_offset) // 64 * (cfg.b_t * 64) + + ov_row_coord * 64 + + swizzle_xor_128b(ov_row_coord, (value_dim_base + ov_col_offset) % 64, elem_bytes=2), 4, nvvm.MMALayout.COL, ) - raw_v_frag1 = nvvm.ldmatrix( + raw_v_frag_hi = nvvm.ldmatrix( sV_ptr - + (value_dim_base + 16 + frag_col_offset) // 64 * (cfg.b_t * 64) - + frag_row_coord * 64 - + swizzle_xor_128b(frag_row_coord, (value_dim_base + 16 + frag_col_offset) % 64, elem_bytes=2), + + (value_dim_base + 16 + ov_col_offset) // 64 * (cfg.b_t * 64) + + ov_row_coord * 64 + + swizzle_xor_128b(ov_row_coord, (value_dim_base + 16 + ov_col_offset) % 64, elem_bytes=2), 4, nvvm.MMALayout.COL, ) bars.mb_w_ready[raw_bar_index.idx].wait(raw_bar_index.phase) - raw_w_frag0 = nvvm.ldmatrix( + raw_w_frag_lo = nvvm.ldmatrix( sW_ptr - + (value_dim_base + frag_col_offset) // 64 * (cfg.b_t * 64) - + frag_row_coord * 64 - + swizzle_xor_128b(frag_row_coord, (value_dim_base + frag_col_offset) % 64, elem_bytes=2), + + (value_dim_base + ov_col_offset) // 64 * (cfg.b_t * 64) + + ov_row_coord * 64 + + swizzle_xor_128b(ov_row_coord, (value_dim_base + ov_col_offset) % 64, elem_bytes=2), 4, nvvm.MMALayout.COL, ) - raw_w_frag1 = nvvm.ldmatrix( + raw_w_frag_hi = nvvm.ldmatrix( sW_ptr - + (value_dim_base + 16 + frag_col_offset) // 64 * (cfg.b_t * 64) - + frag_row_coord * 64 - + swizzle_xor_128b(frag_row_coord, (value_dim_base + 16 + frag_col_offset) % 64, elem_bytes=2), + + (value_dim_base + 16 + ov_col_offset) // 64 * (cfg.b_t * 64) + + ov_row_coord * 64 + + swizzle_xor_128b(ov_row_coord, (value_dim_base + 16 + ov_col_offset) % 64, elem_bytes=2), 4, nvvm.MMALayout.COL, ) - if cutlass.const_expr(mState_init is not None): - bars.mb_state_k_acc_ready.wait(state_k_acc_index.phase) - state_k_vec0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id0 << 16) + projection_col_id, cutlass.Float32), num=2) - state_k_vec1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id1 << 16) + projection_col_id, cutlass.Float32), num=2) - - y_inp_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + y_lo = [cutlass.Int32(0) for _ in range(4)] + y_hi = [cutlass.Int32(0) for _ in range(4)] for reg_idx in cutlass.range_constexpr(4): - raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1) - frag_pair = (reg_idx ^ 2) * 2 - w_pair = raw_w_frag0[raw_matrix] - wv_pair = mul_f16x2( - w_pair, - raw_v_frag0[raw_matrix], - cfg.io_dtype, - ) - if cutlass.const_expr(mState_init is not None): - state_k_val0, state_k_val1 = state_k_vec0[frag_pair], state_k_vec0[frag_pair + 1] - state_k_pair = fp32_to_fp16(state_k_val0, state_k_val1, dtype=cfg.io_dtype) - y_inp_pack0[reg_idx] = sub_f16x2( - wv_pair, - state_k_pair, - cfg.io_dtype, - ) - else: - y_inp_pack0[reg_idx] = wv_pair - - y_inp_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + y_lo[reg_idx] = mul_f16x2(raw_w_frag_lo[reg_idx], raw_v_frag_lo[reg_idx], cfg.io_dtype) + y_hi[reg_idx] = mul_f16x2(raw_w_frag_hi[reg_idx], raw_v_frag_hi[reg_idx], cfg.io_dtype) + if cutlass.const_expr(mState_init is not None): + if seed_from_initial_state: + bars.mb_state_k_acc_ready.wait(state_k_acc_index.phase) + state_k_acc_index = advance(state_k_acc_index, 1) + state_k_vec_lo = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_lo_addr + projection_col_id, cutlass.Float32), num=2) + state_k_vec_hi = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_hi_addr + projection_col_id, cutlass.Float32), num=2) + for reg_idx in cutlass.range_constexpr(4): + frag_pair = reg_idx * 2 + state_k_lo = fp32_to_fp16(state_k_vec_lo[frag_pair], state_k_vec_lo[frag_pair + 1], dtype=cfg.io_dtype) + state_k_hi = fp32_to_fp16(state_k_vec_hi[frag_pair], state_k_vec_hi[frag_pair + 1], dtype=cfg.io_dtype) + y_lo[reg_idx] = sub_f16x2(y_lo[reg_idx], state_k_lo, cfg.io_dtype) + y_hi[reg_idx] = sub_f16x2(y_hi[reg_idx], state_k_hi, cfg.io_dtype) + + y_input_pack_lo = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + y_input_pack_hi = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) for reg_idx in cutlass.range_constexpr(4): - raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1) - frag_pair = (reg_idx ^ 2) * 2 - w_pair = raw_w_frag1[raw_matrix] - wv_pair = mul_f16x2( - w_pair, - raw_v_frag1[raw_matrix], - cfg.io_dtype, - ) - if cutlass.const_expr(mState_init is not None): - state_k_val0, state_k_val1 = state_k_vec1[frag_pair], state_k_vec1[frag_pair + 1] - state_k_pair = fp32_to_fp16(state_k_val0, state_k_val1, dtype=cfg.io_dtype) - y_inp_pack1[reg_idx] = sub_f16x2( - wv_pair, - state_k_pair, - cfg.io_dtype, - ) - else: - y_inp_pack1[reg_idx] = wv_pair - - nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr((tmem_row << 16) + input_col_id, cutlass.Int8), y_inp_pack0[0:4]) - - nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr((tmem_row + 16 << 16) + input_col_id, cutlass.Int8), y_inp_pack1[0:4]) - + y_input_pack_lo[reg_idx] = y_lo[reg_idx] + y_input_pack_hi[reg_idx] = y_hi[reg_idx] + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_lo_addr + input_col_id, cutlass.Int8), y_input_pack_lo[0:4]) + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_hi_addr + input_col_id, cutlass.Int8), y_input_pack_hi[0:4]) nvvm.tcgen05_wait("store") - if cutlass.const_expr(mState_init is not None): - state_k_acc_index = advance(state_k_acc_index, 1) bars.mb_v_done[raw_index.idx].arrive() bars.mb_w_done[raw_index.idx].arrive() - bars.mb_y_inp_ready.arrive() + bars.mb_y_input_ready.arrive() - # ---- U repack: acc TMEM -> packed b16 TMEM -------------------------- + # ---- U stage: acc TMEM -> packed b16 TMEM -------------------------------- bars.mb_u_acc_ready.wait(u_acc_index.phase) u_acc_vals = nvvm.tcgen05_ld( "32x32b", - nvvm.make_tmem_ptr((tmem_row + tmem_subpartition * cfg.threads_per_warp << 16) + (tmem_col + cfg.tmem_u_acc_offset), cutlass.Float32), + nvvm.make_tmem_ptr(row_lo_addr + (tmem_col + cfg.tmem_u_acc_offset), cutlass.Float32), num=cfg.b_t, ) - u_inp_pack = cutlass.Array(cutlass.Int32, (cfg.b_t // 2), alignment=16) + u_input_pack = cutlass.Array(cutlass.Int32, (cfg.b_t // 2), alignment=16) for packed_col in cutlass.range_constexpr((cfg.b_t // 2)): - source_pair = packed_col ^ 4 - token0 = source_pair * 2 + token0 = packed_col * 2 token1 = token0 + 1 - u_inp_pack[packed_col] = fp32_to_fp16(u_acc_vals[token0], u_acc_vals[token1], dtype=cfg.io_dtype) + u_input_pack[packed_col] = fp32_to_fp16(u_acc_vals[token0], u_acc_vals[token1], dtype=cfg.io_dtype) nvvm.tcgen05_st( "32x32b", - nvvm.make_tmem_ptr((tmem_row << 16) + (tmem_col + cfg.tmem_u_inp_offset), cutlass.Int8), - u_inp_pack[0 : (cfg.b_t // 2)], + nvvm.make_tmem_ptr(row_lo_addr + (tmem_col + cfg.tmem_u_input_offset), cutlass.Int8), + u_input_pack[0 : (cfg.b_t // 2)], ) nvvm.tcgen05_wait("store") u_acc_index = advance(u_acc_index, 1) - bars.mb_u_inp_ready.arrive() + bars.mb_u_input_ready.arrive() bars.mb_k_restore_acc_done[k_restore_index.idx].wait(k_restore_index.phase) k_restore_index = advance(k_restore_index, cfg.smem_decay_stages) @@ -1774,9 +1645,9 @@ def compute1_warp_group( if cutlass.const_expr(cfg.enable_checkpoints): cg1_checkpoint_chunks = checkpoint_every_n_tokens // cutlass.Int32(cfg.b_t) - cg1_checkpoint_mod = (cstart + cutlass.Int32(1)) % cg1_checkpoint_chunks + cg1_checkpoint_mod = (compute_start + cutlass.Int32(1)) % cg1_checkpoint_chunks for local_chunk in cutlass.range(1, num_tile_chunks, 1, unroll=1): - chunk_idx = cstart + local_chunk + chunk_idx = compute_start + local_chunk global_chunk = global_chunk_base + local_chunk sV_ptr = sV_raw.data_ptr() + raw_index.idx * (cfg.d_v * cfg.b_t) sW_ptr = sW_raw.data_ptr() + raw_index.idx * (cfg.d_v * cfg.b_t) @@ -1791,44 +1662,41 @@ def compute1_warp_group( do_checkpoint = cg1_checkpoint_mod == 0 cg1_checkpoint_mod = cg1_checkpoint_mod + cutlass.Int32(1) cg1_checkpoint_mod = cutlass.Int32(0) if cg1_checkpoint_mod == cg1_checkpoint_chunks else cg1_checkpoint_mod - do_checkpoint = do_checkpoint and chunk_idx >= wstart - row_addr = (tmem_row + tmem_subpartition * cfg.threads_per_warp) << 16 + do_checkpoint = do_checkpoint and chunk_idx >= write_start state_col_id = tmem_col + cfg.tmem_state_acc_offset - # ---- state repack: acc TMEM -> packed b16 TMEM ---------------------- + # ---- state stage: acc TMEM -> packed b16 TMEM ---------------------------- state_vecs = [] - for k_block in cutlass.range_constexpr(cfg.d_k // 16): - state_vecs.append(nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + state_col_id + k_block * 16, cutlass.Float32), num=16)) + for i in cutlass.range_constexpr(cfg.d_k // 16): + state_vecs.append(nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_lo_addr + state_col_id + i * 16, cutlass.Float32), num=16)) - packed_col_id = tmem_col + cfg.tmem_state_inp_offset - for k_block in cutlass.range_constexpr(cfg.d_k // 16): + packed_col_id = tmem_col + cfg.tmem_state_input_offset + for i in cutlass.range_constexpr(cfg.d_k // 16): state_pack = cutlass.Array(cutlass.Int32, 8, alignment=16) for packed_col in cutlass.range_constexpr(8): - source_pair = packed_col ^ 4 - state_pack[packed_col] = fp32_to_fp16(state_vecs[k_block][2 * source_pair], state_vecs[k_block][2 * source_pair + 1], dtype=cfg.io_dtype) + state_pack[packed_col] = fp32_to_fp16(state_vecs[i][2 * packed_col], state_vecs[i][2 * packed_col + 1], dtype=cfg.io_dtype) nvvm.tcgen05_st( "32x32b", - nvvm.make_tmem_ptr((tmem_row << 16) + packed_col_id + k_block * 8, cutlass.Int8), + nvvm.make_tmem_ptr(row_lo_addr + packed_col_id + i * 8, cutlass.Int8), state_pack[0:8], ) nvvm.tcgen05_wait("store") - bars.mb_state_inp_ready.arrive() + bars.mb_state_input_ready.arrive() - # ---- checkpoint store ----------------------------------------------- + # ---- checkpoint store ---------------------------------------------------- if cutlass.const_expr(cfg.enable_checkpoints): if do_checkpoint: checkpoint_stage = checkpoint_done_index.idx bars.mb_checkpoint_tmastg_done[checkpoint_stage].wait(checkpoint_done_index.phase) checkpoint_done_index = advance(checkpoint_done_index, cfg.smem_checkpoint_stages) checkpoint_stage_base = checkpoint_stage * (cfg.d_k * cfg.d_v) - row16_addr = ((tmem_row + tmem_subpartition * cfg.threads_per_warp) + 16) << 16 checkpoint_vbase = tmem_subpartition * cfg.threads_per_warp - checkpoint_swz_off0 = (checkpoint_vbase + frag_col_offset) // 64 * (cfg.d_k * 64) - checkpoint_swz_col0 = (checkpoint_vbase + frag_col_offset) % 64 - checkpoint_swz_off = (checkpoint_vbase + 16 + frag_col_offset) // 64 * (cfg.d_k * 64) - checkpoint_swz_col = (checkpoint_vbase + 16 + frag_col_offset) % 64 + checkpoint_swz_off_lo = (checkpoint_vbase + ov_col_offset) // 64 * (cfg.d_k * 64) + checkpoint_swz_col_lo = (checkpoint_vbase + ov_col_offset) % 64 + checkpoint_swz_off_hi = (checkpoint_vbase + 16 + ov_col_offset) // 64 * (cfg.d_k * 64) + checkpoint_swz_col_hi = (checkpoint_vbase + 16 + ov_col_offset) % 64 for k_base in cutlass.range_constexpr(0, cfg.d_k, 32): - checkpoint_vec = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + state_col_id + k_base, cutlass.Float32), num=32) + checkpoint_vec = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_lo_addr + state_col_id + k_base, cutlass.Float32), num=32) for g in cutlass.range_constexpr(4): packs = tuple(fp32_to_fp16(checkpoint_vec[g * 8 + 2 * t], checkpoint_vec[g * 8 + 2 * t + 1], dtype=cfg.io_dtype) for t in range(4)) dk = k_base + g * 8 @@ -1847,38 +1715,36 @@ def compute1_warp_group( projection_col_id = tmem_col + cfg.tmem_q_state_acc_offset + prev_q_state_acc_stage * cfg.b_t value_dim_base = tmem_subpartition * cfg.threads_per_warp - row_id0 = tmem_row + value_dim_base - row_id1 = row_id0 + 16 - loaded_vec0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id0 << 16) + projection_col_id, cutlass.Float32), num=2) - loaded_vec1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id1 << 16) + projection_col_id, cutlass.Float32), num=2) + loaded_vec_lo = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_lo_addr + projection_col_id, cutlass.Float32), num=2) + loaded_vec_hi = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_hi_addr + projection_col_id, cutlass.Float32), num=2) - # ---- output drain: O acc TMEM -> scaled b16 SMEM -------------------- - stsm_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) - stsm_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + # ---- output store: O acc TMEM -> scaled b16 SMEM ------------------------- + o_pack_lo = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + o_pack_hi = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) for reg_idx in cutlass.range_constexpr(4): - scaled0_0, scaled0_1 = fmul2(loaded_vec0[2 * reg_idx], loaded_vec0[2 * reg_idx + 1], scale, scale) - scaled1_0, scaled1_1 = fmul2(loaded_vec1[2 * reg_idx], loaded_vec1[2 * reg_idx + 1], scale, scale) - stsm_pack0[reg_idx] = fp32_to_fp16(scaled0_0, scaled0_1, dtype=mO.element_type) - stsm_pack1[reg_idx] = fp32_to_fp16(scaled1_0, scaled1_1, dtype=mO.element_type) + scaled0_0, scaled0_1 = fmul2(loaded_vec_lo[2 * reg_idx], loaded_vec_lo[2 * reg_idx + 1], scale, scale) + scaled1_0, scaled1_1 = fmul2(loaded_vec_hi[2 * reg_idx], loaded_vec_hi[2 * reg_idx + 1], scale, scale) + o_pack_lo[reg_idx] = fp32_to_fp16(scaled0_0, scaled0_1, dtype=mO.element_type) + o_pack_hi[reg_idx] = fp32_to_fp16(scaled1_0, scaled1_1, dtype=mO.element_type) bars.mb_o_tmastg_done[prev_o_stage].wait(((prev_global_chunk // cfg.smem_o_stages) + 1) % 2) nvvm.stmatrix( sO_ptr + prev_o_stage_base - + (value_dim_base + frag_col_offset) // 64 * (cfg.b_t * 64) - + frag_row_coord * 64 - + swizzle_xor_128b(frag_row_coord, (value_dim_base + frag_col_offset) % 64, elem_bytes=2), - stsm_pack0.data_ptr().load(count=4, alignment=4), + + (value_dim_base + ov_col_offset) // 64 * (cfg.b_t * 64) + + ov_row_coord * 64 + + swizzle_xor_128b(ov_row_coord, (value_dim_base + ov_col_offset) % 64, elem_bytes=2), + o_pack_lo.data_ptr().load(count=4, alignment=4), nvvm.MMALayout.COL, shape=nvvm.StoreShape.M8N8, ) nvvm.stmatrix( sO_ptr + prev_o_stage_base - + (value_dim_base + 16 + frag_col_offset) // 64 * (cfg.b_t * 64) - + frag_row_coord * 64 - + swizzle_xor_128b(frag_row_coord, (value_dim_base + 16 + frag_col_offset) % 64, elem_bytes=2), - stsm_pack1.data_ptr().load(count=4, alignment=4), + + (value_dim_base + 16 + ov_col_offset) // 64 * (cfg.b_t * 64) + + ov_row_coord * 64 + + swizzle_xor_128b(ov_row_coord, (value_dim_base + 16 + ov_col_offset) % 64, elem_bytes=2), + o_pack_hi.data_ptr().load(count=4, alignment=4), nvvm.MMALayout.COL, shape=nvvm.StoreShape.M8N8, ) @@ -1886,120 +1752,115 @@ def compute1_warp_group( nvvm.fence_proxy("async.shared", space="cta") bars.mb_o_tmastg_ready[prev_o_stage].arrive() - # ---- Y staging: Y = W*V - state*(Beta*K) ----------------------------- + # ---- Y stage: Y = W*V - state*(Beta*K) ----------------------------------- bars.mb_v_ready[raw_bar_index.idx].wait(raw_bar_index.phase) projection_col_id = tmem_col + cfg.tmem_state_k_acc_offset - input_col_id = tmem_col + cfg.tmem_y_inp_offset + input_col_id = tmem_col + cfg.tmem_y_input_offset value_dim_base = tmem_subpartition * cfg.threads_per_warp - # ---- raw V fragments, then W, then the state*K acc readback ---------- - row_id0 = tmem_row + value_dim_base - row_id1 = row_id0 + 16 - raw_v_frag0 = nvvm.ldmatrix( + # ---- raw V fragments, then W, then the k state acc read ------------------ + raw_v_frag_lo = nvvm.ldmatrix( sV_ptr - + (value_dim_base + frag_col_offset) // 64 * (cfg.b_t * 64) - + frag_row_coord * 64 - + swizzle_xor_128b(frag_row_coord, (value_dim_base + frag_col_offset) % 64, elem_bytes=2), + + (value_dim_base + ov_col_offset) // 64 * (cfg.b_t * 64) + + ov_row_coord * 64 + + swizzle_xor_128b(ov_row_coord, (value_dim_base + ov_col_offset) % 64, elem_bytes=2), 4, nvvm.MMALayout.COL, ) - raw_v_frag1 = nvvm.ldmatrix( + raw_v_frag_hi = nvvm.ldmatrix( sV_ptr - + (value_dim_base + 16 + frag_col_offset) // 64 * (cfg.b_t * 64) - + frag_row_coord * 64 - + swizzle_xor_128b(frag_row_coord, (value_dim_base + 16 + frag_col_offset) % 64, elem_bytes=2), + + (value_dim_base + 16 + ov_col_offset) // 64 * (cfg.b_t * 64) + + ov_row_coord * 64 + + swizzle_xor_128b(ov_row_coord, (value_dim_base + 16 + ov_col_offset) % 64, elem_bytes=2), 4, nvvm.MMALayout.COL, ) bars.mb_w_ready[raw_bar_index.idx].wait(raw_bar_index.phase) - raw_w_frag0 = nvvm.ldmatrix( + raw_w_frag_lo = nvvm.ldmatrix( sW_ptr - + (value_dim_base + frag_col_offset) // 64 * (cfg.b_t * 64) - + frag_row_coord * 64 - + swizzle_xor_128b(frag_row_coord, (value_dim_base + frag_col_offset) % 64, elem_bytes=2), + + (value_dim_base + ov_col_offset) // 64 * (cfg.b_t * 64) + + ov_row_coord * 64 + + swizzle_xor_128b(ov_row_coord, (value_dim_base + ov_col_offset) % 64, elem_bytes=2), 4, nvvm.MMALayout.COL, ) - raw_w_frag1 = nvvm.ldmatrix( + raw_w_frag_hi = nvvm.ldmatrix( sW_ptr - + (value_dim_base + 16 + frag_col_offset) // 64 * (cfg.b_t * 64) - + frag_row_coord * 64 - + swizzle_xor_128b(frag_row_coord, (value_dim_base + 16 + frag_col_offset) % 64, elem_bytes=2), + + (value_dim_base + 16 + ov_col_offset) // 64 * (cfg.b_t * 64) + + ov_row_coord * 64 + + swizzle_xor_128b(ov_row_coord, (value_dim_base + 16 + ov_col_offset) % 64, elem_bytes=2), 4, nvvm.MMALayout.COL, ) bars.mb_state_k_acc_ready.wait(state_k_acc_index.phase) - state_k_vec0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id0 << 16) + projection_col_id, cutlass.Float32), num=2) - state_k_vec1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id1 << 16) + projection_col_id, cutlass.Float32), num=2) + state_k_vec_lo = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_lo_addr + projection_col_id, cutlass.Float32), num=2) + state_k_vec_hi = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_hi_addr + projection_col_id, cutlass.Float32), num=2) - y_inp_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + y_input_pack_lo = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) for reg_idx in cutlass.range_constexpr(4): - raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1) - frag_pair = (reg_idx ^ 2) * 2 - state_k_val0, state_k_val1 = state_k_vec0[frag_pair], state_k_vec0[frag_pair + 1] + frag_pair = reg_idx * 2 + state_k_val0, state_k_val1 = state_k_vec_lo[frag_pair], state_k_vec_lo[frag_pair + 1] state_k_pair = fp32_to_fp16(state_k_val0, state_k_val1, dtype=cfg.io_dtype) - w_pair = raw_w_frag0[raw_matrix] + w_pair = raw_w_frag_lo[reg_idx] wv_pair = mul_f16x2( w_pair, - raw_v_frag0[raw_matrix], + raw_v_frag_lo[reg_idx], cfg.io_dtype, ) - y_inp_pack0[reg_idx] = sub_f16x2( + y_input_pack_lo[reg_idx] = sub_f16x2( wv_pair, state_k_pair, cfg.io_dtype, ) - y_inp_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + y_input_pack_hi = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) for reg_idx in cutlass.range_constexpr(4): - raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1) - frag_pair = (reg_idx ^ 2) * 2 - state_k_val0, state_k_val1 = state_k_vec1[frag_pair], state_k_vec1[frag_pair + 1] + frag_pair = reg_idx * 2 + state_k_val0, state_k_val1 = state_k_vec_hi[frag_pair], state_k_vec_hi[frag_pair + 1] state_k_pair = fp32_to_fp16(state_k_val0, state_k_val1, dtype=cfg.io_dtype) - w_pair = raw_w_frag1[raw_matrix] + w_pair = raw_w_frag_hi[reg_idx] wv_pair = mul_f16x2( w_pair, - raw_v_frag1[raw_matrix], + raw_v_frag_hi[reg_idx], cfg.io_dtype, ) - y_inp_pack1[reg_idx] = sub_f16x2( + y_input_pack_hi[reg_idx] = sub_f16x2( wv_pair, state_k_pair, cfg.io_dtype, ) - nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr((tmem_row << 16) + input_col_id, cutlass.Int8), y_inp_pack0[0:4]) - nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr((tmem_row + 16 << 16) + input_col_id, cutlass.Int8), y_inp_pack1[0:4]) + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_lo_addr + input_col_id, cutlass.Int8), y_input_pack_lo[0:4]) + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_hi_addr + input_col_id, cutlass.Int8), y_input_pack_hi[0:4]) nvvm.tcgen05_wait("store") state_k_acc_index = advance(state_k_acc_index, 1) bars.mb_v_done[raw_index.idx].arrive() bars.mb_w_done[raw_index.idx].arrive() - bars.mb_y_inp_ready.arrive() + bars.mb_y_input_ready.arrive() - # ---- U repack: acc TMEM -> packed b16 TMEM -------------------------- + # ---- U stage: acc TMEM -> packed b16 TMEM -------------------------------- bars.mb_u_acc_ready.wait(u_acc_index.phase) u_acc_vals = nvvm.tcgen05_ld( "32x32b", - nvvm.make_tmem_ptr((tmem_row + tmem_subpartition * cfg.threads_per_warp << 16) + (tmem_col + cfg.tmem_u_acc_offset), cutlass.Float32), + nvvm.make_tmem_ptr(row_lo_addr + (tmem_col + cfg.tmem_u_acc_offset), cutlass.Float32), num=cfg.b_t, ) - u_inp_pack = cutlass.Array(cutlass.Int32, (cfg.b_t // 2), alignment=16) + u_input_pack = cutlass.Array(cutlass.Int32, (cfg.b_t // 2), alignment=16) for packed_col in cutlass.range_constexpr((cfg.b_t // 2)): - source_pair = packed_col ^ 4 - token0 = source_pair * 2 + token0 = packed_col * 2 token1 = token0 + 1 - u_inp_pack[packed_col] = fp32_to_fp16(u_acc_vals[token0], u_acc_vals[token1], dtype=cfg.io_dtype) + u_input_pack[packed_col] = fp32_to_fp16(u_acc_vals[token0], u_acc_vals[token1], dtype=cfg.io_dtype) nvvm.tcgen05_st( "32x32b", - nvvm.make_tmem_ptr((tmem_row << 16) + (tmem_col + cfg.tmem_u_inp_offset), cutlass.Int8), - u_inp_pack[0 : (cfg.b_t // 2)], + nvvm.make_tmem_ptr(row_lo_addr + (tmem_col + cfg.tmem_u_input_offset), cutlass.Int8), + u_input_pack[0 : (cfg.b_t // 2)], ) nvvm.tcgen05_wait("store") u_acc_index = advance(u_acc_index, 1) - bars.mb_u_inp_ready.arrive() + bars.mb_u_input_ready.arrive() bars.mb_k_restore_acc_done[k_restore_index.idx].wait(k_restore_index.phase) k_restore_index = advance(k_restore_index, cfg.smem_decay_stages) @@ -2008,7 +1869,7 @@ def compute1_warp_group( if num_tile_chunks > 0: last_global_chunk = global_chunk_base + num_tile_chunks - cutlass.Int32(1) - output_chunk = wend - cutlass.Int32(1) + output_chunk = write_end - cutlass.Int32(1) final_o_stage = last_global_chunk % cfg.smem_o_stages final_q_state_acc_stage = last_global_chunk % cfg.tmem_q_state_acc_stages final_o_stage_base = final_o_stage * (cfg.b_t * cfg.d_v) @@ -2018,38 +1879,36 @@ def compute1_warp_group( projection_col_id = tmem_col + cfg.tmem_q_state_acc_offset + final_q_state_acc_stage * cfg.b_t value_dim_base = tmem_subpartition * cfg.threads_per_warp - row_id0 = tmem_row + value_dim_base - row_id1 = row_id0 + 16 - loaded_vec0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id0 << 16) + projection_col_id, cutlass.Float32), num=2) - loaded_vec1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id1 << 16) + projection_col_id, cutlass.Float32), num=2) + loaded_vec_lo = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_lo_addr + projection_col_id, cutlass.Float32), num=2) + loaded_vec_hi = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_hi_addr + projection_col_id, cutlass.Float32), num=2) - # ---- output drain: O acc TMEM -> scaled b16 SMEM -------------------- - stsm_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) - stsm_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + # ---- output store: O acc TMEM -> scaled b16 SMEM ------------------------- + o_pack_lo = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + o_pack_hi = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) for reg_idx in cutlass.range_constexpr(4): - scaled0_0, scaled0_1 = fmul2(loaded_vec0[2 * reg_idx], loaded_vec0[2 * reg_idx + 1], scale, scale) - scaled1_0, scaled1_1 = fmul2(loaded_vec1[2 * reg_idx], loaded_vec1[2 * reg_idx + 1], scale, scale) - stsm_pack0[reg_idx] = fp32_to_fp16(scaled0_0, scaled0_1, dtype=mO.element_type) - stsm_pack1[reg_idx] = fp32_to_fp16(scaled1_0, scaled1_1, dtype=mO.element_type) + scaled0_0, scaled0_1 = fmul2(loaded_vec_lo[2 * reg_idx], loaded_vec_lo[2 * reg_idx + 1], scale, scale) + scaled1_0, scaled1_1 = fmul2(loaded_vec_hi[2 * reg_idx], loaded_vec_hi[2 * reg_idx + 1], scale, scale) + o_pack_lo[reg_idx] = fp32_to_fp16(scaled0_0, scaled0_1, dtype=mO.element_type) + o_pack_hi[reg_idx] = fp32_to_fp16(scaled1_0, scaled1_1, dtype=mO.element_type) bars.mb_o_tmastg_done[final_o_stage].wait(((last_global_chunk // cfg.smem_o_stages) + 1) % 2) nvvm.stmatrix( sO_ptr + final_o_stage_base - + (value_dim_base + frag_col_offset) // 64 * (cfg.b_t * 64) - + frag_row_coord * 64 - + swizzle_xor_128b(frag_row_coord, (value_dim_base + frag_col_offset) % 64, elem_bytes=2), - stsm_pack0.data_ptr().load(count=4, alignment=4), + + (value_dim_base + ov_col_offset) // 64 * (cfg.b_t * 64) + + ov_row_coord * 64 + + swizzle_xor_128b(ov_row_coord, (value_dim_base + ov_col_offset) % 64, elem_bytes=2), + o_pack_lo.data_ptr().load(count=4, alignment=4), nvvm.MMALayout.COL, shape=nvvm.StoreShape.M8N8, ) nvvm.stmatrix( sO_ptr + final_o_stage_base - + (value_dim_base + 16 + frag_col_offset) // 64 * (cfg.b_t * 64) - + frag_row_coord * 64 - + swizzle_xor_128b(frag_row_coord, (value_dim_base + 16 + frag_col_offset) % 64, elem_bytes=2), - stsm_pack1.data_ptr().load(count=4, alignment=4), + + (value_dim_base + 16 + ov_col_offset) // 64 * (cfg.b_t * 64) + + ov_row_coord * 64 + + swizzle_xor_128b(ov_row_coord, (value_dim_base + 16 + ov_col_offset) % 64, elem_bytes=2), + o_pack_hi.data_ptr().load(count=4, alignment=4), nvvm.MMALayout.COL, shape=nvvm.StoreShape.M8N8, ) @@ -2057,18 +1916,18 @@ def compute1_warp_group( nvvm.fence_proxy("async.shared", space="cta") bars.mb_o_tmastg_ready[final_o_stage].arrive() - owns_final = wend == num_chunks_b + owns_final = write_end == batch_num_chunks - # ---- final-state drain: state acc TMEM -> GMEM --------------------------- + # ---- final state store: TMEM -> GMEM ----------------------------------------- if cutlass.const_expr(mState_out is not None): - if seqlen_b > 0: + if batch_seqlen > 0: if owns_final: state_vw = 16 // (mState_out.element_type.width // 8) state_dst = (mState_out.iterator + mState_out.layout((batch_idx, head_o, value_dim, 0))).raw_ptr() for key_block_start in cutlass.range_constexpr(0, cfg.d_k, 32): loaded = nvvm.tcgen05_ld( "32x32b", - nvvm.make_tmem_ptr((row_id << 16) + (tmem_col + cfg.tmem_state_acc_offset + key_block_start), cutlass.Float32), + nvvm.make_tmem_ptr(row_lo_addr + (tmem_col + cfg.tmem_state_acc_offset + key_block_start), cutlass.Float32), num=32, ) @@ -2089,7 +1948,7 @@ def compute1_warp_group( else: mState_out[batch_idx, head_o, value_dim, key_dim] = cutlass.Float32(0.0).to(mState_out.element_type) global_chunk_base += num_tile_chunks - tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one) bars.mb_tmem_done[0].arrive() @@ -2111,7 +1970,7 @@ def host( final_state: cute.Tensor | None, work_items: cute.Tensor | None, work_count: cute.Tensor | None, - sched_ctr: cute.Tensor | None, + scheduler_counter: cute.Tensor | None, tensormap_workspace: cute.Tensor, checkpoint_every_n_tokens: cutlass.Int32, scale: cutlass.Float32, @@ -2138,7 +1997,7 @@ def host( final_state, work_items, work_count, - sched_ctr, + scheduler_counter, scale, checkpoint_every_n_tokens, ).launch( @@ -2168,7 +2027,7 @@ def kernel( mState_out: cute.Tensor | None, mWorkItems: cute.Tensor, mCount: cute.Tensor, - mSched: cute.Tensor | None, + mScheduler: cute.Tensor | None, scale: cutlass.Float32, checkpoint_every_n_tokens: cutlass.Int32, ) -> None: @@ -2184,11 +2043,11 @@ def kernel( bidx = cute.arch.block_idx()[0] num_ctas = cute.arch.grid_dim()[0] warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - lane = tidx % cfg.threads_per_warp + lane_idx = tidx % cfg.threads_per_warp total_tiles = mCount[0] - if cutlass.const_expr(cfg.dyn_sched): - assert mSched is not None and mSched.element_type == cutlass.Int32 + if cutlass.const_expr(cfg.dynamic_scheduling): + assert mScheduler is not None and mScheduler.element_type == cutlass.Int32 assert mQ.element_type == cfg.io_dtype and mK.element_type == cfg.io_dtype and mV.element_type == cfg.io_dtype assert mGate.element_type == cutlass.Float32 assert mBeta.element_type == cfg.io_dtype and mW.element_type == cfg.io_dtype, "channel-wise beta/w must match the io dtype" @@ -2215,11 +2074,10 @@ def kernel( desc_o_base = desc_base_words + cutlass.Int32(6) * arr_words desc_checkpoint_base = desc_base_words + cutlass.Int32(7) * arr_words - # Buffers are declaration-ordered and intentionally non-aliased. SMEM = cutlass.AddressSpace.smem bars = make_gdn2_bars(cfg) sTmem_base = cutlass.Array(cutlass.Int32, 1, space=SMEM, alignment=4) - sSched = cutlass.Array(cutlass.Int32, cfg.sched_stages, space=SMEM, alignment=16) + sScheduler = cutlass.Array(cutlass.Int32, cfg.scheduler_stages, space=SMEM, alignment=16) sK_decay_raw = cutlass.Array(cfg.io_dtype, cfg.k_decay_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) sQ_decay_raw = cutlass.Array(cfg.io_dtype, cfg.q_decay_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) sK_restore_raw = cutlass.Array(cfg.io_dtype, cfg.k_restore_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) @@ -2259,7 +2117,7 @@ def kernel( stride_byte_offset=1024, layout=nvvm.Tcgen05SmemSwizzle.SWIZZLE_128B, ) - sK_restore = SmemTile( + sK_restore_trans = SmemTile( base=sK_restore_raw, elems_per_stage=(cfg.d_k * cfg.b_t), stages=cfg.smem_decay_stages, @@ -2308,15 +2166,15 @@ def kernel( bars.mb_o_acc_done[stage].init() bars.mb_state_k_acc_ready.init() bars.mb_u_acc_ready.init() - bars.mb_state_inp_ready.init() + bars.mb_state_input_ready.init() for stage in cutlass.range_constexpr(cfg.smem_state_scale_diag_stages): bars.mb_state_scale_diag_done[stage].init() for stage in cutlass.range_constexpr(cfg.smem_decay_stages): bars.mb_decay_tcgen05_done[stage].init() bars.mb_decay_super_done[stage].init() bars.mb_k_restore_acc_done[stage].init() - bars.mb_y_inp_ready.init() - bars.mb_u_inp_ready.init() + bars.mb_y_input_ready.init() + bars.mb_u_input_ready.init() bars.mb_tmem_done[0].init() elif warp_idx == cfg.super_mma_warp_id: if elect_one: @@ -2333,9 +2191,9 @@ def kernel( for stage in cutlass.range_constexpr(cfg.smem_o_stages): bars.mb_o_tmastg_ready[stage].init() bars.mb_o_tmastg_done[stage].init() - for stage in cutlass.range_constexpr(cfg.sched_stages): - bars.mb_sched_ready[stage].init() - bars.mb_sched_done[stage].init() + for stage in cutlass.range_constexpr(cfg.scheduler_stages): + bars.mb_scheduler_ready[stage].init() + bars.mb_scheduler_done[stage].init() if cutlass.const_expr(cfg.enable_checkpoints): for stage in cutlass.range_constexpr(cfg.smem_checkpoint_stages): bars.mb_checkpoint_tmastg_ready[stage].init() @@ -2354,9 +2212,9 @@ def kernel( num_ctas, cu_seqlens, mWorkItems, - mSched, - sSched, - lane, + mScheduler, + sScheduler, + lane_idx, sBeta_raw, sGate_raw, sK_raw, @@ -2379,8 +2237,8 @@ def kernel( num_ctas, cu_seqlens, mWorkItems, - sSched, - lane, + sScheduler, + lane_idx, sK_inv_raw, sIntermediate_raw, sK_decay_raw, @@ -2394,11 +2252,11 @@ def kernel( num_ctas, cu_seqlens, mWorkItems, - sSched, + sScheduler, sTmem_base, sIntermediate, sK_decay, - sK_restore, + sK_restore_trans, sQ_decay, sState_scale_diag, bars, @@ -2411,8 +2269,8 @@ def kernel( num_ctas, cu_seqlens, mWorkItems, - sSched, - lane, + sScheduler, + lane_idx, mO, sK_inv_raw, sO_raw, @@ -2432,8 +2290,8 @@ def kernel( num_ctas, cu_seqlens, mWorkItems, - sSched, - lane, + sScheduler, + lane_idx, warp_idx, mQ, mA_log, @@ -2459,8 +2317,8 @@ def kernel( num_ctas, cu_seqlens, mWorkItems, - sSched, - lane, + sScheduler, + lane_idx, sTmem_base, warp_idx, mState_out, @@ -2497,8 +2355,8 @@ class Gdn2Cfg: v_ratio: int n_heads_out: int max_active_clusters: int - dyn_sched: bool = False - sched_stages: int = CFG.SMEM_SCHED_STAGES + dynamic_scheduling: bool = False + scheduler_stages: int = CFG.SMEM_SCHEDULER_STAGES compute_group_0_warp_ids: tuple[int, ...] = CFG.COMPUTE_GROUP_0_WARP_IDS compute_group_1_warp_ids: tuple[int, ...] = CFG.COMPUTE_GROUP_1_WARP_IDS @@ -2515,7 +2373,7 @@ class Gdn2Cfg: cg0_group_count: int = 2 cg0_warps_per_group: int = 4 cg0_threads_per_group: int = 0 - cg0_group_sync_barrier_base_id: int = 1 # CG0 group g syncs on nbar id 1 + g + cg0_group_sync_barrier_base_id: int = 1 # CG0 group g syncs on named-barrier id 1 + g cg0_tile_entry_barrier_id: int = 5 # CG0-wide (both groups) work-item entry sync tmem_user_threads: int = 0 tmem_lifecycle_barrier_id: int = 3 @@ -2523,9 +2381,9 @@ class Gdn2Cfg: num_regs_compute_group_1: int = CFG.NUM_REGS_COMPUTE_GROUP_1 num_regs_other: int = CFG.NUM_REGS_OTHER - # ---- SMEM / TMEM ring stage counts ------------------------------------------- + # ---- SMEM / TMEM ring stage counts ----------------------------------------------- smem_raw_stages: int = CFG.SMEM_RAW_STAGES - smem_raw_bar_stages: int = 0 # ready-ring mbar depth: raw rounded up to even (CG0 ping-pong parity) + smem_raw_bar_stages: int = 0 # ready-ring mbar depth: raw rounded up to even smem_checkpoint_stages: int = 1 smem_o_stages: int = CFG.SMEM_O_STAGES smem_decay_stages: int = CFG.SMEM_DECAY_STAGES @@ -2534,16 +2392,16 @@ class Gdn2Cfg: qk_scale_ready_stages: int = CFG.QK_SCALE_READY_STAGES tmem_q_state_acc_stages: int = CFG.TMEM_Q_STATE_ACC_STAGES - # ---- TMEM column offsets (state doubles as the final_state acc) -------------- + # ---- TMEM column offsets (state doubles as the final state acc) ------------------ tmem_state_acc_offset: int = 0 - tmem_state_inp_offset: int = 0 + tmem_state_input_offset: int = 0 tmem_q_state_acc_offset: int = 0 tmem_state_k_acc_offset: int = 0 tmem_u_acc_offset: int = 0 - tmem_y_inp_offset: int = 0 - tmem_u_inp_offset: int = 0 + tmem_y_input_offset: int = 0 + tmem_u_input_offset: int = 0 - # ---- SMEM buffer cosizes ----------------------------------------------------- + # ---- SMEM buffer cosizes --------------------------------------------------------- q_cosize: int = 0 k_cosize: int = 0 v_cosize: int = 0 @@ -2583,7 +2441,7 @@ def build_cfg( v_ratio: int, n_heads_out: int, max_active_clusters: int, - dyn_sched: bool = False, + dynamic_scheduling: bool = False, ) -> Gdn2Cfg: """Build the per-compile ``Gdn2Cfg`` (io_dtype in {Float16, BFloat16}); fills the derived TMEM column offsets and SMEM buffer cosizes.""" @@ -2604,7 +2462,7 @@ def build_cfg( v_ratio=v_ratio, n_heads_out=n_heads_out, max_active_clusters=max_active_clusters, - dyn_sched=dyn_sched, + dynamic_scheduling=dynamic_scheduling, ) if enable_checkpoints: cfg.smem_raw_stages = 3 @@ -2616,13 +2474,13 @@ def build_cfg( if cfg.smem_state_scale_diag_stages != cfg.qk_scale_ready_stages: raise ValueError("diag and qk-scale ready rings must share their rolling stage") - cfg.tmem_state_inp_offset = cfg.tmem_state_acc_offset + cfg.d_k - cfg.tmem_q_state_acc_offset = cfg.tmem_state_inp_offset + (cfg.d_k // 2) + cfg.tmem_state_input_offset = cfg.tmem_state_acc_offset + cfg.d_k + cfg.tmem_q_state_acc_offset = cfg.tmem_state_input_offset + (cfg.d_k // 2) cfg.tmem_state_k_acc_offset = cfg.tmem_q_state_acc_offset + cfg.tmem_q_state_acc_stages * cfg.b_t cfg.tmem_u_acc_offset = cfg.tmem_state_k_acc_offset + cfg.b_t - cfg.tmem_y_inp_offset = cfg.tmem_u_acc_offset + cfg.b_t - cfg.tmem_u_inp_offset = cfg.tmem_y_inp_offset + (cfg.b_t // 2) - assert (cfg.tmem_u_inp_offset + (cfg.b_t // 2)) <= 512 + cfg.tmem_y_input_offset = cfg.tmem_u_acc_offset + cfg.b_t + cfg.tmem_u_input_offset = cfg.tmem_y_input_offset + (cfg.b_t // 2) + assert (cfg.tmem_u_input_offset + (cfg.b_t // 2)) <= 512 cfg.q_cosize = cfg.smem_raw_stages * cfg.d_k * cfg.b_t cfg.k_cosize = cfg.smem_raw_stages * cfg.d_k * cfg.b_t @@ -2661,7 +2519,7 @@ def build_descs_body( base_w, base_o, base_checkpoint, - desc_ws: cute.Tensor, + desc_workspace: cute.Tensor, cu_seqlens: cute.Tensor, q: cute.Tensor, k: cute.Tensor, @@ -2689,14 +2547,14 @@ def build_descs_body( main kernel's CTA 0 prologue (warps past the array count fall through the widx guards).""" arr_words = n_batch * cutlass.Int32(TENSOR_MAP_QWORDS) - sub0 = cute.make_tensor(desc_ws.iterator, cute.make_layout((arr_words,), stride=(1,))) - sub1 = cute.make_tensor(desc_ws.iterator + arr_words, cute.make_layout((arr_words,), stride=(1,))) - sub2 = cute.make_tensor(desc_ws.iterator + 2 * arr_words, cute.make_layout((arr_words,), stride=(1,))) - sub3 = cute.make_tensor(desc_ws.iterator + 3 * arr_words, cute.make_layout((arr_words,), stride=(1,))) - sub4 = cute.make_tensor(desc_ws.iterator + 4 * arr_words, cute.make_layout((arr_words,), stride=(1,))) - sub5 = cute.make_tensor(desc_ws.iterator + 5 * arr_words, cute.make_layout((arr_words,), stride=(1,))) - sub6 = cute.make_tensor(desc_ws.iterator + 6 * arr_words, cute.make_layout((arr_words,), stride=(1,))) - sub7 = cute.make_tensor(desc_ws.iterator + 7 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub0 = cute.make_tensor(desc_workspace.iterator, cute.make_layout((arr_words,), stride=(1,))) + sub1 = cute.make_tensor(desc_workspace.iterator + arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub2 = cute.make_tensor(desc_workspace.iterator + 2 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub3 = cute.make_tensor(desc_workspace.iterator + 3 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub4 = cute.make_tensor(desc_workspace.iterator + 4 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub5 = cute.make_tensor(desc_workspace.iterator + 5 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub6 = cute.make_tensor(desc_workspace.iterator + 6 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub7 = cute.make_tensor(desc_workspace.iterator + 7 * arr_words, cute.make_layout((arr_words,), stride=(1,))) if widx == 0: if nvvm.elect_sync(): @@ -2736,7 +2594,7 @@ def build_descs_body( @cute.kernel def prologue_kernel( order_gen: cutlass.Constexpr[bool], - has_sched: cutlass.Constexpr[bool], + has_scheduler: cutlass.Constexpr[bool], b_t: cutlass.Constexpr[int], base_q: cutlass.GridConstant[cuda.tensor_map.TensorMap], base_k: cutlass.GridConstant[cuda.tensor_map.TensorMap], @@ -2746,7 +2604,7 @@ def prologue_kernel( base_w: cutlass.GridConstant[cuda.tensor_map.TensorMap], base_o: cutlass.GridConstant[cuda.tensor_map.TensorMap], base_checkpoint: cutlass.GridConstant[cuda.tensor_map.TensorMap], - desc_ws: cute.Tensor, + desc_workspace: cute.Tensor, cu_seqlens: cute.Tensor, q: cute.Tensor, k: cute.Tensor, @@ -2759,7 +2617,7 @@ def prologue_kernel( mStaging: cute.Tensor | None, mCount: cute.Tensor, mWorkItems: cute.Tensor, - mSched: cute.Tensor | None, + mScheduler: cute.Tensor | None, n_batch: cutlass.Int32, q_row_stride: cutlass.Int32, k_row_stride: cutlass.Int32, @@ -2771,7 +2629,7 @@ def prologue_kernel( checkpoint_row_stride: cutlass.Int32, checkpoint_every_n: cutlass.Int32, ) -> None: - """Single-CTA prologue: LPT-order the work-item table and zero the sched + """Single-CTA prologue: LPT-order the work-item table and zero the scheduler rings via :func:`order_body`, then build the per-batch TMA-descriptor arrays via :func:`build_descs_body`, one warp per array (the extra warps only take part in the order phase).""" @@ -2784,10 +2642,10 @@ def prologue_kernel( n_heads_out = cutlass.Int32(gate.shape[1]) order_body( order_gen, - has_sched, + has_scheduler, b_t, ORDER_THREADS, - ORDER_ELEMS, + ORDER_ELEMENTS, tidx, n_heads_out, n_heads_out * n_batch, @@ -2795,7 +2653,7 @@ def prologue_kernel( mStaging, mCount, mWorkItems, - mSched, + mScheduler, sKey, sIdx, sSpread, @@ -2810,7 +2668,7 @@ def prologue_kernel( base_w, base_o, base_checkpoint, - desc_ws, + desc_workspace, cu_seqlens, q, k, @@ -2838,7 +2696,7 @@ def prologue( io_dtype: cutlass.Constexpr, b_t: cutlass.Constexpr[int], order_gen: cutlass.Constexpr[bool], - has_sched: cutlass.Constexpr[bool], + has_scheduler: cutlass.Constexpr[bool], q: cute.Tensor, k: cute.Tensor, v: cute.Tensor, @@ -2851,7 +2709,7 @@ def prologue( work_item_staging: cute.Tensor | None, work_count: cute.Tensor, work_items: cute.Tensor, - sched_ctr: cute.Tensor | None, + scheduler_counter: cute.Tensor | None, tensormap_workspace: cute.Tensor, checkpoint_every_n: cutlass.Int32, stream: cuda_driver.CUstream, @@ -2910,7 +2768,7 @@ def prologue( base_checkpoint = cuda.create_tensor_map_tiled_from_view(checkpoint_view, box_dims=(tma_granu_elems, d_k, 1, 1), stride_order=(0, 1, 2, 3), swizzle=swz) prologue_kernel( order_gen, - has_sched, + has_scheduler, b_t, base_q, base_k, @@ -2933,7 +2791,7 @@ def prologue( work_item_staging, work_count, work_items, - sched_ctr, + scheduler_counter, cutlass.Int32(batch_size), cutlass.Int32(q.stride[0]), cutlass.Int32(k.stride[0]), @@ -2947,7 +2805,7 @@ def prologue( ).launch(grid=(1, 1, 1), block=(ORDER_THREADS, 1, 1), stream=stream) -# ---- Torch adapter / host-side compilation --------------------------------------- +# ---- Torch adapter / host-side compilation ------------------------------------------- @lru_cache(maxsize=None) @@ -2965,7 +2823,7 @@ def get_compiled_cache( safe_gate: bool, gate_lower_bound: float, beta_sigmoid: bool, - dyn_sched: bool, + dynamic_scheduling: bool, order_gen: bool, ): """Return a mutable dict that lazily stores the compiled kernel.""" @@ -2986,7 +2844,7 @@ def compile( k_ratio: int, v_ratio: int, n_heads_out: int, - dyn_sched: bool = False, + dynamic_scheduling: bool = False, *, num_sm: int, q_cute, @@ -3003,8 +2861,8 @@ def compile( state_out_cute, work_items_cute=None, work_count_cute=None, - sched_ctr_cute=None, - tensormap_ws_cute, + scheduler_counter_cute=None, + tensormap_workspace_cute, checkpoint_every_n_tokens, scale, stream, @@ -3025,7 +2883,7 @@ def compile( v_ratio=v_ratio, n_heads_out=n_heads_out, max_active_clusters=num_sm, - dyn_sched=dyn_sched, + dynamic_scheduling=dynamic_scheduling, ) return cute.compile( @@ -3045,8 +2903,8 @@ def compile( state_out_cute, work_items_cute, work_count_cute, - sched_ctr_cute, - tensormap_ws_cute, + scheduler_counter_cute, + tensormap_workspace_cute, checkpoint_every_n_tokens, scale, stream, @@ -3076,7 +2934,7 @@ def chunk_gdn2_sm100( use_beta_sigmoid: bool = False, work_items=None, work_count=None, - sched_ctr=None, + scheduler_counter=None, work_item_scratch=None, *, tensormap_workspace, @@ -3118,8 +2976,8 @@ def chunk_gdn2_sm100( use_beta_sigmoid: ``beta`` holds logits; sigmoid in-kernel work_items: ``(max_items, 8)`` int32 work-item table from ``common/split_k.py`` (REQUIRED; an uncut table row is the whole - (b, h) sequence). Each item computes chunks ``[cstart, wend)`` - and writes O/checkpoints only for ``[wstart, wend)``. + (b, h) sequence). Each item computes chunks ``[compute_start, write_end)`` + and writes O/checkpoints only for ``[write_start, write_end)``. work_count: ``(1,)`` int32 device-side item count (REQUIRED) """ HQ = q.shape[1] @@ -3137,8 +2995,10 @@ def chunk_gdn2_sm100( f"output_state_checkpoints dtype must match the io dtype (fp32 state belongs to output_state): got {output_state_checkpoints.dtype} with io {q.dtype}" ) if work_items is None or work_count is None: - raise ValueError("work_items/work_count are required (the split-table stage builds them for every launch)") - dyn_sched = sched_ctr is not None + raise ValueError( + "work_items/work_count are required (built by the split-table stage, or by an order-generating prologue when work_item_scratch is None)" + ) + dynamic_scheduling = scheduler_counter is not None order_gen = work_item_scratch is None if initial_state is not None: @@ -3177,7 +3037,7 @@ def chunk_gdn2_sm100( safe_gate, gate_lower_bound, use_beta_sigmoid, - dyn_sched, + dynamic_scheduling, order_gen, ) @@ -3207,11 +3067,11 @@ def chunk_gdn2_sm100( work_items_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1), divisibility=1) work_count_cute = from_dlpack(work_count, assumed_align=4).mark_layout_dynamic() - sched_ctr_cute = None - if dyn_sched: - sched_ctr_cute = from_dlpack(sched_ctr, assumed_align=4).mark_layout_dynamic() + scheduler_counter_cute = None + if dynamic_scheduling: + scheduler_counter_cute = from_dlpack(scheduler_counter, assumed_align=4).mark_layout_dynamic() - tensormap_ws_cute = from_dlpack(tensormap_workspace, assumed_align=128).mark_layout_dynamic() + tensormap_workspace_cute = from_dlpack(tensormap_workspace, assumed_align=128).mark_layout_dynamic() cache["compiled"] = compile( io_dtype, @@ -3227,7 +3087,7 @@ def chunk_gdn2_sm100( k_ratio, v_ratio, HO, - dyn_sched, + dynamic_scheduling, num_sm=multiprocessor_count(current_device()), q_cute=q_cute, k_cute=k_cute, @@ -3243,8 +3103,8 @@ def chunk_gdn2_sm100( state_out_cute=state_out_cute, work_items_cute=work_items_cute, work_count_cute=work_count_cute, - sched_ctr_cute=sched_ctr_cute, - tensormap_ws_cute=tensormap_ws_cute, + scheduler_counter_cute=scheduler_counter_cute, + tensormap_workspace_cute=tensormap_workspace_cute, checkpoint_every_n_tokens=checkpoint_every_n_tokens, scale=scale, stream=cu_stream, @@ -3261,7 +3121,7 @@ def chunk_gdn2_sm100( w_pl = from_dlpack(w, assumed_align=16).mark_layout_dynamic(leading_dim=2) o_pl = from_dlpack(output, assumed_align=16).mark_layout_dynamic(leading_dim=2) cu_pl = from_dlpack(cu_seqlens, assumed_align=8).mark_layout_dynamic() - ws_pl = from_dlpack(tensormap_workspace, assumed_align=128).mark_layout_dynamic() + workspace_pl = from_dlpack(tensormap_workspace, assumed_align=128).mark_layout_dynamic() state_checkpoints_pl = None if state_checkpoints_for_descs is not None: state_checkpoints_pl = from_dlpack(state_checkpoints_for_descs, assumed_align=16).mark_layout_dynamic(leading_dim=3) @@ -3272,15 +3132,15 @@ def chunk_gdn2_sm100( work_items_pl = from_dlpack(work_items, assumed_align=16) work_items_pl.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1), divisibility=1) work_count_pl = from_dlpack(work_count, assumed_align=4).mark_layout_dynamic() - sched_pl = None - if dyn_sched: - sched_pl = from_dlpack(sched_ctr, assumed_align=4).mark_layout_dynamic() + scheduler_pl = None + if dynamic_scheduling: + scheduler_pl = from_dlpack(scheduler_counter, assumed_align=4).mark_layout_dynamic() cache["prologue"] = cute.compile( prologue, io_dtype, CFG.B_T, order_gen, - dyn_sched, + dynamic_scheduling, q_pl, k_pl, v_pl, @@ -3293,8 +3153,8 @@ def chunk_gdn2_sm100( staging_pl, work_count_pl, work_items_pl, - sched_pl, - ws_pl, + scheduler_pl, + workspace_pl, cutlass.Int32(checkpoint_every_n_tokens), cu_stream, options="--enable-tvm-ffi", @@ -3312,7 +3172,7 @@ def chunk_gdn2_sm100( work_item_scratch if not order_gen else None, work_count, work_items, - sched_ctr, + scheduler_counter, tensormap_workspace, checkpoint_every_n_tokens, cu_stream, @@ -3332,7 +3192,7 @@ def chunk_gdn2_sm100( output_state if store_final_state else None, work_items, work_count, - sched_ctr, + scheduler_counter, tensormap_workspace, checkpoint_every_n_tokens, scale, @@ -3358,7 +3218,7 @@ def run_prefill( output_state_checkpoints, work_items, work_count, - sched_ctr, + scheduler_counter, work_item_scratch, tensormap_workspace, checkpoint_every_n_tokens, @@ -3382,7 +3242,7 @@ def run_prefill( work_item_scratch, work_count, work_items, - sched_ctr, + scheduler_counter, tensormap_workspace, checkpoint_every_n_tokens, cu_stream, @@ -3402,7 +3262,7 @@ def run_prefill( output_state, work_items, work_count, - sched_ctr, + scheduler_counter, tensormap_workspace, checkpoint_every_n_tokens, scale, diff --git a/python/cudnn/linear_attention/frost/kernel/gdn2_recompute_config.py b/python/cudnn/linear_attention/frost/kernel/gdn2_recompute_config.py index 2b723e05b..07a74ba7f 100644 --- a/python/cudnn/linear_attention/frost/kernel/gdn2_recompute_config.py +++ b/python/cudnn/linear_attention/frost/kernel/gdn2_recompute_config.py @@ -54,7 +54,7 @@ class Cfg: # --- SMEM / TMEM ring stage counts --- SMEM_RAW_STAGES: int = 5 - SMEM_SCHED_STAGES: int = 8 + SMEM_SCHEDULER_STAGES: int = 8 SMEM_DECAY_STAGES: int = 2 SMEM_INTERMEDIATE_STAGES: int = 2 SMEM_STATE_SCALE_DIAG_STAGES: int = 4 diff --git a/python/cudnn/linear_attention/frost/kernel/gdn2_recompute_f16.py b/python/cudnn/linear_attention/frost/kernel/gdn2_recompute_f16.py index 8e5c5dd23..91a007161 100644 --- a/python/cudnn/linear_attention/frost/kernel/gdn2_recompute_f16.py +++ b/python/cudnn/linear_attention/frost/kernel/gdn2_recompute_f16.py @@ -15,70 +15,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Chunked Gated DeltaNet v2 (GDN-2) recompute (state/checkpoint-only) kernel for -Blackwell SM100/SM103 (Cutlass DSL), BT=16 tiling with per-key-channel decay + -per-key erase gate (beta) + per-value write gate (w), using direct CUTLASS -primitives. Framework-neutral entry ``chunk_gdn2_recompute_sm100``. - -A copy of the prefill (``gdn2_prefill_f16.py``) with the Q/O path removed: -it reproduces the per-chunk checkpoint series and the final state -byte-for-byte while skipping the attention output. - - S_t = S_{t-1} * diag(exp(g_t)) - U_t = W_t * V_t - (Beta_t * K_t)^T S_t - S_t += K_t (x) U_t - -vs KDA: the erase gate Beta and write gate W are per-channel tensors. Beta -is folded into the K_decay operand (feeds KK^T and state*K), the strict-lower -tile loses its per-row Beta scale, and Y becomes `W*V - state*K`; Beta/W -arrive by TMA alongside K/V. - -ABI: k `[T, HK, DK]`, v `[T, HV, DV]`, gate -`[T, HO, DK]` fp32 (natural-log decay unless SAFE_GATE), beta `[T, HO, DK]` -and w `[T, HO, DV]` in the io dtype, cu_seqlens int32, states/checkpoints -`[N, HO, DV, DK]` (VK, k contiguous). GQA/GVA head broadcast follows repeat_interleave: -source head = head_idx // (HO // H_x). State presence, L2NORM, SAFE_GATE, -checkpoints, and the head ratios are compile-time specializations. - -Warp assignments (16 warps = 512 threads): - warps 0-7 : compute group 0 - Gate prefix scan + decay/restore operands - warps 8-11 : compute group 1 - TMEM value side, state stores - warp 12 : super-MMA - register-MMA KK^T + Neumann inverse - warp 13 : tcgen05-MMA - the four state GEMMs + the TMEM lifecycle - warp 14 : TMA load - per-chunk input G->S loads - warp 15 : epilogue - the checkpoint TMA store - -SMEM layout: - Buffer Bytes Stages - K / V raw 20480 5 <-- SW128 TMA ring (io dtype) - Beta / W raw 2x 20480 5 <-- per-channel gates, same ring - Gate raw 40960 5 <-- fp32 prefix-scan source - dt_bias (+a_log slot) 516 1 <-- SAFE_GATE only - K_inv 8192 2 <-- token-major ldmatrix/tcgen05 B operand - K decay 8192 2 <-- tcgen05 SW128 K-box-major A/B operands - K restore 8192 2 <-- tcgen05 B operand for the state update - state-scale diag 12288 3 <-- per-k-atom decay diagonal blocks - intermediate (T_inv) 1024 2 <-- SW32 16x16 register-MMA tiles - -TMEM layout (240 of 512 columns): - Buffer Cols Purpose - state 0-127 state[DK,DV] fp32 recurrent state - state inp 128-191 packed b16 A operand view of the state - state_k_acc 192-207 state*K fp32 accumulator - u_acc 208-223 U fp32 accumulator - Y 224-231 packed b16 A operand: W*V - state*(Beta*K) - U input 232-239 packed b16 A operand: the U readback - -GEMM schedule (tcgen05-MMA warp, in issue order per chunk): - state*K -> state_k_acc - state decay (diag blocks) - U = T_inv @ Y -> u_acc - final_state += U @ K_restore - -Requires a cutlass DSL build providing `cutlass.experimental.*`; not available -in the pip nvidia-cutlass-dsl releases. -""" - from dataclasses import dataclass from functools import lru_cache from typing import Callable, NamedTuple, Optional, Type @@ -90,7 +26,7 @@ import cutlass.cute as cute from cutlass.cute.runtime import from_dlpack -from ..common.split_k import ORDER_CAPACITY, ORDER_ELEMS, ORDER_THREADS, decode_work_item, order_body +from ..common.split_k import ORDER_CAPACITY, ORDER_ELEMENTS, ORDER_THREADS, decode_work_item, order_body from ..common.host import get_dtype from cudnn.frost.buffers import data_ptr from cudnn.frost.device import current_device, multiprocessor_count @@ -105,9 +41,10 @@ ) from cudnn.frost.tile_dsl.handles import GmemTileTma, MmaDesc, SmemTile, tma_slice_runtime_desc from cudnn.frost.tile_dsl.mma import mma_step, mma_ts_step -from cudnn.frost.tile_dsl.swizzle import swizzle_lin_128b, swizzle_lin_S, swizzle_xor_128b +from cudnn.frost.tile_dsl.swizzle import swizzle_xor_128b, swizzle_xor_32b from cudnn.frost.tile_dsl.tma import tma_load_tile, tma_store_commit, tma_store_tile, tma_store_wait, tma_tensormap_acquire from cudnn.frost.tile_dsl.pointwise import ( + sigmoid, f16x2_to_f32, fadd2, fmul2, @@ -120,14 +57,7 @@ ) LOG2_E: float = 1.4426950408889634 - - DEFAULT_GATE_LOWER_BOUND: float = -5.0 - - -# Host-side API defaults. - - L2_NORM_EPS: float = 1.0e-12 @@ -149,9 +79,9 @@ class Gdn2Bars(NamedTuple): mb_state_k_acc_ready: MBarrier mb_u_acc_ready: MBarrier - mb_state_inp_ready: MBarrier - mb_y_inp_ready: MBarrier - mb_u_inp_ready: MBarrier + mb_state_input_ready: MBarrier + mb_y_input_ready: MBarrier + mb_u_input_ready: MBarrier mb_k_decay_inv_cg0_ready: MBarrier mb_decay_tcgen05_done: MBarrier @@ -168,13 +98,12 @@ class Gdn2Bars(NamedTuple): mb_checkpoint_tmastg_ready: MBarrier mb_checkpoint_tmastg_done: MBarrier - mb_sched_ready: MBarrier - mb_sched_done: MBarrier + mb_scheduler_ready: MBarrier + mb_scheduler_done: MBarrier def make_gdn2_bars(cfg) -> Gdn2Bars: - """Bars factory. MUST be called from inside ``kernel`` (allocates the - mbarrier rings in SMEM ahead of the data buffers).""" + """Gdn2Bars factory.""" def alloc(n): return cutlass.Array(cutlass.Int64, n, space=cutlass.AddressSpace.smem, alignment=8) @@ -196,9 +125,9 @@ def alloc(n): mb_beta_done=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=CG0_GROUP_THREADS, producer=Producer.THREAD), mb_state_k_acc_ready=MBarrier(alloc(1), stages=1, init_count=1, producer=Producer.MMA_COMMIT), mb_u_acc_ready=MBarrier(alloc(1), stages=1, init_count=1, producer=Producer.MMA_COMMIT), - mb_state_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), - mb_y_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), - mb_u_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_state_input_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_y_input_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_u_input_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), mb_k_decay_inv_cg0_ready=MBarrier(alloc(cfg.smem_decay_stages), stages=cfg.smem_decay_stages, init_count=CG0_GROUP_THREADS, producer=Producer.THREAD), mb_decay_tcgen05_done=MBarrier(alloc(cfg.smem_decay_stages), stages=cfg.smem_decay_stages, init_count=1, producer=Producer.MMA_COMMIT), mb_decay_super_done=MBarrier(alloc(cfg.smem_decay_stages), stages=cfg.smem_decay_stages, init_count=WARP, producer=Producer.THREAD), @@ -223,40 +152,40 @@ def alloc(n): alloc(cfg.smem_checkpoint_stages), stages=cfg.smem_checkpoint_stages, init_count=CG1_THREADS, producer=Producer.THREAD ), mb_checkpoint_tmastg_done=MBarrier(alloc(cfg.smem_checkpoint_stages), stages=cfg.smem_checkpoint_stages, init_count=WARP, producer=Producer.THREAD), - mb_sched_ready=MBarrier(alloc(cfg.sched_stages), stages=cfg.sched_stages, init_count=1, producer=Producer.THREAD), - mb_sched_done=MBarrier(alloc(cfg.sched_stages), stages=cfg.sched_stages, init_count=15, producer=Producer.THREAD), + mb_scheduler_ready=MBarrier(alloc(cfg.scheduler_stages), stages=cfg.scheduler_stages, init_count=1, producer=Producer.THREAD), + mb_scheduler_done=MBarrier(alloc(cfg.scheduler_stages), stages=cfg.scheduler_stages, init_count=15, producer=Producer.THREAD), ) -# ---- Dynamic tile scheduler ------------------------------------------------------ +# ---- Dynamic tile scheduler ---------------------------------------------------------- @cute.jit -def sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ctas): +def scheduler_publish_next(cfg, bars, sScheduler, mScheduler, scheduler_state, tile_idx, num_ctas, elect_one): """TMA-warp side: pull the next tile off the global ticket, publish it.""" - if cutlass.const_expr(cfg.dyn_sched): - bars.mb_sched_done[sched_state.idx].wait(sched_state.phase) - if nvvm.elect_sync(): - fetched = cutlass.Int32(nvvm.atomicrmw("add", mSched.iterator, cutlass.Int32(1), mem_order="relaxed", syncscope="gpu")) - sSched[sched_state.idx] = num_ctas + fetched + if cutlass.const_expr(cfg.dynamic_scheduling): + bars.mb_scheduler_done[scheduler_state.idx].wait(scheduler_state.phase) + if elect_one: + fetched = cutlass.Int32(nvvm.atomicrmw("add", mScheduler.iterator, cutlass.Int32(1), mem_order="relaxed", syncscope="gpu")) + sScheduler[scheduler_state.idx] = num_ctas + fetched nvvm.bar_warp_sync(cute.arch.FULL_MASK) - next_tile = sSched[sched_state.idx] - if nvvm.elect_sync(): - bars.mb_sched_ready[sched_state.idx].arrive() - return next_tile, advance(sched_state, cfg.sched_stages) - return tile_idx + num_ctas, sched_state + next_tile = sScheduler[scheduler_state.idx] + if elect_one: + bars.mb_scheduler_ready[scheduler_state.idx].arrive() + return next_tile, advance(scheduler_state, cfg.scheduler_stages) + return tile_idx + num_ctas, scheduler_state @cute.jit -def sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas): +def scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one): """Consumer side: read the TMA warp's published next tile.""" - if cutlass.const_expr(cfg.dyn_sched): - bars.mb_sched_ready[sched_state.idx].wait(sched_state.phase) - next_tile = sSched[sched_state.idx] - if nvvm.elect_sync(): - bars.mb_sched_done[sched_state.idx].arrive() - return next_tile, advance(sched_state, cfg.sched_stages) - return tile_idx + num_ctas, sched_state + if cutlass.const_expr(cfg.dynamic_scheduling): + bars.mb_scheduler_ready[scheduler_state.idx].wait(scheduler_state.phase) + next_tile = sScheduler[scheduler_state.idx] + if elect_one: + bars.mb_scheduler_done[scheduler_state.idx].arrive() + return next_tile, advance(scheduler_state, cfg.scheduler_stages) + return tile_idx + num_ctas, scheduler_state @cute.jit @@ -265,9 +194,7 @@ def diag_idx(cfg, key_dim): block = key_dim // cutlass.Int32(16) coord = key_dim - block * cutlass.Int32(16) - storage_col = coord ^ cutlass.Int32((cfg.b_t // 2)) - linear_idx = block * cutlass.Int32(256) + coord * cutlass.Int32(16) + storage_col - return swizzle_lin_S(linear_idx, bbits=1, mbase=3, sshift=3) + return block * cutlass.Int32(256) + coord * cutlass.Int32(16) + swizzle_xor_32b(key_dim, coord) @cute.jit @@ -278,9 +205,9 @@ def tmaldg_warp( num_ctas, cu_seqlens, mWorkItems, - mSched, - sSched, - lane, + mScheduler, + sScheduler, + lane_idx, sBeta_raw, sGate_raw, sK_raw, @@ -295,8 +222,13 @@ def tmaldg_warp( ) -> None: """TMA-LDG warp role (warp 14): persistent scheduler loop issuing the per-chunk K/V/Beta/W/Gate G->S loads.""" - elect_one = nvvm.elect_sync() nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) + + raw_index = PipelineState.start(phase=1) + raw_bar_index = PipelineState.start(phase=0) + scheduler_state = PipelineState.start(phase=1) + + elect_one = nvvm.elect_sync() sK_tma = SmemTile( base=sK_raw, elems_per_stage=(cfg.d_k * cfg.b_t), @@ -352,12 +284,11 @@ def tmaldg_warp( tma_granu_elems=32, tma_subtile_stride_elems=(cfg.b_t * 32), ) - raw_index = PipelineState.start(phase=1) - raw_bar_index = PipelineState.start(phase=0) - sched_state = PipelineState.start(phase=1) tile_idx = cutlass.Int32(bidx) while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) head_o = head_idx head_k = head_idx if cfg.k_ratio == 1 else head_idx // cutlass.Int32(cfg.k_ratio) head_v = head_idx if cfg.v_ratio == 1 else head_idx // cutlass.Int32(cfg.v_ratio) @@ -373,38 +304,38 @@ def tmaldg_warp( tma_tensormap_acquire(desc_gate_slot) tma_tensormap_acquire(desc_beta_slot) tma_tensormap_acquire(desc_w_slot) - for chunk_idx in cutlass.range(cstart, wend, 1, unroll=1): + for chunk_idx in cutlass.range(compute_start, write_end, 1, unroll=1): chunk_start = chunk_idx * cfg.b_t - # ---- K load ---------------------------------------------------------- + # ---- K load -------------------------------------------------------------- bars.mb_k_done[raw_index.idx].wait(raw_index.phase) if elect_one: bars.mb_k_ready[raw_bar_index.idx].arrive(n_bytes=cfg.tma_k_bytes) k_slice = tma_slice_runtime_desc(desc_k_slot, cutlass.Int32(0), head_k, chunk_start) tma_load_tile(sK_tma[raw_index.idx], k_slice, bars.mb_k_ready[raw_bar_index.idx].smem_ptr, acquire=False) - # ---- V load ---------------------------------------------------------- + # ---- V load -------------------------------------------------------------- bars.mb_v_done[raw_index.idx].wait(raw_index.phase) if elect_one: bars.mb_v_ready[raw_bar_index.idx].arrive(n_bytes=cfg.tma_v_bytes) v_slice = tma_slice_runtime_desc(desc_v_slot, cutlass.Int32(0), head_v, chunk_start) tma_load_tile(sV_tma[raw_index.idx], v_slice, bars.mb_v_ready[raw_bar_index.idx].smem_ptr, acquire=False) - # ---- Beta load ------------------------------------------------------- + # ---- Beta load: GMEM -> SMEM --------------------------------------------- bars.mb_beta_done[raw_index.idx].wait(raw_index.phase) if elect_one: bars.mb_beta_ready[raw_bar_index.idx].arrive(n_bytes=cfg.tma_beta_bytes) beta_slice = tma_slice_runtime_desc(desc_beta_slot, cutlass.Int32(0), head_o, chunk_start) tma_load_tile(sBeta_tma[raw_index.idx], beta_slice, bars.mb_beta_ready[raw_bar_index.idx].smem_ptr, acquire=False) - # ---- W load ---------------------------------------------------------- + # ---- W load -------------------------------------------------------------- bars.mb_w_done[raw_index.idx].wait(raw_index.phase) if elect_one: bars.mb_w_ready[raw_bar_index.idx].arrive(n_bytes=cfg.tma_w_bytes) w_slice = tma_slice_runtime_desc(desc_w_slot, cutlass.Int32(0), head_o, chunk_start) tma_load_tile(sW_tma[raw_index.idx], w_slice, bars.mb_w_ready[raw_bar_index.idx].smem_ptr, acquire=False) - # ---- Gate load ------------------------------------------------------- + # ---- Gate load: GMEM -> SMEM --------------------------------------------- bars.mb_gate_done[raw_index.idx].wait(raw_index.phase) if elect_one: bars.mb_gate_ready[raw_bar_index.idx].arrive(n_bytes=cfg.tma_gate_bytes) @@ -412,7 +343,7 @@ def tmaldg_warp( tma_load_tile(sGate_tma[raw_index.idx], gate_slice, bars.mb_gate_ready[raw_bar_index.idx].smem_ptr, acquire=False) raw_index = advance(raw_index, cfg.smem_raw_stages) raw_bar_index = advance(raw_bar_index, cfg.smem_raw_bar_stages) - tile_idx, sched_state = sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_publish_next(cfg, bars, sScheduler, mScheduler, scheduler_state, tile_idx, num_ctas, elect_one) @cute.jit @@ -423,8 +354,8 @@ def super_mma_warp( num_ctas, cu_seqlens, mWorkItems, - sSched, - lane, + sScheduler, + lane_idx, sK_inv_raw, sIntermediate_raw, sK_decay_raw, @@ -433,26 +364,28 @@ def super_mma_warp( """Super-MMA warp role (warp 12): persistent scheduler loop computing the register-MMA T_inv.""" nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) + elect_one = nvvm.elect_sync() - # ---- ldmatrix/stmatrix lane decode ------------------------------------------- - rhs_row_coord = lane % 8 + (cutlass.Int32(8) if (lane // 16) else cutlass.Int32(0)) - rhs_col_offset = cutlass.Int32(8) if ((lane // 8) % 2) else cutlass.Int32(0) - lhs_row_coord = lane % 8 + (cutlass.Int32(8) if ((lane // 8) % 2) else cutlass.Int32(0)) - lhs_col_offset = cutlass.Int32(8) if ((lane // 8) // 2) else cutlass.Int32(0) - decay_key_mask = cutlass.Int32(8) - stsm_row_coord = lane & 7 - stsm_col_coord = cutlass.Int32(0) - if (lane // 8) & 1: - stsm_row_coord = stsm_row_coord + cutlass.Int32(8) - if lane // 8 >= 2: - stsm_col_coord = cutlass.Int32(8) - stsm_idx = swizzle_lin_S(stsm_row_coord * cfg.b_t + (stsm_col_coord ^ (cfg.b_t // 2)), bbits=1, mbase=3, sshift=3) + # ---- ldmatrix/stmatrix lane decode ----------------------------------------------- + k_inv_row_coord = lane_idx % 8 + (cutlass.Int32(8) if (lane_idx // 16) else cutlass.Int32(0)) + k_inv_col_offset = cutlass.Int32(8) if ((lane_idx // 8) % 2) else cutlass.Int32(0) + k_decay_row_coord = lane_idx % 8 + (cutlass.Int32(8) if ((lane_idx // 8) % 2) else cutlass.Int32(0)) + k_decay_col_offset = cutlass.Int32(8) if ((lane_idx // 8) // 2) else cutlass.Int32(0) + t_inv_row_coord = lane_idx & 7 + t_inv_col_coord = cutlass.Int32(0) + if (lane_idx // 8) & 1: + t_inv_row_coord = t_inv_row_coord + cutlass.Int32(8) + if lane_idx // 8 >= 2: + t_inv_col_coord = cutlass.Int32(8) + t_inv_idx = t_inv_row_coord * cfg.b_t + swizzle_xor_32b(t_inv_row_coord, t_inv_col_coord) global_chunk_base = cutlass.Int32(0) - sched_state = PipelineState.start(phase=0) + scheduler_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) - num_chunks_tile = wend - cstart # processed chunks; ring bookkeeping runs on global_chunk_base + local_chunk + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) + num_chunks_tile = write_end - compute_start for local_chunk in cutlass.range(num_chunks_tile, unroll=1): global_chunk = global_chunk_base + local_chunk decay_stage = global_chunk % cfg.smem_decay_stages @@ -463,53 +396,51 @@ def super_mma_warp( bars.mb_k_decay_inv_cg0_ready[decay_stage].wait((global_chunk // cfg.smem_decay_stages) % 2) - # ---- KK = K_decay @ K_inv^T ------------------------------------------ + # ---- KK = K decay @ K inv^T ---------------------------------------------- kk_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): kk_acc[accum_idx] = cutlass.Float32(0.0) - for k_block in cutlass.range_constexpr((cfg.d_k // 16)): - # Load B operand - k_inv_col = k_block * 16 + rhs_col_offset + for i in cutlass.range_constexpr((cfg.d_k // 16)): + k_inv_col = i * 16 + k_inv_col_offset k_inv_segment = k_inv_col // 64 - rhs_frag = nvvm.ldmatrix( + k_inv_frag = nvvm.ldmatrix( sK_inv_ptr + k_inv_segment * (cfg.b_t * 64) - + rhs_row_coord * 64 - + swizzle_xor_128b(rhs_row_coord, k_inv_col - k_inv_segment * 64, elem_bytes=2), + + k_inv_row_coord * 64 + + swizzle_xor_128b(k_inv_row_coord, k_inv_col - k_inv_segment * 64, elem_bytes=2), 4, nvvm.MMALayout.ROW, ) - # Load A operand - storage_key = (k_block * 16 + lhs_col_offset) ^ decay_key_mask - storage_slice = storage_key // 64 - kk_lhs_frag = nvvm.ldmatrix( + k_decay_col = i * 16 + k_decay_col_offset + k_decay_segment = k_decay_col // 64 + k_decay_frag = nvvm.ldmatrix( sK_decay_ptr - + storage_slice * (cfg.b_t * 64) - + swizzle_xor_128b(lhs_row_coord, lhs_row_coord * 64 + storage_key - storage_slice * 64, elem_bytes=2), + + k_decay_segment * (cfg.b_t * 64) + + swizzle_xor_128b(k_decay_row_coord, k_decay_row_coord * 64 + k_decay_col - k_decay_segment * 64, elem_bytes=2), 4, nvvm.MMALayout.ROW, ) mma_step( kk_acc, - (kk_lhs_frag[0], kk_lhs_frag[1], kk_lhs_frag[2], kk_lhs_frag[3]), - (rhs_frag[0], rhs_frag[1], rhs_frag[2], rhs_frag[3]), + (k_decay_frag[0], k_decay_frag[1], k_decay_frag[2], k_decay_frag[3]), + (k_inv_frag[0], k_inv_frag[1], k_inv_frag[2], k_inv_frag[3]), k_step=0, M=16, N=16, ab_dtype=cfg.io_dtype, ) - # ---- L = tril(KK, -1) fragment --------------------------------------- - row_lo = lane // 4 + # ---- L = tril(KK, -1) ---------------------------------------------------- + row_lo = lane_idx // 4 row_hi = row_lo + cutlass.Int32(8) l_regs = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): row_coord = row_lo if cutlass.const_expr(accum_idx % 4 >= 2): row_coord = row_hi - col_coord = (accum_idx // 4) * 8 + 2 * (lane % 4) + col_coord = (accum_idx // 4) * 8 + 2 * (lane_idx % 4) if cutlass.const_expr(accum_idx % 2 == 1): col_coord = col_coord + cutlass.Int32(1) l_regs[accum_idx] = kk_acc[accum_idx] if row_coord > col_coord else cutlass.Float32(0.0) @@ -519,13 +450,13 @@ def super_mma_warp( l_a3 = fp32_to_fp16(l_regs[6], l_regs[7], dtype=cfg.io_dtype) l_values = cutlass.Vector.from_elements((l_a0, l_a1, l_a2, l_a3), cutlass.Int32).bitcast(cfg.io_dtype).to(cutlass.Float32) - # ---- T_inv = I - L, then three Neumann doubling rounds --------------- + # ---- T^-1 = I - L, then three Neumann doubling rounds -------------------- tinv_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): row_coord = row_lo if cutlass.const_expr(accum_idx % 4 >= 2): row_coord = row_hi - col_coord = (accum_idx // 4) * 8 + 2 * (lane % 4) + col_coord = (accum_idx // 4) * 8 + 2 * (lane_idx % 4) if cutlass.const_expr(accum_idx % 2 == 1): col_coord = col_coord + cutlass.Int32(1) eye = cutlass.Float32(1.0) if row_coord == col_coord else cutlass.Float32(0.0) @@ -533,8 +464,8 @@ def super_mma_warp( lpow_a0, lpow_a1, lpow_a2, lpow_a3 = l_a0, l_a1, l_a2, l_a3 mov_lpow0, mov_lpow1, mov_lpow2, mov_lpow3 = movmatrix_16b(l_a0), movmatrix_16b(l_a1), movmatrix_16b(l_a2), movmatrix_16b(l_a3) - for _round in cutlass.range_constexpr(3): - # ---- Lpow = Lpow @ Lpow ------------------------------------------ + for neumann_round in cutlass.range_constexpr(3): + # ---- Lpow = Lpow @ Lpow ---------------------------------------------- sq_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): sq_acc[accum_idx] = cutlass.Float32(0.0) @@ -552,7 +483,7 @@ def super_mma_warp( lpow_a2 = fp32_to_fp16(sq_acc[4], sq_acc[5], dtype=cfg.io_dtype) lpow_a3 = fp32_to_fp16(sq_acc[6], sq_acc[7], dtype=cfg.io_dtype) mov_lpow0, mov_lpow1, mov_lpow2, mov_lpow3 = movmatrix_16b(lpow_a0), movmatrix_16b(lpow_a1), movmatrix_16b(lpow_a2), movmatrix_16b(lpow_a3) - # ---- T_inv += T_inv @ Lpow --------------------------------------- + # ---- T^-1 += T^-1 @ Lpow --------------------------------------------- upd_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): upd_acc[accum_idx] = cutlass.Float32(0.0) @@ -580,7 +511,7 @@ def super_mma_warp( bars.mb_t_inv_done[intermediate_stage].wait(((global_chunk // cfg.smem_intermediate_stages) + 1) % 2) nvvm.stmatrix( - sIntermediate_ptr + stsm_idx, + sIntermediate_ptr + t_inv_idx, [ fp32_to_fp16(tinv_acc[0], tinv_acc[1], dtype=cfg.io_dtype), fp32_to_fp16(tinv_acc[2], tinv_acc[3], dtype=cfg.io_dtype), @@ -594,7 +525,7 @@ def super_mma_warp( bars.mb_t_inv_ready[intermediate_stage].arrive() bars.mb_decay_super_done[decay_stage].arrive() global_chunk_base += num_chunks_tile - tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one) @cute.jit @@ -605,11 +536,11 @@ def tcgen05_mma_warp( num_ctas, cu_seqlens, mWorkItems, - sSched, + sScheduler, tmem_base_holder, sIntermediate, sK_decay, - sK_restore, + sK_restore_trans, sState_scale_diag, bars, ) -> None: @@ -620,20 +551,20 @@ def tcgen05_mma_warp( nvvm.tcgen05_alloc(tmem_base_holder, cutlass.Int32(512), group=nvvm.CTAGroup.CTA_1) nvvm.barrier_cta_sync(cfg.tmem_lifecycle_barrier_id, thread_count=cfg.tmem_user_threads) tmem_base = tmem_base_holder.load() - state_inp_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_state_inp_offset, cutlass.Int8) + state_input_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_state_input_offset, cutlass.Int8) state_dsts = tuple(nvvm.make_tmem_ptr(tmem_base + cfg.tmem_state_acc_offset + k * 16, cutlass.Float32) for k in range(cfg.d_k // 16)) state_k_acc_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_state_k_acc_offset, cutlass.Float32) u_acc_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_u_acc_offset, cutlass.Float32) - y_inp_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_y_inp_offset, cutlass.Int8) - u_inp_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_u_inp_offset, cutlass.Int8) + y_input_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_y_input_offset, cutlass.Int8) + u_input_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_u_input_offset, cutlass.Int8) state_dst_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_state_acc_offset, cutlass.Float32) - state_inp_index = PipelineState.start(phase=0) + state_input_index = PipelineState.start(phase=0) state_read_index = PipelineState.start(phase=0) - y_inp_index = PipelineState.start(phase=0) - u_inp_index = PipelineState.start(phase=0) + y_input_index = PipelineState.start(phase=0) + u_input_index = PipelineState.start(phase=0) qk_scale_index = PipelineState.start(phase=0) - # ---- chunk-invariant GEMM descriptors ---------------------------------------- + # ---- chunk-invariant GEMM descriptors -------------------------------------------- bpe = cfg.io_dtype.width // 8 idesc_acc = nvvm.Tcgen05InstrDesc.build( c_dtype=cutlass.Float32, @@ -659,7 +590,7 @@ def tcgen05_mma_warp( m_dim=cfg.d_v, b_major=1, ) - bmm_state_desc = MmaDesc( + bmm_state_k_decay_desc = MmaDesc( M=cfg.d_v, N=cfg.b_t, K=cfg.d_k, @@ -671,7 +602,7 @@ def tcgen05_mma_warp( idesc=idesc_acc, kind=nvvm.Tcgen05MMAKind.F16, ) - bmm_diag_desc = MmaDesc( + bmm_state_diag_desc = MmaDesc( M=cfg.d_v, N=16, K=16, @@ -683,7 +614,7 @@ def tcgen05_mma_warp( idesc=idesc_diag, kind=nvvm.Tcgen05MMAKind.F16, ) - bmm_t_inv_desc = MmaDesc( + bmm_y_t_inv_desc = MmaDesc( M=cfg.d_v, N=cfg.b_t, K=cfg.b_t, @@ -695,7 +626,7 @@ def tcgen05_mma_warp( idesc=idesc_acc, kind=nvvm.Tcgen05MMAKind.F16, ) - bmm_final_state_desc = MmaDesc( + bmm_u_k_restore_desc = MmaDesc( M=cfg.d_v, N=cfg.d_k, K=cfg.b_t, @@ -707,37 +638,44 @@ def tcgen05_mma_warp( idesc=idesc_final_state, kind=nvvm.Tcgen05MMAKind.F16, ) - STATE_A_SEG = bmm_state_desc.sps_B * bmm_state_desc.tmem_advance_A - STATE_B_SEG = bmm_state_desc.smem_subtile_B >> 4 + STATE_A_SEG = bmm_state_k_decay_desc.sps_B * bmm_state_k_decay_desc.tmem_advance_A + STATE_B_SEG = bmm_state_k_decay_desc.smem_subtile_B >> 4 global_chunk_base = cutlass.Int32(0) - sched_state = PipelineState.start(phase=0) + scheduler_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) - num_chunks_tile = wend - cstart + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) + num_chunks_tile = write_end - compute_start + if cutlass.const_expr(cfg.use_initial_state): + seed_state = compute_start == 0 for local_chunk in cutlass.range(num_chunks_tile, unroll=1): global_chunk = global_chunk_base + local_chunk - have_state = cutlass.Boolean(True) if cutlass.const_expr(cfg.use_initial_state) else local_chunk > 0 + if cutlass.const_expr(cfg.use_initial_state): + have_state = local_chunk > 0 or seed_state + else: + have_state = local_chunk > 0 decay_stage = global_chunk % cfg.smem_decay_stages state_scale_diag_stage = qk_scale_index.idx intermediate_stage = global_chunk % cfg.smem_intermediate_stages sK_decay_stage = sK_decay[decay_stage] - sK_restore_stage = sK_restore[decay_stage] + sK_restore_stage = sK_restore_trans[decay_stage] sState_scale_diag_stage = sState_scale_diag[state_scale_diag_stage] sIntermediate_stage = sIntermediate[intermediate_stage] - # ---- state_k_acc = state(T) @ K_decay^T ------------------------------ + # ---- k state = state(T) @ K decay^T -------------------------------------- bars.mb_k_decay_inv_cg0_ready[decay_stage].wait((global_chunk // cfg.smem_decay_stages) % 2) if have_state: - bars.mb_state_inp_ready.wait(state_inp_index.phase) - state_inp_index = advance(state_inp_index, 1) + bars.mb_state_input_ready.wait(state_input_index.phase) + state_input_index = advance(state_input_index, 1) desc_k_decay = sK_decay_stage.desc() - for s in cutlass.range_constexpr(bmm_state_desc.num_subtiles_B): - for k in cutlass.range_constexpr(bmm_state_desc.sps_B): + for s in cutlass.range_constexpr(bmm_state_k_decay_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_state_k_decay_desc.sps_B): mma_ts_step( - bmm_state_desc, - state_inp_ptr.subview(s * STATE_A_SEG), + bmm_state_k_decay_desc, + state_input_ptr.subview(s * STATE_A_SEG), desc_k_decay + s * STATE_B_SEG, state_k_acc_ptr, k, @@ -755,16 +693,16 @@ def tcgen05_mma_warp( bars.mb_state_acc_read_done.wait(state_read_index.phase) state_read_index = advance(state_read_index, 1) - # ---- state decay = state(T) @ diag(exp2(g_last)) (per-k-atom blocks) -- + # ---- state decay = state(T) @ diag(exp2(g last)) (per-k-atom blocks) ----- bars.mb_qk_scale_ready[qk_scale_index.idx].wait(qk_scale_index.phase) if have_state: desc_diag = sState_scale_diag_stage.desc() - for k_block in cutlass.range_constexpr(cfg.d_k // 16): + for i in cutlass.range_constexpr(cfg.d_k // 16): mma_ts_step( - bmm_diag_desc, - state_inp_ptr.subview(k_block * bmm_diag_desc.tmem_advance_A), - desc_diag.advance_start_address(k_block * 256 * 2), - state_dsts[k_block], + bmm_state_diag_desc, + state_input_ptr.subview(i * bmm_state_diag_desc.tmem_advance_A), + desc_diag.advance_start_address(i * 256 * 2), + state_dsts[i], 0, cutlass.Boolean(False), ) @@ -772,29 +710,29 @@ def tcgen05_mma_warp( if elect_one: bars.mb_state_scale_diag_done[state_scale_diag_stage].arrive(cta_group=1) - # ---- u_acc = Y(T) @ T_inv^T ------------------------------------------ + # ---- u acc = Y(T) @ (T^-1)^T --------------------------------------------- bars.mb_t_inv_ready[intermediate_stage].wait((global_chunk // cfg.smem_intermediate_stages) % 2) - bars.mb_y_inp_ready.wait(y_inp_index.phase) - y_inp_index = advance(y_inp_index, 1) + bars.mb_y_input_ready.wait(y_input_index.phase) + y_input_index = advance(y_input_index, 1) desc_t_inv = sIntermediate_stage.desc() - mma_ts_step(bmm_t_inv_desc, y_inp_ptr, desc_t_inv, u_acc_ptr, 0, cutlass.Boolean(False)) + mma_ts_step(bmm_y_t_inv_desc, y_input_ptr, desc_t_inv, u_acc_ptr, 0, cutlass.Boolean(False)) if elect_one: bars.mb_u_acc_ready.arrive(cta_group=1) bars.mb_t_inv_done[intermediate_stage].arrive(cta_group=1) - # ---- final_state += U(T) @ K_restore --------------------------------- - bars.mb_u_inp_ready.wait(u_inp_index.phase) - u_inp_index = advance(u_inp_index, 1) + # ---- final state += U(T) @ K restore ------------------------------------- + bars.mb_u_input_ready.wait(u_input_index.phase) + u_input_index = advance(u_input_index, 1) desc_k_restore = sK_restore_stage.desc() - mma_ts_step(bmm_final_state_desc, u_inp_ptr, desc_k_restore, state_dst_ptr, 0, have_state) + mma_ts_step(bmm_u_k_restore_desc, u_input_ptr, desc_k_restore, state_dst_ptr, 0, have_state) if elect_one: bars.mb_k_restore_acc_done[decay_stage].arrive(cta_group=1) qk_scale_index = advance(qk_scale_index, cfg.smem_state_scale_diag_stages) global_chunk_base += num_chunks_tile - tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one) bars.mb_tmem_done[0].wait(0) nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_1) nvvm.tcgen05_dealloc( @@ -812,7 +750,7 @@ def epilogue_warp( num_ctas, cu_seqlens, mWorkItems, - sSched, + sScheduler, sCheckpoint_raw, desc_checkpoint_base, checkpoint_every_n_tokens, @@ -834,21 +772,23 @@ def epilogue_warp( tma_subtile_stride_elems=cfg.d_k * 64, ) checkpoint_ready_index = PipelineState.start(phase=0) - sched_state = PipelineState.start(phase=0) + scheduler_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) head_o = head_idx - num_chunks_tile = wend - cstart + num_chunks_tile = write_end - compute_start if cutlass.const_expr(cfg.enable_checkpoints): checkpoint_slot = batch_idx * cutlass.Int32(TENSOR_MAP_QWORDS) desc_checkpoint_slot = (desc_checkpoint_base + checkpoint_slot).tospace(cutlass.AddressSpace.generic) checkpoint_chunks = checkpoint_every_n_tokens // cutlass.Int32(cfg.b_t) - checkpoint_quot = (cstart + cutlass.Int32(1)) // checkpoint_chunks - checkpoint_mod = (cstart + cutlass.Int32(1)) % checkpoint_chunks + checkpoint_quot = (compute_start + cutlass.Int32(1)) // checkpoint_chunks + checkpoint_mod = (compute_start + cutlass.Int32(1)) % checkpoint_chunks if elect_one: tma_tensormap_acquire(desc_checkpoint_slot) - if num_chunks_tile > 0 and wstart == 0: + if num_chunks_tile > 0 and write_start == 0: checkpoint_stage = checkpoint_ready_index.idx bars.mb_checkpoint_tmastg_ready[checkpoint_stage].wait(checkpoint_ready_index.phase) checkpoint_ready_index = advance(checkpoint_ready_index, cfg.smem_checkpoint_stages) @@ -858,11 +798,11 @@ def epilogue_warp( tma_store_wait(0) bars.mb_checkpoint_tmastg_done[checkpoint_stage].arrive() for local_chunk in cutlass.range(num_chunks_tile, unroll=1): - chunk_idx = cstart + local_chunk + chunk_idx = compute_start + local_chunk if local_chunk > 0: - # ---- checkpoint store ---------------------------------------- + # ---- checkpoint store -------------------------------------------- do_checkpoint = checkpoint_mod == 0 - do_checkpoint = do_checkpoint and chunk_idx >= wstart + do_checkpoint = do_checkpoint and chunk_idx >= write_start if do_checkpoint: checkpoint_stage = checkpoint_ready_index.idx bars.mb_checkpoint_tmastg_ready[checkpoint_stage].wait(checkpoint_ready_index.phase) @@ -877,7 +817,7 @@ def epilogue_warp( if checkpoint_mod == checkpoint_chunks: checkpoint_mod = cutlass.Int32(0) checkpoint_quot = checkpoint_quot + cutlass.Int32(1) - tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one) @cute.jit @@ -885,9 +825,7 @@ def gate_scale(cfg, raw_gate: cutlass.Float32) -> cutlass.Float32: """Map raw gate to the log2-domain decay increment.""" if cutlass.const_expr(cfg.safe_gate): - half = cutlass.Float32(0.5) - sigmoid = cute.math.tanh(raw_gate * half, approx=True) * half + half - return cfg.gate_scale_log2 * sigmoid + return cfg.gate_scale_log2 * sigmoid(raw_gate) # Default ABI: Gate arrives in natural-log space return raw_gate * cutlass.Float32(LOG2_E) @@ -900,8 +838,8 @@ def compute0_warp_group( num_ctas, cu_seqlens, mWorkItems, - sSched, - lane, + sScheduler, + lane_idx, warp_idx, mA_log, mDt_bias, @@ -919,26 +857,30 @@ def compute0_warp_group( """CG0 warp-group role (warps 0-7): persistent scheduler loop computing the Gate prefix scan and the decay/restore operands.""" nvvm.setmaxregister(cfg.num_regs_compute_group_0, nvvm.SetMaxRegisterAction.INCREASE) + elect_one = nvvm.elect_sync() + + scheduler_state = PipelineState.start(phase=0) + cg0_warp = warp_idx - cfg.compute_group_0_warp_ids[0] - cg0_group_id = cg0_warp // cfg.cg0_warps_per_group cg0_local_warp = cg0_warp % cfg.cg0_warps_per_group - prefix_dim = cg0_local_warp * cfg.threads_per_warp + lane + + cg0_group_id = cg0_warp // cfg.cg0_warps_per_group + channel_dim = cg0_local_warp * cfg.threads_per_warp + lane_idx cg0_a_log_exp = cutlass.Float32(1.0) cg0_dt_bias_value = cutlass.Float32(0.0) global_chunk_base = cutlass.Int32(0) - sched_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) opaque_one = opaque_f32_zero() + cutlass.Float32(1.0) while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) head_o = head_idx - num_chunks_tile = wend - cstart + num_chunks_tile = write_end - compute_start if cutlass.const_expr(cfg.safe_gate): if num_chunks_tile > 0: cg0_a_log_exp = cute.math.exp2(mA_log[head_o].to(cutlass.Float32) * LOG2_E, fastmath=True) - cg0_dt_bias_value = mDt_bias[head_o, prefix_dim].to(cutlass.Float32) - # non-pow2 ring cursors: one divmod per tile each, +group-stride with wrap per chunk - # tile entry: both ping-pong groups inherit each other's delivery proofs (parity-swap guard) + cg0_dt_bias_value = mDt_bias[head_o, channel_dim].to(cutlass.Float32) nvvm.barrier_cta_sync(cfg.cg0_tile_entry_barrier_id, thread_count=cfg.cg0_group_count * cfg.cg0_threads_per_group) group_chunk_base = global_chunk_base + cutlass.Int32(cg0_group_id) diag_stage_cursor = group_chunk_base % cutlass.Int32(cfg.smem_state_scale_diag_stages) @@ -947,7 +889,7 @@ def compute0_warp_group( raw_bar_cursor = group_chunk_base % cutlass.Int32(cfg.smem_raw_bar_stages) raw_bar_phase = (group_chunk_base // cutlass.Int32(cfg.smem_raw_bar_stages)) % cutlass.Int32(2) for local_chunk in cutlass.range(cg0_group_id, num_chunks_tile, cfg.cg0_group_count, unroll=1): - chunk_idx = cstart + local_chunk + chunk_idx = compute_start + local_chunk global_chunk = global_chunk_base + local_chunk chunk_start = chunk_idx * cfg.b_t decay_stage = global_chunk % cfg.smem_decay_stages @@ -967,48 +909,36 @@ def compute0_warp_group( bars.mb_gate_ready[raw_bar_cursor].wait(raw_bar_phase) row_group_start = cg0_local_warp * (cfg.b_t // cfg.cg0_warps_per_group) - lane_row_group = lane // 8 - lane_in_row_group = lane - lane_row_group * 8 + lane_row_group = lane_idx // 8 + lane_in_row_group = lane_idx - lane_row_group * 8 decay_row = row_group_start + lane_row_group - decay_key_mask = cutlass.Int32(8) - prefix_dim = cg0_local_warp * cfg.threads_per_warp + lane + channel_dim = cg0_local_warp * cfg.threads_per_warp + lane_idx - # ---- Gate prefix scan ----------------------------------------------- - f32_segment = prefix_dim // 32 + # ---- Gate prefix scan ---------------------------------------------------- + f32_segment = channel_dim // 32 prefix_seg_base = f32_segment * (cfg.b_t * 32) - prefix_col = prefix_dim - f32_segment * 32 + prefix_col = channel_dim - f32_segment * 32 g_prefix_regs = cutlass.Array(cutlass.Float32, cfg.b_t, alignment=16) if cutlass.const_expr(cfg.safe_gate): - valid_rows = seqlen_b - chunk_idx * cutlass.Int32(cfg.b_t) - valid_mask = cutlass.vector.create_mask([cfg.b_t], [valid_rows]) - for row_pair in cutlass.range_constexpr(cfg.b_t // 2): - row0 = row_pair * 2 - row1 = row0 + 1 - prefix_idx0 = prefix_seg_base + swizzle_xor_128b(row0, row0 * 32 + prefix_col, elem_bytes=4) - prefix_idx1 = prefix_seg_base + swizzle_xor_128b(row1, row1 * 32 + prefix_col, elem_bytes=4) - gate0 = (sGate_ptr + prefix_idx0).load() - gate1 = (sGate_ptr + prefix_idx1).load() - gate0 = cg0_a_log_exp * (gate0 + cg0_dt_bias_value) - gate1 = cg0_a_log_exp * (gate1 + cg0_dt_bias_value) - gate0 = gate_scale( - cfg, - gate0, - ) - gate1 = gate_scale( - cfg, - gate1, - ) - gate_pair = cutlass.Vector.from_elements((gate0, gate1), cutlass.Float32) - gate_pair = cutlass.vector.where(valid_mask[row0 : row1 + 1], gate_pair, 0.0) - g_prefix_regs[row0] = gate_pair[0] - g_prefix_regs[row1] = gate_pair[1] + for row in cutlass.range_constexpr(cfg.b_t): + prefix_idx = prefix_seg_base + swizzle_xor_128b(row, row * 32 + prefix_col, elem_bytes=4) + gate = (sGate_ptr + prefix_idx).load() + token_idx = chunk_idx * cutlass.Int32(cfg.b_t) + cutlass.Int32(row) + if token_idx < batch_seqlen: + gate = gate_scale( + cfg, + cg0_a_log_exp * (gate + cg0_dt_bias_value), + ) + else: + gate = cutlass.Float32(0.0) + g_prefix_regs[row] = gate else: for row in cutlass.range_constexpr(cfg.b_t): prefix_idx = prefix_seg_base + swizzle_xor_128b(row, row * 32 + prefix_col, elem_bytes=4) gate = (sGate_ptr + prefix_idx).load() token_idx = chunk_idx * cutlass.Int32(cfg.b_t) + cutlass.Int32(row) - if token_idx < seqlen_b: + if token_idx < batch_seqlen: gate = gate_scale( cfg, gate, @@ -1029,7 +959,7 @@ def compute0_warp_group( g_prefix_regs[row1] = prefix1 prefix_acc = prefix1 - # ---- exp2(g): stage prefixes + final-token decay --------------------- + # ---- exp2(g): stage prefixes + final-token decay ------------------------- for row in cutlass.range_constexpr(cfg.b_t): g_prefix_regs[row] = cute.math.exp2(g_prefix_regs[row], fastmath=True) @@ -1038,9 +968,9 @@ def compute0_warp_group( prefix_idx = prefix_seg_base + swizzle_xor_128b(row, row * 32 + prefix_col, elem_bytes=4) (sGate_ptr + prefix_idx).store(g_prefix_regs[row]) - # ---- state-scale diag: stage exp2(g_last) decay blocks --------------- + # ---- state-scale diag: stage exp2(g last) decay blocks ------------------- bars.mb_state_scale_diag_done[state_scale_diag_stage].wait(diag_phase ^ cutlass.Int32(1)) - sState_scale_diag_ptr[diag_idx(cfg, prefix_dim)] = exp_g_last.to(cfg.io_dtype) + sState_scale_diag_ptr[diag_idx(cfg, channel_dim)] = exp_g_last.to(cfg.io_dtype) nvvm.barrier_cta_sync(cfg.cg0_group_sync_barrier_base_id + cg0_group_id, thread_count=cfg.cg0_threads_per_group) @@ -1050,7 +980,7 @@ def compute0_warp_group( raw_k_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) raw_beta_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) - # ---- optional K L2-norm + K_inv staging ------------------------------ + # ---- optional K L2-norm + K inv stage ------------------------------------ if cutlass.const_expr(cfg.l2norm): kk_lo = opaque_f32_zero() kk_hi = opaque_f32_zero() @@ -1069,8 +999,7 @@ def compute0_warp_group( raw_k_regs[reg_base + dim_offset] = k_val beta_val = raw_beta_frag_f32[dim_offset] if cutlass.const_expr(cfg.beta_sigmoid): - half = cutlass.Float32(0.5) - beta_val = (cute.math.tanh(beta_val * half, approx=True) * half + half).to(cfg.io_dtype).to(cutlass.Float32) + beta_val = sigmoid(beta_val).to(cfg.io_dtype).to(cutlass.Float32) raw_beta_regs[reg_base + dim_offset] = beta_val if cutlass.const_expr(cfg.l2norm): # even dims in the lo lane, odd in the hi lane: same terms, same order @@ -1088,7 +1017,7 @@ def compute0_warp_group( norm_floor_sq = cutlass.Float32(L2_NORM_EPS * L2_NORM_EPS) k_inv_norm = cute.math.rsqrt(cute.math.max(k_sum_sq, norm_floor_sq), fastmath=True) - # ---- decay/restore operands: exp2(+-g) applied per key channel ------- + # ---- decay/restore operands: exp2(+-g) applied per key channel ----------- exp_g_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) exp_g_last_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) for dim_half in cutlass.range_constexpr(2): @@ -1120,7 +1049,7 @@ def compute0_warp_group( exp_g_last_regs[f32_reg_base + 2] = exp_g_last_frag[2] exp_g_last_regs[f32_reg_base + 3] = exp_g_last_frag[3] - # ---- K_decay + K_inv operands: K * exp2(+g) and K * exp2(-g) ----- + # ---- K decay + K inv operands: K * exp2(+g) and K * exp2(-g) --------- k_decay_pack = cutlass.Array(cutlass.Int32, 4, alignment=16) for pair_idx in cutlass.range_constexpr(4): dim0 = pair_idx * 2 @@ -1162,11 +1091,9 @@ def compute0_warp_group( f16_segment_dim = dim_base - f16_segment * 64 k_inv_swizzled_idx = f16_segment * (cfg.b_t * 64) + decay_row * 64 + swizzle_xor_128b(decay_row, f16_segment_dim, elem_bytes=2) (sK_inv_ptr + k_inv_swizzled_idx).store(k_inv_vec, alignment=16) - storage_key = dim_base ^ decay_key_mask - storage_slice = storage_key // 64 - decay_swizzled_idx = storage_slice * (cfg.b_t * 64) + swizzle_xor_128b( - decay_row, decay_row * 64 + storage_key - storage_slice * 64, elem_bytes=2 - ) + decay_col = dim_base + decay_segment = decay_col // 64 + decay_swizzled_idx = decay_segment * (cfg.b_t * 64) + swizzle_xor_128b(decay_row, decay_row * 64 + decay_col - decay_segment * 64, elem_bytes=2) (sK_decay_ptr + decay_swizzled_idx).store(k_decay_vec, alignment=16) nvvm.fence_proxy("async.shared", space="cta") bars.mb_k_decay_inv_cg0_ready[decay_stage].arrive() @@ -1174,7 +1101,7 @@ def compute0_warp_group( bars.mb_gate_done[raw_stage].arrive() bars.mb_beta_done[raw_stage].arrive() - # ---- K_restore operand: K_inv * exp_g_last -------------------------- + # ---- K restore operand: K inv * exp2(g last) ----------------------------- bars.mb_k_restore_acc_done[decay_stage].wait(((global_chunk // cfg.smem_decay_stages + 1) % 2)) for dim_half in cutlass.range_constexpr(2): dim_base = dim_half * (cfg.d_k // 2) + lane_in_row_group * 8 @@ -1185,10 +1112,9 @@ def compute0_warp_group( dim1 = dim0 + 1 exp_g_last_pair = fp32_to_fp16(exp_g_last_regs[reg_base + dim0], exp_g_last_regs[reg_base + dim1], dtype=cfg.io_dtype) k_restore_pack[pair_idx] = mul_f16x2(k_inv_pack[dim_half * 4 + pair_idx], exp_g_last_pair, cfg.io_dtype) - storage_row = decay_row ^ (cfg.b_t // 2) f16_segment = dim_base // 64 f16_segment_dim = dim_base - f16_segment * 64 - k_restore_idx = f16_segment * (cfg.b_t * 64) + storage_row * 64 + swizzle_xor_128b(storage_row, f16_segment_dim, elem_bytes=2) + k_restore_idx = f16_segment * (cfg.b_t * 64) + decay_row * 64 + swizzle_xor_128b(decay_row, f16_segment_dim, elem_bytes=2) k_restore_vec = cutlass.Vector.from_elements( ( k_restore_pack[0], @@ -1213,7 +1139,7 @@ def compute0_warp_group( raw_bar_cursor = raw_bar_cursor - cutlass.Int32(cfg.smem_raw_bar_stages) if raw_bar_wrapped else raw_bar_cursor raw_bar_phase = raw_bar_phase ^ (cutlass.Int32(1) if raw_bar_wrapped else cutlass.Int32(0)) global_chunk_base += num_chunks_tile - tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one) @cute.jit @@ -1224,8 +1150,8 @@ def compute1_warp_group( num_ctas, cu_seqlens, mWorkItems, - sSched, - lane, + sScheduler, + lane_idx, tmem_base_holder, warp_idx, mState_out, @@ -1239,124 +1165,105 @@ def compute1_warp_group( """CG1 warp-group role (warps 8-11): persistent scheduler loop staging the value-side TMEM operands and storing checkpoints/final state.""" nvvm.setmaxregister(cfg.num_regs_compute_group_1, nvvm.SetMaxRegisterAction.INCREASE) - sCheckpoint_ptr = sCheckpoint_raw.data_ptr() + elect_one = nvvm.elect_sync() + checkpoint_done_index = PipelineState.start(phase=1) + + sCheckpoint_ptr = sCheckpoint_raw.data_ptr() nvvm.barrier_cta_sync(cfg.tmem_lifecycle_barrier_id, thread_count=cfg.tmem_user_threads) tmem_base = tmem_base_holder.load() tmem_col = tmem_base & 0xFFFF tmem_row = tmem_base >> 16 + row_lo_addr = tmem_row << 16 + row_hi_addr = (tmem_row + 16) << 16 tmem_subpartition = warp_idx % (cfg.d_v // cfg.threads_per_warp) # ldmatrix.x4/stmatrix.x4 COL lane decode for the V/W loads - ldsm_row_coord = (lane // 16) * 8 + (lane & 7) - ldsm_col_offset = ((lane // 8) & 1) * 8 - row_id = tmem_row + tmem_subpartition * cfg.threads_per_warp - value_dim = tmem_subpartition * cfg.threads_per_warp + lane + ov_row_coord = (lane_idx // 16) * 8 + (lane_idx & 7) + ov_col_offset = ((lane_idx // 8) & 1) * 8 + value_dim = tmem_subpartition * cfg.threads_per_warp + lane_idx state_k_acc_index = PipelineState.start(phase=0) u_acc_index = PipelineState.start(phase=0) k_restore_index = PipelineState.start(phase=0) # CG1's per-chunk mb_k_restore_acc_done wait slot raw_index = PipelineState.start(phase=0) # raw-ring slot for the sV/sW reads + inputs_done arrives - raw_bar_index = PipelineState.start(phase=0) # even-depth ready-ring slot (decoupled from the data ring) - sched_state = PipelineState.start(phase=0) + raw_bar_index = PipelineState.start(phase=0) # even-depth ready-ring slot + scheduler_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) head_o = head_idx - num_chunks_tile = wend - cstart + num_chunks_tile = write_end - compute_start if num_chunks_tile > 0: - # ---- first chunk: seed state TMEM from mState_init ---------- - seed_from_initial_state = cstart == 0 + # ---- first chunk: seed state TMEM from mState init ----------------------- + seed_from_initial_state = compute_start == 0 + sV_ptr = sV_raw.data_ptr() + raw_index.idx * (cfg.d_v * cfg.b_t) + sW_ptr = sW_raw.data_ptr() + raw_index.idx * (cfg.d_v * cfg.b_t) + + # ---- state seed: initial state GMEM -> state TMEM ------------------------ + state_col_id = tmem_col + cfg.tmem_state_acc_offset + packed_col_id = tmem_col + cfg.tmem_state_input_offset if cutlass.const_expr(mState_init is not None): if seed_from_initial_state: seed_vw = 16 // (mState_init.element_type.width // 8) seed_src = (mState_init.iterator + mState_init.layout((batch_idx, head_o, value_dim, 0))).raw_ptr() - for key_block_start in cutlass.range_constexpr(0, cfg.d_k, 32): - state_block = cutlass.Array(cutlass.Float32, 32, alignment=16) - for g in cutlass.range_constexpr(32 // seed_vw): - seed_chunk = (seed_src + key_block_start + g * seed_vw).load(count=seed_vw, alignment=16) + state_vecs = [] + for i in cutlass.range_constexpr(cfg.d_k // 16): + state_block = [] + for g in cutlass.range_constexpr(16 // seed_vw): + seed_chunk = (seed_src + i * 16 + g * seed_vw).load(count=seed_vw, alignment=16) for t in cutlass.range_constexpr(seed_vw): - state_block[g * seed_vw + t] = seed_chunk[t].to(cutlass.Float32) - - nvvm.tcgen05_st( - "32x32b", - nvvm.make_tmem_ptr((row_id << 16) + (tmem_col + cfg.tmem_state_acc_offset + key_block_start), cutlass.Float32), - state_block[0:32], - ) - else: - for key_block_start in cutlass.range_constexpr(0, cfg.d_k, 32): - state_block = cutlass.Array(cutlass.Float32, 32, alignment=16) - for col in cutlass.range_constexpr(32): - state_block[col] = cutlass.Float32(0.0) + state_block.append(seed_chunk[t].to(cutlass.Float32)) + state_vecs.append(state_block) + for i in cutlass.range_constexpr(cfg.d_k // 16): + state_pack = cutlass.Array(cutlass.Int32, 8, alignment=16) + for packed_col in cutlass.range_constexpr(8): + state_pack[packed_col] = fp32_to_fp16(state_vecs[i][2 * packed_col], state_vecs[i][2 * packed_col + 1], dtype=cfg.io_dtype) nvvm.tcgen05_st( "32x32b", - nvvm.make_tmem_ptr((row_id << 16) + (tmem_col + cfg.tmem_state_acc_offset + key_block_start), cutlass.Float32), - state_block[0:32], - ) - if cutlass.const_expr(mState_init is not None): - nvvm.tcgen05_wait("store") - sV_ptr = sV_raw.data_ptr() + raw_index.idx * (cfg.d_v * cfg.b_t) - sW_ptr = sW_raw.data_ptr() + raw_index.idx * (cfg.d_v * cfg.b_t) - - row_addr = (tmem_row + tmem_subpartition * cfg.threads_per_warp) << 16 - state_col_id = tmem_col + cfg.tmem_state_acc_offset - - # ---- state repack: acc TMEM -> packed b16 TMEM ---------------------- - packed_col_id = tmem_col + cfg.tmem_state_inp_offset - if cutlass.const_expr(mState_init is not None): - state_vecs = [] - for k_block in cutlass.range_constexpr(cfg.d_k // 16): - state_vecs.append(nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + state_col_id + k_block * 16, cutlass.Float32), num=16)) - - for k_block in cutlass.range_constexpr(cfg.d_k // 16): - state_pack = cutlass.Array(cutlass.Int32, 8, alignment=16) - for packed_col in cutlass.range_constexpr(8): - source_pair = packed_col ^ 4 - state_pack[packed_col] = fp32_to_fp16( - state_vecs[k_block][2 * source_pair], state_vecs[k_block][2 * source_pair + 1], dtype=cfg.io_dtype + nvvm.make_tmem_ptr(row_lo_addr + packed_col_id + i * 8, cutlass.Int8), + state_pack[0:8], ) - nvvm.tcgen05_st( - "32x32b", - nvvm.make_tmem_ptr((tmem_row << 16) + packed_col_id + k_block * 8, cutlass.Int8), - state_pack[0:8], - ) - nvvm.tcgen05_wait("store") - bars.mb_state_inp_ready.arrive() - if cutlass.const_expr(cfg.enable_checkpoints): - if wstart == 0: - checkpoint_stage = checkpoint_done_index.idx - bars.mb_checkpoint_tmastg_done[checkpoint_stage].wait(checkpoint_done_index.phase) - checkpoint_done_index = advance(checkpoint_done_index, cfg.smem_checkpoint_stages) - checkpoint_stage_base = checkpoint_stage * (cfg.d_k * cfg.d_v) - for k_block in cutlass.range_constexpr(cfg.d_k // 16): - for g in cutlass.range_constexpr(2): - packs = tuple( - fp32_to_fp16(state_vecs[k_block][g * 8 + 2 * t], state_vecs[k_block][g * 8 + 2 * t + 1], dtype=cfg.io_dtype) - for t in range(4) - ) - dk = k_block * 16 + g * 8 - checkpoint_addr = ( - checkpoint_stage_base + (dk // 64) * (cfg.d_v * 64) + value_dim * 64 + swizzle_xor_128b(value_dim, dk % 64, elem_bytes=2) - ) - (sCheckpoint_ptr + checkpoint_addr).store( - cutlass.Vector.from_elements(packs, cutlass.Int32).bitcast(cfg.io_dtype), alignment=16 - ) - nvvm.tcgen05_wait("load") - nvvm.fence_proxy("async.shared", space="cta") - bars.mb_checkpoint_tmastg_ready[checkpoint_stage].arrive() - bars.mb_state_acc_read_done.arrive() + nvvm.tcgen05_wait("store") + bars.mb_state_input_ready.arrive() + if cutlass.const_expr(cfg.enable_checkpoints): + if write_start == 0: + checkpoint_stage = checkpoint_done_index.idx + bars.mb_checkpoint_tmastg_done[checkpoint_stage].wait(checkpoint_done_index.phase) + checkpoint_done_index = advance(checkpoint_done_index, cfg.smem_checkpoint_stages) + checkpoint_stage_base = checkpoint_stage * (cfg.d_k * cfg.d_v) + for i in cutlass.range_constexpr(cfg.d_k // 16): + for g in cutlass.range_constexpr(2): + packs = tuple( + fp32_to_fp16(state_vecs[i][g * 8 + 2 * t], state_vecs[i][g * 8 + 2 * t + 1], dtype=cfg.io_dtype) for t in range(4) + ) + dk = i * 16 + g * 8 + checkpoint_addr = ( + checkpoint_stage_base + + (dk // 64) * (cfg.d_v * 64) + + value_dim * 64 + + swizzle_xor_128b(value_dim, dk % 64, elem_bytes=2) + ) + (sCheckpoint_ptr + checkpoint_addr).store( + cutlass.Vector.from_elements(packs, cutlass.Int32).bitcast(cfg.io_dtype), alignment=16 + ) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_checkpoint_tmastg_ready[checkpoint_stage].arrive() + bars.mb_state_acc_read_done.arrive() if cutlass.const_expr(cfg.enable_checkpoints and mState_init is None): - if wstart == 0: + if write_start == 0: checkpoint_stage = checkpoint_done_index.idx bars.mb_checkpoint_tmastg_done[checkpoint_stage].wait(checkpoint_done_index.phase) checkpoint_done_index = advance(checkpoint_done_index, cfg.smem_checkpoint_stages) checkpoint_stage_base = checkpoint_stage * (cfg.d_k * cfg.d_v) zero_packs = tuple(cutlass.Int32(0) for _ in range(4)) - for k_block in cutlass.range_constexpr(cfg.d_k // 16): + for i in cutlass.range_constexpr(cfg.d_k // 16): for g in cutlass.range_constexpr(2): - dk = k_block * 16 + g * 8 + dk = i * 16 + g * 8 checkpoint_addr = ( checkpoint_stage_base + (dk // 64) * (cfg.d_v * 64) + value_dim * 64 + swizzle_xor_128b(value_dim, dk % 64, elem_bytes=2) ) @@ -1366,130 +1273,99 @@ def compute1_warp_group( nvvm.fence_proxy("async.shared", space="cta") bars.mb_checkpoint_tmastg_ready[checkpoint_stage].arrive() - # ---- Y staging: Y = W*V - state*(Beta*K) ----------------------------- + # ---- Y stage: Y = W*V - state*(Beta*K) ----------------------------------- bars.mb_v_ready[raw_bar_index.idx].wait(raw_bar_index.phase) projection_col_id = tmem_col + cfg.tmem_state_k_acc_offset - input_col_id = tmem_col + cfg.tmem_y_inp_offset + input_col_id = tmem_col + cfg.tmem_y_input_offset value_dim_base = tmem_subpartition * cfg.threads_per_warp - # ---- raw V fragments, then W, then the state*K acc readback ---------- - row_id0 = tmem_row + value_dim_base - row_id1 = row_id0 + 16 - raw_v_frag0 = nvvm.ldmatrix( + # ---- raw V fragments, then W, then the k state acc read ------------------ + raw_v_frag_lo = nvvm.ldmatrix( sV_ptr - + (value_dim_base + ldsm_col_offset) // 64 * (cfg.b_t * 64) - + ldsm_row_coord * 64 - + swizzle_xor_128b(ldsm_row_coord, (value_dim_base + ldsm_col_offset) % 64, elem_bytes=2), + + (value_dim_base + ov_col_offset) // 64 * (cfg.b_t * 64) + + ov_row_coord * 64 + + swizzle_xor_128b(ov_row_coord, (value_dim_base + ov_col_offset) % 64, elem_bytes=2), 4, nvvm.MMALayout.COL, ) - raw_v_frag1 = nvvm.ldmatrix( + raw_v_frag_hi = nvvm.ldmatrix( sV_ptr - + (value_dim_base + 16 + ldsm_col_offset) // 64 * (cfg.b_t * 64) - + ldsm_row_coord * 64 - + swizzle_xor_128b(ldsm_row_coord, (value_dim_base + 16 + ldsm_col_offset) % 64, elem_bytes=2), + + (value_dim_base + 16 + ov_col_offset) // 64 * (cfg.b_t * 64) + + ov_row_coord * 64 + + swizzle_xor_128b(ov_row_coord, (value_dim_base + 16 + ov_col_offset) % 64, elem_bytes=2), 4, nvvm.MMALayout.COL, ) bars.mb_w_ready[raw_bar_index.idx].wait(raw_bar_index.phase) - raw_w_frag0 = nvvm.ldmatrix( + raw_w_frag_lo = nvvm.ldmatrix( sW_ptr - + (value_dim_base + ldsm_col_offset) // 64 * (cfg.b_t * 64) - + ldsm_row_coord * 64 - + swizzle_xor_128b(ldsm_row_coord, (value_dim_base + ldsm_col_offset) % 64, elem_bytes=2), + + (value_dim_base + ov_col_offset) // 64 * (cfg.b_t * 64) + + ov_row_coord * 64 + + swizzle_xor_128b(ov_row_coord, (value_dim_base + ov_col_offset) % 64, elem_bytes=2), 4, nvvm.MMALayout.COL, ) - raw_w_frag1 = nvvm.ldmatrix( + raw_w_frag_hi = nvvm.ldmatrix( sW_ptr - + (value_dim_base + 16 + ldsm_col_offset) // 64 * (cfg.b_t * 64) - + ldsm_row_coord * 64 - + swizzle_xor_128b(ldsm_row_coord, (value_dim_base + 16 + ldsm_col_offset) % 64, elem_bytes=2), + + (value_dim_base + 16 + ov_col_offset) // 64 * (cfg.b_t * 64) + + ov_row_coord * 64 + + swizzle_xor_128b(ov_row_coord, (value_dim_base + 16 + ov_col_offset) % 64, elem_bytes=2), 4, nvvm.MMALayout.COL, ) - if cutlass.const_expr(mState_init is not None): - bars.mb_state_k_acc_ready.wait(state_k_acc_index.phase) - state_k_vec0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id0 << 16) + projection_col_id, cutlass.Float32), num=2) - state_k_vec1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id1 << 16) + projection_col_id, cutlass.Float32), num=2) - - y_inp_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + y_lo = [cutlass.Int32(0) for _ in range(4)] + y_hi = [cutlass.Int32(0) for _ in range(4)] for reg_idx in cutlass.range_constexpr(4): - raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1) - frag_pair = (reg_idx ^ 2) * 2 - w_pair = raw_w_frag0[raw_matrix] - wv_pair = mul_f16x2( - w_pair, - raw_v_frag0[raw_matrix], - cfg.io_dtype, - ) - if cutlass.const_expr(mState_init is not None): - state_k_val0, state_k_val1 = state_k_vec0[frag_pair], state_k_vec0[frag_pair + 1] - state_k_pair = fp32_to_fp16(state_k_val0, state_k_val1, dtype=cfg.io_dtype) - y_inp_pack0[reg_idx] = sub_f16x2( - wv_pair, - state_k_pair, - cfg.io_dtype, - ) - else: - y_inp_pack0[reg_idx] = wv_pair - - y_inp_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + y_lo[reg_idx] = mul_f16x2(raw_w_frag_lo[reg_idx], raw_v_frag_lo[reg_idx], cfg.io_dtype) + y_hi[reg_idx] = mul_f16x2(raw_w_frag_hi[reg_idx], raw_v_frag_hi[reg_idx], cfg.io_dtype) + if cutlass.const_expr(mState_init is not None): + if seed_from_initial_state: + bars.mb_state_k_acc_ready.wait(state_k_acc_index.phase) + state_k_acc_index = advance(state_k_acc_index, 1) + state_k_vec_lo = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_lo_addr + projection_col_id, cutlass.Float32), num=2) + state_k_vec_hi = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_hi_addr + projection_col_id, cutlass.Float32), num=2) + for reg_idx in cutlass.range_constexpr(4): + frag_pair = reg_idx * 2 + state_k_lo = fp32_to_fp16(state_k_vec_lo[frag_pair], state_k_vec_lo[frag_pair + 1], dtype=cfg.io_dtype) + state_k_hi = fp32_to_fp16(state_k_vec_hi[frag_pair], state_k_vec_hi[frag_pair + 1], dtype=cfg.io_dtype) + y_lo[reg_idx] = sub_f16x2(y_lo[reg_idx], state_k_lo, cfg.io_dtype) + y_hi[reg_idx] = sub_f16x2(y_hi[reg_idx], state_k_hi, cfg.io_dtype) + + y_input_pack_lo = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + y_input_pack_hi = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) for reg_idx in cutlass.range_constexpr(4): - raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1) - frag_pair = (reg_idx ^ 2) * 2 - w_pair = raw_w_frag1[raw_matrix] - wv_pair = mul_f16x2( - w_pair, - raw_v_frag1[raw_matrix], - cfg.io_dtype, - ) - if cutlass.const_expr(mState_init is not None): - state_k_val0, state_k_val1 = state_k_vec1[frag_pair], state_k_vec1[frag_pair + 1] - state_k_pair = fp32_to_fp16(state_k_val0, state_k_val1, dtype=cfg.io_dtype) - y_inp_pack1[reg_idx] = sub_f16x2( - wv_pair, - state_k_pair, - cfg.io_dtype, - ) - else: - y_inp_pack1[reg_idx] = wv_pair - - nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr((tmem_row << 16) + input_col_id, cutlass.Int8), y_inp_pack0[0:4]) - - nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr((tmem_row + 16 << 16) + input_col_id, cutlass.Int8), y_inp_pack1[0:4]) - + y_input_pack_lo[reg_idx] = y_lo[reg_idx] + y_input_pack_hi[reg_idx] = y_hi[reg_idx] + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_lo_addr + input_col_id, cutlass.Int8), y_input_pack_lo[0:4]) + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_hi_addr + input_col_id, cutlass.Int8), y_input_pack_hi[0:4]) nvvm.tcgen05_wait("store") - if cutlass.const_expr(mState_init is not None): - state_k_acc_index = advance(state_k_acc_index, 1) bars.mb_v_done[raw_index.idx].arrive() bars.mb_w_done[raw_index.idx].arrive() - bars.mb_y_inp_ready.arrive() + bars.mb_y_input_ready.arrive() - # ---- U repack: acc TMEM -> packed b16 TMEM -------------------------- + # ---- U stage: acc TMEM -> packed b16 TMEM -------------------------------- bars.mb_u_acc_ready.wait(u_acc_index.phase) u_vals = nvvm.tcgen05_ld( "32x32b", - nvvm.make_tmem_ptr((tmem_row + tmem_subpartition * cfg.threads_per_warp << 16) + (tmem_col + cfg.tmem_u_acc_offset), cutlass.Float32), + nvvm.make_tmem_ptr(row_lo_addr + (tmem_col + cfg.tmem_u_acc_offset), cutlass.Float32), num=cfg.b_t, ) - u_inp_pack = cutlass.Array(cutlass.Int32, (cfg.b_t // 2), alignment=16) + u_input_pack = cutlass.Array(cutlass.Int32, (cfg.b_t // 2), alignment=16) for packed_col in cutlass.range_constexpr((cfg.b_t // 2)): - source_pair = packed_col ^ 4 - token0 = source_pair * 2 + token0 = packed_col * 2 token1 = token0 + 1 - u_inp_pack[packed_col] = fp32_to_fp16(u_vals[token0], u_vals[token1], dtype=cfg.io_dtype) + u_input_pack[packed_col] = fp32_to_fp16(u_vals[token0], u_vals[token1], dtype=cfg.io_dtype) nvvm.tcgen05_st( "32x32b", - nvvm.make_tmem_ptr((tmem_row << 16) + (tmem_col + cfg.tmem_u_inp_offset), cutlass.Int8), - u_inp_pack[0 : (cfg.b_t // 2)], + nvvm.make_tmem_ptr(row_lo_addr + (tmem_col + cfg.tmem_u_input_offset), cutlass.Int8), + u_input_pack[0 : (cfg.b_t // 2)], ) nvvm.tcgen05_wait("store") u_acc_index = advance(u_acc_index, 1) - bars.mb_u_inp_ready.arrive() + bars.mb_u_input_ready.arrive() bars.mb_k_restore_acc_done[k_restore_index.idx].wait(k_restore_index.phase) k_restore_index = advance(k_restore_index, cfg.smem_decay_stages) @@ -1498,9 +1374,9 @@ def compute1_warp_group( if cutlass.const_expr(cfg.enable_checkpoints): cg1_checkpoint_chunks = checkpoint_every_n_tokens // cutlass.Int32(cfg.b_t) - cg1_checkpoint_mod = (cstart + cutlass.Int32(1)) % cg1_checkpoint_chunks + cg1_checkpoint_mod = (compute_start + cutlass.Int32(1)) % cg1_checkpoint_chunks for local_chunk in cutlass.range(1, num_chunks_tile, 1, unroll=1): - chunk_idx = cstart + local_chunk + chunk_idx = compute_start + local_chunk sV_ptr = sV_raw.data_ptr() + raw_index.idx * (cfg.d_v * cfg.b_t) sW_ptr = sW_raw.data_ptr() + raw_index.idx * (cfg.d_v * cfg.b_t) @@ -1509,44 +1385,42 @@ def compute1_warp_group( do_checkpoint = cg1_checkpoint_mod == 0 cg1_checkpoint_mod = cg1_checkpoint_mod + cutlass.Int32(1) cg1_checkpoint_mod = cutlass.Int32(0) if cg1_checkpoint_mod == cg1_checkpoint_chunks else cg1_checkpoint_mod - do_checkpoint = do_checkpoint and chunk_idx >= wstart - row_addr = (tmem_row + tmem_subpartition * cfg.threads_per_warp) << 16 + do_checkpoint = do_checkpoint and chunk_idx >= write_start state_col_id = tmem_col + cfg.tmem_state_acc_offset - # ---- state repack: acc TMEM -> packed b16 TMEM ---------------------- + # ---- state stage: acc TMEM -> packed b16 TMEM ---------------------------- state_vecs = [] - for k_block in cutlass.range_constexpr(cfg.d_k // 16): - state_vecs.append(nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + state_col_id + k_block * 16, cutlass.Float32), num=16)) + for i in cutlass.range_constexpr(cfg.d_k // 16): + state_vecs.append(nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_lo_addr + state_col_id + i * 16, cutlass.Float32), num=16)) - packed_col_id = tmem_col + cfg.tmem_state_inp_offset - for k_block in cutlass.range_constexpr(cfg.d_k // 16): + packed_col_id = tmem_col + cfg.tmem_state_input_offset + for i in cutlass.range_constexpr(cfg.d_k // 16): state_pack = cutlass.Array(cutlass.Int32, 8, alignment=16) for packed_col in cutlass.range_constexpr(8): - source_pair = packed_col ^ 4 - state_pack[packed_col] = fp32_to_fp16(state_vecs[k_block][2 * source_pair], state_vecs[k_block][2 * source_pair + 1], dtype=cfg.io_dtype) + state_pack[packed_col] = fp32_to_fp16(state_vecs[i][2 * packed_col], state_vecs[i][2 * packed_col + 1], dtype=cfg.io_dtype) nvvm.tcgen05_st( "32x32b", - nvvm.make_tmem_ptr((tmem_row << 16) + packed_col_id + k_block * 8, cutlass.Int8), + nvvm.make_tmem_ptr(row_lo_addr + packed_col_id + i * 8, cutlass.Int8), state_pack[0:8], ) nvvm.tcgen05_wait("store") - bars.mb_state_inp_ready.arrive() + bars.mb_state_input_ready.arrive() - # ---- checkpoint: post-publish f32 fragment read; the decay GEMM's ------ + # ---- checkpoint: post-publish f32 fragment read, ordered after the decay + # GEMM --------------------------------------------------------------------- if cutlass.const_expr(cfg.enable_checkpoints): if do_checkpoint: checkpoint_stage = checkpoint_done_index.idx bars.mb_checkpoint_tmastg_done[checkpoint_stage].wait(checkpoint_done_index.phase) checkpoint_done_index = advance(checkpoint_done_index, cfg.smem_checkpoint_stages) checkpoint_stage_base = checkpoint_stage * (cfg.d_k * cfg.d_v) - row16_addr = ((tmem_row + tmem_subpartition * cfg.threads_per_warp) + 16) << 16 checkpoint_vbase = tmem_subpartition * cfg.threads_per_warp - checkpoint_swz_off0 = (checkpoint_vbase + ldsm_col_offset) // 64 * (cfg.d_k * 64) - checkpoint_swz_col0 = (checkpoint_vbase + ldsm_col_offset) % 64 - checkpoint_swz_off = (checkpoint_vbase + 16 + ldsm_col_offset) // 64 * (cfg.d_k * 64) - checkpoint_swz_col = (checkpoint_vbase + 16 + ldsm_col_offset) % 64 + checkpoint_swz_off_lo = (checkpoint_vbase + ov_col_offset) // 64 * (cfg.d_k * 64) + checkpoint_swz_col_lo = (checkpoint_vbase + ov_col_offset) % 64 + checkpoint_swz_off_hi = (checkpoint_vbase + 16 + ov_col_offset) // 64 * (cfg.d_k * 64) + checkpoint_swz_col_hi = (checkpoint_vbase + 16 + ov_col_offset) % 64 for k_base in cutlass.range_constexpr(0, cfg.d_k, 32): - checkpoint_vec = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + state_col_id + k_base, cutlass.Float32), num=32) + checkpoint_vec = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_lo_addr + state_col_id + k_base, cutlass.Float32), num=32) for g in cutlass.range_constexpr(4): packs = tuple(fp32_to_fp16(checkpoint_vec[g * 8 + 2 * t], checkpoint_vec[g * 8 + 2 * t + 1], dtype=cfg.io_dtype) for t in range(4)) dk = k_base + g * 8 @@ -1560,140 +1434,135 @@ def compute1_warp_group( bars.mb_checkpoint_tmastg_ready[checkpoint_stage].arrive() else: bars.mb_state_acc_read_done.arrive() - # ---- Y staging: Y = W*V - state*(Beta*K) ----------------------------- + # ---- Y stage: Y = W*V - state*(Beta*K) ----------------------------------- bars.mb_v_ready[raw_bar_index.idx].wait(raw_bar_index.phase) projection_col_id = tmem_col + cfg.tmem_state_k_acc_offset - input_col_id = tmem_col + cfg.tmem_y_inp_offset + input_col_id = tmem_col + cfg.tmem_y_input_offset value_dim_base = tmem_subpartition * cfg.threads_per_warp - # ---- raw V fragments, then W, then the state*K acc readback ---------- - row_id0 = tmem_row + value_dim_base - row_id1 = row_id0 + 16 - raw_v_frag0 = nvvm.ldmatrix( + # ---- raw V fragments, then W, then the k state acc read ------------------ + raw_v_frag_lo = nvvm.ldmatrix( sV_ptr - + (value_dim_base + ldsm_col_offset) // 64 * (cfg.b_t * 64) - + ldsm_row_coord * 64 - + swizzle_xor_128b(ldsm_row_coord, (value_dim_base + ldsm_col_offset) % 64, elem_bytes=2), + + (value_dim_base + ov_col_offset) // 64 * (cfg.b_t * 64) + + ov_row_coord * 64 + + swizzle_xor_128b(ov_row_coord, (value_dim_base + ov_col_offset) % 64, elem_bytes=2), 4, nvvm.MMALayout.COL, ) - raw_v_frag1 = nvvm.ldmatrix( + raw_v_frag_hi = nvvm.ldmatrix( sV_ptr - + (value_dim_base + 16 + ldsm_col_offset) // 64 * (cfg.b_t * 64) - + ldsm_row_coord * 64 - + swizzle_xor_128b(ldsm_row_coord, (value_dim_base + 16 + ldsm_col_offset) % 64, elem_bytes=2), + + (value_dim_base + 16 + ov_col_offset) // 64 * (cfg.b_t * 64) + + ov_row_coord * 64 + + swizzle_xor_128b(ov_row_coord, (value_dim_base + 16 + ov_col_offset) % 64, elem_bytes=2), 4, nvvm.MMALayout.COL, ) bars.mb_w_ready[raw_bar_index.idx].wait(raw_bar_index.phase) - raw_w_frag0 = nvvm.ldmatrix( + raw_w_frag_lo = nvvm.ldmatrix( sW_ptr - + (value_dim_base + ldsm_col_offset) // 64 * (cfg.b_t * 64) - + ldsm_row_coord * 64 - + swizzle_xor_128b(ldsm_row_coord, (value_dim_base + ldsm_col_offset) % 64, elem_bytes=2), + + (value_dim_base + ov_col_offset) // 64 * (cfg.b_t * 64) + + ov_row_coord * 64 + + swizzle_xor_128b(ov_row_coord, (value_dim_base + ov_col_offset) % 64, elem_bytes=2), 4, nvvm.MMALayout.COL, ) - raw_w_frag1 = nvvm.ldmatrix( + raw_w_frag_hi = nvvm.ldmatrix( sW_ptr - + (value_dim_base + 16 + ldsm_col_offset) // 64 * (cfg.b_t * 64) - + ldsm_row_coord * 64 - + swizzle_xor_128b(ldsm_row_coord, (value_dim_base + 16 + ldsm_col_offset) % 64, elem_bytes=2), + + (value_dim_base + 16 + ov_col_offset) // 64 * (cfg.b_t * 64) + + ov_row_coord * 64 + + swizzle_xor_128b(ov_row_coord, (value_dim_base + 16 + ov_col_offset) % 64, elem_bytes=2), 4, nvvm.MMALayout.COL, ) bars.mb_state_k_acc_ready.wait(state_k_acc_index.phase) - state_k_vec0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id0 << 16) + projection_col_id, cutlass.Float32), num=2) - state_k_vec1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id1 << 16) + projection_col_id, cutlass.Float32), num=2) + state_k_vec_lo = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_lo_addr + projection_col_id, cutlass.Float32), num=2) + state_k_vec_hi = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_hi_addr + projection_col_id, cutlass.Float32), num=2) - y_inp_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + y_input_pack_lo = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) for reg_idx in cutlass.range_constexpr(4): - raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1) - frag_pair = (reg_idx ^ 2) * 2 - state_k_val0, state_k_val1 = state_k_vec0[frag_pair], state_k_vec0[frag_pair + 1] + frag_pair = reg_idx * 2 + state_k_val0, state_k_val1 = state_k_vec_lo[frag_pair], state_k_vec_lo[frag_pair + 1] state_k_pair = fp32_to_fp16(state_k_val0, state_k_val1, dtype=cfg.io_dtype) - w_pair = raw_w_frag0[raw_matrix] + w_pair = raw_w_frag_lo[reg_idx] wv_pair = mul_f16x2( w_pair, - raw_v_frag0[raw_matrix], + raw_v_frag_lo[reg_idx], cfg.io_dtype, ) - y_inp_pack0[reg_idx] = sub_f16x2( + y_input_pack_lo[reg_idx] = sub_f16x2( wv_pair, state_k_pair, cfg.io_dtype, ) - y_inp_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + y_input_pack_hi = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) for reg_idx in cutlass.range_constexpr(4): - raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1) - frag_pair = (reg_idx ^ 2) * 2 - state_k_val0, state_k_val1 = state_k_vec1[frag_pair], state_k_vec1[frag_pair + 1] + frag_pair = reg_idx * 2 + state_k_val0, state_k_val1 = state_k_vec_hi[frag_pair], state_k_vec_hi[frag_pair + 1] state_k_pair = fp32_to_fp16(state_k_val0, state_k_val1, dtype=cfg.io_dtype) - w_pair = raw_w_frag1[raw_matrix] + w_pair = raw_w_frag_hi[reg_idx] wv_pair = mul_f16x2( w_pair, - raw_v_frag1[raw_matrix], + raw_v_frag_hi[reg_idx], cfg.io_dtype, ) - y_inp_pack1[reg_idx] = sub_f16x2( + y_input_pack_hi[reg_idx] = sub_f16x2( wv_pair, state_k_pair, cfg.io_dtype, ) - nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr((tmem_row << 16) + input_col_id, cutlass.Int8), y_inp_pack0[0:4]) + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_lo_addr + input_col_id, cutlass.Int8), y_input_pack_lo[0:4]) - nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr((tmem_row + 16 << 16) + input_col_id, cutlass.Int8), y_inp_pack1[0:4]) + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_hi_addr + input_col_id, cutlass.Int8), y_input_pack_hi[0:4]) nvvm.tcgen05_wait("store") state_k_acc_index = advance(state_k_acc_index, 1) bars.mb_v_done[raw_index.idx].arrive() bars.mb_w_done[raw_index.idx].arrive() - bars.mb_y_inp_ready.arrive() + bars.mb_y_input_ready.arrive() - # ---- U repack: acc TMEM -> packed b16 TMEM -------------------------- + # ---- U stage: acc TMEM -> packed b16 TMEM -------------------------------- bars.mb_u_acc_ready.wait(u_acc_index.phase) u_vals = nvvm.tcgen05_ld( "32x32b", - nvvm.make_tmem_ptr((tmem_row + tmem_subpartition * cfg.threads_per_warp << 16) + (tmem_col + cfg.tmem_u_acc_offset), cutlass.Float32), + nvvm.make_tmem_ptr(row_lo_addr + (tmem_col + cfg.tmem_u_acc_offset), cutlass.Float32), num=cfg.b_t, ) - u_inp_pack = cutlass.Array(cutlass.Int32, (cfg.b_t // 2), alignment=16) + u_input_pack = cutlass.Array(cutlass.Int32, (cfg.b_t // 2), alignment=16) for packed_col in cutlass.range_constexpr((cfg.b_t // 2)): - source_pair = packed_col ^ 4 - token0 = source_pair * 2 + token0 = packed_col * 2 token1 = token0 + 1 - u_inp_pack[packed_col] = fp32_to_fp16(u_vals[token0], u_vals[token1], dtype=cfg.io_dtype) + u_input_pack[packed_col] = fp32_to_fp16(u_vals[token0], u_vals[token1], dtype=cfg.io_dtype) nvvm.tcgen05_st( "32x32b", - nvvm.make_tmem_ptr((tmem_row << 16) + (tmem_col + cfg.tmem_u_inp_offset), cutlass.Int8), - u_inp_pack[0 : (cfg.b_t // 2)], + nvvm.make_tmem_ptr(row_lo_addr + (tmem_col + cfg.tmem_u_input_offset), cutlass.Int8), + u_input_pack[0 : (cfg.b_t // 2)], ) nvvm.tcgen05_wait("store") u_acc_index = advance(u_acc_index, 1) - bars.mb_u_inp_ready.arrive() + bars.mb_u_input_ready.arrive() bars.mb_k_restore_acc_done[k_restore_index.idx].wait(k_restore_index.phase) k_restore_index = advance(k_restore_index, cfg.smem_decay_stages) raw_index = advance(raw_index, cfg.smem_raw_stages) raw_bar_index = advance(raw_bar_index, cfg.smem_raw_bar_stages) - owns_final = wend == num_chunks_b + owns_final = write_end == batch_num_chunks - # ---- final-state drain: state acc TMEM -> GMEM --------------------------- + # ---- final state store: TMEM -> GMEM ----------------------------------------- if cutlass.const_expr(mState_out is not None): - if seqlen_b > 0: + if batch_seqlen > 0: if owns_final: state_vw = 16 // (mState_out.element_type.width // 8) state_dst = (mState_out.iterator + mState_out.layout((batch_idx, head_o, value_dim, 0))).raw_ptr() for key_block_start in cutlass.range_constexpr(0, cfg.d_k, 32): state_vec = nvvm.tcgen05_ld( "32x32b", - nvvm.make_tmem_ptr((row_id << 16) + (tmem_col + cfg.tmem_state_acc_offset + key_block_start), cutlass.Float32), + nvvm.make_tmem_ptr(row_lo_addr + (tmem_col + cfg.tmem_state_acc_offset + key_block_start), cutlass.Float32), num=32, ) @@ -1713,7 +1582,7 @@ def compute1_warp_group( mState_out[batch_idx, head_o, value_dim, key_dim] = mState_init[batch_idx, head_o, value_dim, key_dim] else: mState_out[batch_idx, head_o, value_dim, key_dim] = cutlass.Float32(0.0).to(mState_out.element_type) - tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one) bars.mb_tmem_done[0].arrive() @@ -1733,7 +1602,7 @@ def host( final_state: cute.Tensor | None, work_items: cute.Tensor | None, work_count: cute.Tensor | None, - sched_ctr: cute.Tensor | None, + scheduler_counter: cute.Tensor | None, tensormap_workspace: cute.Tensor, checkpoint_every_n_tokens: cutlass.Int32, stream, @@ -1757,7 +1626,7 @@ def host( final_state, work_items, work_count, - sched_ctr, + scheduler_counter, checkpoint_every_n_tokens, ).launch( grid=grid_shape, @@ -1784,7 +1653,7 @@ def kernel( mState_out: cute.Tensor | None, mWorkItems: cute.Tensor, mCount: cute.Tensor, - mSched: cute.Tensor | None, + mScheduler: cute.Tensor | None, checkpoint_every_n_tokens: cutlass.Int32, ) -> None: """BT=16 GDN-2 recompute device kernel (persistent); grid @@ -1794,11 +1663,11 @@ def kernel( bidx = cute.arch.block_idx()[0] num_ctas = cute.arch.grid_dim()[0] warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - lane = tidx % cfg.threads_per_warp + lane_idx = tidx % cfg.threads_per_warp total_tiles = mCount[0] - if cutlass.const_expr(cfg.dyn_sched): - assert mSched is not None and mSched.element_type == cutlass.Int32 + if cutlass.const_expr(cfg.dynamic_scheduling): + assert mScheduler is not None and mScheduler.element_type == cutlass.Int32 assert mK.element_type == cfg.io_dtype and mV.element_type == cfg.io_dtype assert mGate.element_type == cutlass.Float32 assert mBeta.element_type == cfg.io_dtype and mW.element_type == cfg.io_dtype, "channel-wise beta/w must match the io dtype" @@ -1823,11 +1692,10 @@ def kernel( desc_w_base = desc_base_words + cutlass.Int32(4) * arr_words desc_checkpoint_base = desc_base_words + cutlass.Int32(5) * arr_words - # Buffers are declaration-ordered and intentionally non-aliased. SMEM = cutlass.AddressSpace.smem bars = make_gdn2_bars(cfg) tmem_base_holder = cutlass.Array(cutlass.Int32, 1, space=SMEM, alignment=4) - sSched = cutlass.Array(cutlass.Int32, cfg.sched_stages, space=SMEM, alignment=16) + sScheduler = cutlass.Array(cutlass.Int32, cfg.scheduler_stages, space=SMEM, alignment=16) sK_decay_raw = cutlass.Array(cfg.io_dtype, cfg.k_decay_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) sK_restore_raw = cutlass.Array(cfg.io_dtype, cfg.k_restore_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) sIntermediate_raw = cutlass.Array(cfg.io_dtype, cfg.intermediate_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) @@ -1851,7 +1719,7 @@ def kernel( stride_byte_offset=1024, layout=nvvm.Tcgen05SmemSwizzle.SWIZZLE_128B, ) - sK_restore = SmemTile( + sK_restore_trans = SmemTile( base=sK_restore_raw, elems_per_stage=(cfg.d_k * cfg.b_t), stages=cfg.smem_decay_stages, @@ -1895,15 +1763,15 @@ def kernel( if elect_one: bars.mb_state_k_acc_ready.init() bars.mb_u_acc_ready.init() - bars.mb_state_inp_ready.init() + bars.mb_state_input_ready.init() for stage in cutlass.range_constexpr(cfg.smem_state_scale_diag_stages): bars.mb_state_scale_diag_done[stage].init() for stage in cutlass.range_constexpr(cfg.smem_decay_stages): bars.mb_decay_tcgen05_done[stage].init() bars.mb_decay_super_done[stage].init() bars.mb_k_restore_acc_done[stage].init() - bars.mb_y_inp_ready.init() - bars.mb_u_inp_ready.init() + bars.mb_y_input_ready.init() + bars.mb_u_input_ready.init() bars.mb_tmem_done[0].init() elif warp_idx == cfg.super_mma_warp_id: if elect_one: @@ -1916,9 +1784,9 @@ def kernel( bars.mb_k_decay_inv_cg0_ready[stage].init() elif warp_idx == cfg.epilogue_warp_id: if elect_one: - for stage in cutlass.range_constexpr(cfg.sched_stages): - bars.mb_sched_ready[stage].init() - bars.mb_sched_done[stage].init() + for stage in cutlass.range_constexpr(cfg.scheduler_stages): + bars.mb_scheduler_ready[stage].init() + bars.mb_scheduler_done[stage].init() if cutlass.const_expr(cfg.enable_checkpoints): for stage in cutlass.range_constexpr(cfg.smem_checkpoint_stages): bars.mb_checkpoint_tmastg_ready[stage].init() @@ -1937,9 +1805,9 @@ def kernel( num_ctas, cu_seqlens, mWorkItems, - mSched, - sSched, - lane, + mScheduler, + sScheduler, + lane_idx, sBeta_raw, sGate_raw, sK_raw, @@ -1960,8 +1828,8 @@ def kernel( num_ctas, cu_seqlens, mWorkItems, - sSched, - lane, + sScheduler, + lane_idx, sK_inv_raw, sIntermediate_raw, sK_decay_raw, @@ -1975,11 +1843,11 @@ def kernel( num_ctas, cu_seqlens, mWorkItems, - sSched, + sScheduler, tmem_base_holder, sIntermediate, sK_decay, - sK_restore, + sK_restore_trans, sState_scale_diag, bars, ) @@ -1991,7 +1859,7 @@ def kernel( num_ctas, cu_seqlens, mWorkItems, - sSched, + sScheduler, sCheckpoint_raw, desc_checkpoint_base, checkpoint_every_n_tokens, @@ -2005,8 +1873,8 @@ def kernel( num_ctas, cu_seqlens, mWorkItems, - sSched, - lane, + sScheduler, + lane_idx, warp_idx, mA_log, mDt_bias, @@ -2029,8 +1897,8 @@ def kernel( num_ctas, cu_seqlens, mWorkItems, - sSched, - lane, + sScheduler, + lane_idx, tmem_base_holder, warp_idx, mState_out, @@ -2063,8 +1931,8 @@ class Gdn2RecomputeCfg: v_ratio: int n_heads_out: int max_active_clusters: int - dyn_sched: bool = False - sched_stages: int = CFG.SMEM_SCHED_STAGES + dynamic_scheduling: bool = False + scheduler_stages: int = CFG.SMEM_SCHEDULER_STAGES compute_group_0_warp_ids: tuple[int, ...] = CFG.COMPUTE_GROUP_0_WARP_IDS compute_group_1_warp_ids: tuple[int, ...] = CFG.COMPUTE_GROUP_1_WARP_IDS @@ -2081,7 +1949,7 @@ class Gdn2RecomputeCfg: cg0_group_count: int = 2 cg0_warps_per_group: int = 4 cg0_threads_per_group: int = 0 - cg0_group_sync_barrier_base_id: int = 1 # CG0 group g syncs on nbar id 1 + g + cg0_group_sync_barrier_base_id: int = 1 # CG0 group g syncs on named-barrier id 1 + g cg0_tile_entry_barrier_id: int = 5 # CG0-wide (both groups) work-item entry sync tmem_user_threads: int = 0 tmem_lifecycle_barrier_id: int = 3 @@ -2089,24 +1957,24 @@ class Gdn2RecomputeCfg: num_regs_compute_group_1: int = CFG.NUM_REGS_COMPUTE_GROUP_1 num_regs_other: int = CFG.NUM_REGS_OTHER - # ---- SMEM / TMEM ring stage counts ------------------------------------------- + # ---- SMEM / TMEM ring stage counts ----------------------------------------------- smem_raw_stages: int = CFG.SMEM_RAW_STAGES - smem_raw_bar_stages: int = 0 # ready-ring mbar depth: raw rounded up to even (CG0 ping-pong parity) + smem_raw_bar_stages: int = 0 # ready-ring mbar depth: raw rounded up to even smem_checkpoint_stages: int = 1 smem_decay_stages: int = CFG.SMEM_DECAY_STAGES smem_intermediate_stages: int = CFG.SMEM_INTERMEDIATE_STAGES smem_state_scale_diag_stages: int = CFG.SMEM_STATE_SCALE_DIAG_STAGES qk_scale_ready_stages: int = CFG.QK_SCALE_READY_STAGES - # ---- TMEM column offsets (state doubles as the final_state acc) -------------- + # ---- TMEM column offsets (state doubles as the final state acc) ------------------ tmem_state_acc_offset: int = 0 - tmem_state_inp_offset: int = 0 + tmem_state_input_offset: int = 0 tmem_state_k_acc_offset: int = 0 tmem_u_acc_offset: int = 0 - tmem_y_inp_offset: int = 0 - tmem_u_inp_offset: int = 0 + tmem_y_input_offset: int = 0 + tmem_u_input_offset: int = 0 - # ---- SMEM buffer cosizes ----------------------------------------------------- + # ---- SMEM buffer cosizes --------------------------------------------------------- k_cosize: int = 0 v_cosize: int = 0 gate_cosize: int = 0 @@ -2141,7 +2009,7 @@ def build_cfg( v_ratio: int, n_heads_out: int, max_active_clusters: int, - dyn_sched: bool = False, + dynamic_scheduling: bool = False, ) -> Gdn2RecomputeCfg: """Build the per-compile ``Gdn2RecomputeCfg`` (io_dtype in {Float16, BFloat16}); fills the derived TMEM column offsets and SMEM buffer cosizes.""" @@ -2161,7 +2029,7 @@ def build_cfg( v_ratio=v_ratio, n_heads_out=n_heads_out, max_active_clusters=max_active_clusters, - dyn_sched=dyn_sched, + dynamic_scheduling=dynamic_scheduling, ) if enable_checkpoints: cfg.smem_raw_stages = 4 @@ -2173,12 +2041,12 @@ def build_cfg( if cfg.smem_state_scale_diag_stages != cfg.qk_scale_ready_stages: raise ValueError("diag and qk-scale ready rings must share their rolling stage") - cfg.tmem_state_inp_offset = cfg.tmem_state_acc_offset + cfg.d_k - cfg.tmem_state_k_acc_offset = cfg.tmem_state_inp_offset + (cfg.d_k // 2) + cfg.tmem_state_input_offset = cfg.tmem_state_acc_offset + cfg.d_k + cfg.tmem_state_k_acc_offset = cfg.tmem_state_input_offset + (cfg.d_k // 2) cfg.tmem_u_acc_offset = cfg.tmem_state_k_acc_offset + cfg.b_t - cfg.tmem_y_inp_offset = cfg.tmem_u_acc_offset + cfg.b_t - cfg.tmem_u_inp_offset = cfg.tmem_y_inp_offset + (cfg.b_t // 2) - assert (cfg.tmem_u_inp_offset + (cfg.b_t // 2)) <= 512 + cfg.tmem_y_input_offset = cfg.tmem_u_acc_offset + cfg.b_t + cfg.tmem_u_input_offset = cfg.tmem_y_input_offset + (cfg.b_t // 2) + assert (cfg.tmem_u_input_offset + (cfg.b_t // 2)) <= 512 cfg.k_cosize = cfg.smem_raw_stages * cfg.d_k * cfg.b_t cfg.v_cosize = cfg.smem_raw_stages * cfg.d_v * cfg.b_t @@ -2270,7 +2138,7 @@ def build_descs_body( def prologue_kernel( run_order: cutlass.Constexpr[bool], order_gen: cutlass.Constexpr[bool], - has_sched: cutlass.Constexpr[bool], + has_scheduler: cutlass.Constexpr[bool], b_t: cutlass.Constexpr[int], base_k: cutlass.GridConstant[cuda.tensor_map.TensorMap], base_v: cutlass.GridConstant[cuda.tensor_map.TensorMap], @@ -2289,7 +2157,7 @@ def prologue_kernel( mStaging: cute.Tensor | None, mCount: cute.Tensor, mWorkItems: cute.Tensor, - mSched: cute.Tensor | None, + mScheduler: cute.Tensor | None, n_batch: cutlass.Int32, k_row_stride: cutlass.Int32, v_row_stride: cutlass.Int32, @@ -2301,7 +2169,7 @@ def prologue_kernel( ) -> None: """Single-CTA prologue. Under ``run_order`` this kernel is the first work-item-table consumer, so it LPT-orders the table and zeroes both - consumers' sched rings via :func:`order_body`; it then builds the + consumers' scheduler rings via :func:`order_body`; it then builds the per-batch TMA-descriptor arrays via :func:`build_descs_body`, one warp per array (the extra warps only take part in the order phase).""" tidx, _, _ = cute.arch.thread_idx() @@ -2314,10 +2182,10 @@ def prologue_kernel( n_heads_out = cutlass.Int32(gate.shape[1]) order_body( order_gen, - has_sched, + has_scheduler, b_t, ORDER_THREADS, - ORDER_ELEMS, + ORDER_ELEMENTS, tidx, n_heads_out, n_heads_out * n_batch, @@ -2325,7 +2193,7 @@ def prologue_kernel( mStaging, mCount, mWorkItems, - mSched, + mScheduler, sKey, sIdx, sSpread, @@ -2363,7 +2231,7 @@ def prologue( b_t: cutlass.Constexpr[int], run_order: cutlass.Constexpr[bool], order_gen: cutlass.Constexpr[bool], - has_sched: cutlass.Constexpr[bool], + has_scheduler: cutlass.Constexpr[bool], k: cute.Tensor, v: cute.Tensor, gate: cute.Tensor, @@ -2374,7 +2242,7 @@ def prologue( work_item_staging: cute.Tensor | None, work_count: cute.Tensor, work_items: cute.Tensor, - sched_all: cute.Tensor | None, + scheduler_all: cute.Tensor | None, tensormap_workspace: cute.Tensor, checkpoint_every_n: cutlass.Int32, stream: cuda_driver.CUstream, @@ -2390,7 +2258,7 @@ def prologue( d_k = k.shape[2] d_v = v.shape[2] bpe = io_dtype.width // 8 - tma_box_elems = 128 // bpe + tma_box_elements = 128 // bpe seqlen = k.shape[0] k_headed = cute.make_tensor(k.iterator, cute.make_layout((d_k, h_k, seqlen), stride=(1, k.stride[1], k.stride[0]))) @@ -2400,11 +2268,11 @@ def prologue( w_headed = cute.make_tensor(w.iterator, cute.make_layout((d_v, ho, seqlen), stride=(1, w.stride[1], w.stride[0]))) swz = cuda.TensorMapSwizzle.s128b - base_k = cuda.create_tensor_map_tiled_from_view(k_headed, box_dims=(tma_box_elems, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) - base_v = cuda.create_tensor_map_tiled_from_view(v_headed, box_dims=(tma_box_elems, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_k = cuda.create_tensor_map_tiled_from_view(k_headed, box_dims=(tma_box_elements, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_v = cuda.create_tensor_map_tiled_from_view(v_headed, box_dims=(tma_box_elements, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) base_gate = cuda.create_tensor_map_tiled_from_view(gate_headed, box_dims=(32, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) - base_beta = cuda.create_tensor_map_tiled_from_view(beta_headed, box_dims=(tma_box_elems, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) - base_w = cuda.create_tensor_map_tiled_from_view(w_headed, box_dims=(tma_box_elems, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_beta = cuda.create_tensor_map_tiled_from_view(beta_headed, box_dims=(tma_box_elements, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_w = cuda.create_tensor_map_tiled_from_view(w_headed, box_dims=(tma_box_elements, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) base_checkpoint = base_k if cutlass.const_expr(state_checkpoints is not None): @@ -2415,11 +2283,13 @@ def prologue( stride=(state_checkpoints.stride[3], state_checkpoints.stride[2], state_checkpoints.stride[0], state_checkpoints.stride[1]), ), ) - base_checkpoint = cuda.create_tensor_map_tiled_from_view(checkpoint_view, box_dims=(tma_box_elems, d_k, 1, 1), stride_order=(0, 1, 2, 3), swizzle=swz) + base_checkpoint = cuda.create_tensor_map_tiled_from_view( + checkpoint_view, box_dims=(tma_box_elements, d_k, 1, 1), stride_order=(0, 1, 2, 3), swizzle=swz + ) prologue_kernel( run_order, order_gen, - has_sched, + has_scheduler, b_t, base_k, base_v, @@ -2438,7 +2308,7 @@ def prologue( work_item_staging, work_count, work_items, - sched_all, + scheduler_all, cutlass.Int32(batch_size), cutlass.Int32(k.stride[0]), cutlass.Int32(v.stride[0]), @@ -2450,7 +2320,7 @@ def prologue( ).launch(grid=(1, 1, 1), block=(ORDER_THREADS, 1, 1), stream=stream) -# ---- Torch adapter / host-side compilation --------------------------------------- +# ---- Torch adapter / host-side compilation ------------------------------------------- @lru_cache(maxsize=None) @@ -2468,10 +2338,10 @@ def get_compiled_cache( safe_gate: bool, gate_lower_bound: float, beta_sigmoid: bool, - dyn_sched: bool, + dynamic_scheduling: bool, order_in_prologue: bool, order_gen: bool, - has_sched: bool, + has_scheduler: bool, ): """Return a mutable dict that lazily stores the compiled kernel.""" return {} @@ -2490,7 +2360,7 @@ def compile( k_ratio: int, v_ratio: int, n_heads_out: int, - dyn_sched: bool = False, + dynamic_scheduling: bool = False, *, num_sm: int, k_cute, @@ -2505,8 +2375,8 @@ def compile( state_out_cute, work_items_cute=None, work_count_cute=None, - sched_ctr_cute=None, - tensormap_ws_cute, + scheduler_counter_cute=None, + tensormap_workspace_cute, checkpoint_every_n_tokens, stream, ): @@ -2525,7 +2395,7 @@ def compile( v_ratio=v_ratio, n_heads_out=n_heads_out, max_active_clusters=num_sm, - dyn_sched=dyn_sched, + dynamic_scheduling=dynamic_scheduling, ) return cute.compile( @@ -2543,8 +2413,8 @@ def compile( state_out_cute, work_items_cute, work_count_cute, - sched_ctr_cute, - tensormap_ws_cute, + scheduler_counter_cute, + tensormap_workspace_cute, checkpoint_every_n_tokens, stream, options="--enable-tvm-ffi --opt-level 2", @@ -2570,8 +2440,8 @@ def chunk_gdn2_recompute_sm100( use_beta_sigmoid: bool = False, work_items=None, work_count=None, - sched_ctr=None, - sched_all=None, + scheduler_counter=None, + scheduler_all=None, work_item_scratch=None, order_in_prologue: bool = False, *, @@ -2611,8 +2481,8 @@ def chunk_gdn2_recompute_sm100( use_beta_sigmoid: ``beta`` holds logits; sigmoid in-kernel work_items: ``(max_items, 8)`` int32 work-item table from ``common/split_k.py`` (REQUIRED; an uncut table row is the whole - (b, h) sequence). Each item computes chunks ``[cstart, wend)`` - and writes checkpoints only for ``[wstart, wend)``. + (b, h) sequence). Each item computes chunks ``[compute_start, write_end)`` + and writes checkpoints only for ``[write_start, write_end)``. work_count: ``(1,)`` int32 device-side item count (REQUIRED) """ HK = k.shape[1] @@ -2629,11 +2499,13 @@ def chunk_gdn2_recompute_sm100( f"output_state_checkpoints dtype must match the io dtype (fp32 state belongs to output_state): got {output_state_checkpoints.dtype} with io {k.dtype}" ) if work_items is None or work_count is None: - raise ValueError("work_items/work_count are required (the split-table stage builds them for every launch)") - dyn_sched = sched_ctr is not None + raise ValueError( + "work_items/work_count are required (built by the split-table stage, or by an order-generating prologue when work_item_scratch is None)" + ) + dynamic_scheduling = scheduler_counter is not None order_gen = work_item_scratch is None - if order_in_prologue and sched_all is None: - raise ValueError("order_in_prologue requires sched_all (the prologue zeroes both consumers' sched rings)") + if order_in_prologue and scheduler_all is None: + raise ValueError("order_in_prologue requires scheduler_all (the prologue zeroes both consumers' scheduler rings)") if initial_state is not None: state_dtype_src = initial_state.dtype @@ -2670,10 +2542,10 @@ def chunk_gdn2_recompute_sm100( safe_gate, gate_lower_bound, use_beta_sigmoid, - dyn_sched, + dynamic_scheduling, order_in_prologue, order_gen, - sched_all is not None, + scheduler_all is not None, ) if "compiled" not in cache: @@ -2700,11 +2572,11 @@ def chunk_gdn2_recompute_sm100( work_items_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1), divisibility=1) work_count_cute = from_dlpack(work_count, assumed_align=4).mark_layout_dynamic() - sched_ctr_cute = None - if dyn_sched: - sched_ctr_cute = from_dlpack(sched_ctr, assumed_align=4).mark_layout_dynamic() + scheduler_counter_cute = None + if dynamic_scheduling: + scheduler_counter_cute = from_dlpack(scheduler_counter, assumed_align=4).mark_layout_dynamic() - tensormap_ws_cute = from_dlpack(tensormap_workspace, assumed_align=128).mark_layout_dynamic() + tensormap_workspace_cute = from_dlpack(tensormap_workspace, assumed_align=128).mark_layout_dynamic() cache["compiled"] = compile( io_dtype, @@ -2719,7 +2591,7 @@ def chunk_gdn2_recompute_sm100( k_ratio, v_ratio, HO, - dyn_sched, + dynamic_scheduling, num_sm=multiprocessor_count(current_device()), k_cute=k_cute, v_cute=v_cute, @@ -2733,8 +2605,8 @@ def chunk_gdn2_recompute_sm100( state_out_cute=state_out_cute, work_items_cute=work_items_cute, work_count_cute=work_count_cute, - sched_ctr_cute=sched_ctr_cute, - tensormap_ws_cute=tensormap_ws_cute, + scheduler_counter_cute=scheduler_counter_cute, + tensormap_workspace_cute=tensormap_workspace_cute, checkpoint_every_n_tokens=checkpoint_every_n_tokens, stream=cu_stream, ) @@ -2748,7 +2620,7 @@ def chunk_gdn2_recompute_sm100( beta_pl = from_dlpack(beta, assumed_align=16).mark_layout_dynamic(leading_dim=2) w_pl = from_dlpack(w, assumed_align=16).mark_layout_dynamic(leading_dim=2) cu_pl = from_dlpack(cu_seqlens, assumed_align=8).mark_layout_dynamic() - ws_pl = from_dlpack(tensormap_workspace, assumed_align=128).mark_layout_dynamic() + workspace_pl = from_dlpack(tensormap_workspace, assumed_align=128).mark_layout_dynamic() state_checkpoints_pl = None if state_checkpoints_for_descs is not None: state_checkpoints_pl = from_dlpack(state_checkpoints_for_descs, assumed_align=16).mark_layout_dynamic(leading_dim=3) @@ -2759,16 +2631,16 @@ def chunk_gdn2_recompute_sm100( work_items_pl = from_dlpack(work_items, assumed_align=16) work_items_pl.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1), divisibility=1) work_count_pl = from_dlpack(work_count, assumed_align=4).mark_layout_dynamic() - sched_all_pl = None - if sched_all is not None: - sched_all_pl = from_dlpack(sched_all, assumed_align=4).mark_layout_dynamic() + scheduler_all_pl = None + if scheduler_all is not None: + scheduler_all_pl = from_dlpack(scheduler_all, assumed_align=4).mark_layout_dynamic() cache["prologue"] = cute.compile( prologue, io_dtype, CFG.B_T, order_in_prologue, order_gen, - sched_all is not None, + scheduler_all is not None, k_pl, v_pl, gate_pl, @@ -2779,8 +2651,8 @@ def chunk_gdn2_recompute_sm100( staging_pl, work_count_pl, work_items_pl, - sched_all_pl, - ws_pl, + scheduler_all_pl, + workspace_pl, cutlass.Int32(checkpoint_every_n_tokens), cu_stream, options="--enable-tvm-ffi", @@ -2796,7 +2668,7 @@ def chunk_gdn2_recompute_sm100( work_item_scratch if not order_gen else None, work_count, work_items, - sched_all, + scheduler_all, tensormap_workspace, checkpoint_every_n_tokens, cu_stream, @@ -2814,7 +2686,7 @@ def chunk_gdn2_recompute_sm100( output_state if store_final_state else None, work_items, work_count, - sched_ctr, + scheduler_counter, tensormap_workspace, checkpoint_every_n_tokens, cu_stream, @@ -2837,8 +2709,8 @@ def run_recompute( output_state_checkpoints, work_items, work_count, - sched_ctr, - sched_all, + scheduler_counter, + scheduler_all, work_item_scratch, tensormap_workspace, checkpoint_every_n_tokens, @@ -2859,7 +2731,7 @@ def run_recompute( work_item_scratch, work_count, work_items, - sched_all, + scheduler_all, tensormap_workspace, checkpoint_every_n_tokens, cu_stream, @@ -2877,7 +2749,7 @@ def run_recompute( output_state, work_items, work_count, - sched_ctr, + scheduler_counter, tensormap_workspace, checkpoint_every_n_tokens, cu_stream, diff --git a/python/cudnn/linear_attention/frost/kernel/gdn_bprop_config.py b/python/cudnn/linear_attention/frost/kernel/gdn_bprop_config.py index 9536c5a85..6854084a3 100644 --- a/python/cudnn/linear_attention/frost/kernel/gdn_bprop_config.py +++ b/python/cudnn/linear_attention/frost/kernel/gdn_bprop_config.py @@ -35,23 +35,25 @@ class Cfg: # --- TMA descriptor pool --- - # --- warp assignments (12 warps total) --- + # --- warp assignments (16 warps total) --- COMPUTE_GROUP_0_WARP_IDS: Tuple[int, ...] = (0, 1, 2, 3) # T-pairwise / kk_epi / qk_epi / inverse / parts COMPUTE_GROUP_1_WARP_IDS: Tuple[int, ...] = (4, 5, 6, 7) # dH prep / dV-dK-dQ epilogues / dq dot - MMA_WARP_ID: int = 8 - TMA_QKV_WARP_ID: int = 9 - LOAD_GATE_BETA_WARP_ID: int = 10 - EPILOGUE_WARP_ID: int = 11 + COMPUTE_GROUP_2_WARP_IDS: Tuple[int, ...] = (8, 9, 10, 11) # dK inter rescale / attn read / dGate K parts / dK fold + TCGEN05_MMA_WARP_ID: int = 12 + TMA_QKV_WARP_ID: int = 13 + LOAD_GATE_BETA_WARP_ID: int = 14 + EPILOGUE_WARP_ID: int = 15 # --- register split --- - NUM_REGS_COMPUTE_GROUP_0: int = 224 - NUM_REGS_COMPUTE_GROUP_1: int = 248 + NUM_REGS_COMPUTE_GROUP_0: int = 208 + NUM_REGS_COMPUTE_GROUP_1: int = 144 + NUM_REGS_COMPUTE_GROUP_2: int = 128 NUM_REGS_OTHER: int = 32 THREADS_PER_WARP: int = 32 CLUSTER_SHAPE_MNK: Tuple[int, int, int] = (1, 1, 1) - SMEM_SCHED_STAGES: int = 2 + SMEM_SCHEDULER_STAGES: int = 2 # --- SMEM stage counts --- SMEM_Q_STAGES: int = 1 diff --git a/python/cudnn/linear_attention/frost/kernel/gdn_bprop_f16.py b/python/cudnn/linear_attention/frost/kernel/gdn_bprop_f16.py index 8f34c5461..23e62a6ab 100644 --- a/python/cudnn/linear_attention/frost/kernel/gdn_bprop_f16.py +++ b/python/cudnn/linear_attention/frost/kernel/gdn_bprop_f16.py @@ -73,7 +73,7 @@ KK (strict-masked M_kk) 8192 1 <-- KK epi's only store; inverse input + dGate/dBeta A staging / sDa 8192 1 <-- ALIAS: A then the masked dA dM staging (sDm) 8192 1 <-- Step 8 -> dK dM-terms - dstate_entry (sDstate) 32768 1 <-- f16 restage, dK-inter's A + dstate_entry (sDstate) 32768 1 <-- f16 restage, dK-inter's A dQ store staging 16384 1 dK store staging 16384 1 dV store staging 16384 1 @@ -84,7 +84,7 @@ cols 128-192 : dV/dK accumulator (fp32) <-- one slot, five sequential per-chunk productions (dV inter -> dU intra -> dK inter -> dK attn -> dK dM-terms), each with its own mbar pair - cols 192-256 : dstate input (f16 packed) <-- prefill: state_inp + cols 192-256 : dstate input (f16 packed) <-- prefill: state_input cols 256-384 : shared accumulators x2 <-- KK / A / k_state / U / dY / dA / dM core cols 384-448 : shared inputs x2 (f16 packed) <-- dO' / dU / dY'; the dK @@ -117,8 +117,7 @@ from cutlass.cute.runtime import from_dlpack from ..common.thd import emit_checkpoint_seq_descs, emit_seq_descs, TENSOR_MAP_QWORDS -from ..common.split_k import ORDER_CAPACITY, ORDER_ELEMS, ORDER_THREADS, decode_work_item, order_body -from ..common.elementwise import softplus +from ..common.split_k import ORDER_CAPACITY, ORDER_ELEMENTS, ORDER_THREADS, decode_work_item, order_body from ..common.host import get_dtype from cudnn.frost.buffers import data_ptr from cudnn.frost.device import current_device, multiprocessor_count @@ -132,8 +131,8 @@ ) from cudnn.frost.tile_dsl.handles import MmaDesc, SmemTile, tma_slice_runtime_desc from cudnn.frost.tile_dsl.mma import mma_ss, mma_step_k8, mma_ts_step, mma_step -from cudnn.frost.tile_dsl.pointwise import fp32_to_fp16, f16x2_to_f32, fmul2, fadd2, ffma2, opaque_f32_zero, sub_f16x2 -from cudnn.frost.tile_dsl.swizzle import swizzle_lin_128b, swizzle_xor_128b +from cudnn.frost.tile_dsl.pointwise import f16x2_to_f32, fadd2, ffma2, fmul2, fp32_to_fp16, opaque_f32_zero, sigmoid, softplus, softplus2, sub_f16x2 +from cudnn.frost.tile_dsl.swizzle import swizzle_xor_128b from cudnn.frost.tile_dsl.tma import ( tma_load_tile, tma_store_tile, @@ -145,25 +144,14 @@ class GdnBwdBars(NamedTuple): - """GDN bprop pipeline mbarrier inventory. - - Every pipeline is a ``_ready``/``_done`` MBarrier pair over one ring: a - slot is acquired for filling by waiting ``_done`` and committed by - arriving ``_ready``; the reading side waits ``_ready`` and releases the - slot by arriving ``_done``. - - Operand buffers read by both the MMA warp and a compute group carry a - SPLIT done pair (``_mma_done`` MMA_COMMIT + ``_cg?_done`` THREAD) and the - TMA warp waits BOTH: a plain arrive from the MMA warp fires at MMA issue, - not completion, which would release the buffer for reload mid-GEMM. - """ + """Every inter-warp handoff as an ``MBarrier`` over its ring.""" mb_q_ready: MBarrier mb_q_mma_done: MBarrier mb_q_cg1_done: MBarrier mb_k_ready: MBarrier mb_k_mma_done: MBarrier - mb_k_cg0_done: MBarrier + mb_k_cg2_done: MBarrier mb_v_ready: MBarrier mb_v_mma_done: MBarrier mb_do_ready: MBarrier @@ -202,11 +190,11 @@ class GdnBwdBars(NamedTuple): mb_dm_acc_done: MBarrier mb_dk_state_path_acc_ready: MBarrier - mb_dstate_inp_ready: MBarrier - mb_dstate_inp_done: MBarrier - mb_do_prime_inp_ready: MBarrier - mb_du_inp_ready: MBarrier - mb_dyp_inp_ready: MBarrier + mb_dstate_input_ready: MBarrier + mb_dstate_input_done: MBarrier + mb_do_prime_input_ready: MBarrier + mb_du_input_ready: MBarrier + mb_dyp_input_ready: MBarrier mb_y_ready: MBarrier mb_t_inv_ready: MBarrier @@ -219,6 +207,8 @@ class GdnBwdBars(NamedTuple): mb_dbeta_cg1_ready: MBarrier mb_dgate_cg1_ready: MBarrier + mb_dgate_cg0_ready: MBarrier + mb_cg2_chunk_ready: MBarrier mb_dq_tmastg_ready: MBarrier mb_dq_tmastg_done: MBarrier @@ -229,20 +219,21 @@ class GdnBwdBars(NamedTuple): mb_sdv_done: MBarrier mb_tmem_done: MBarrier - mb_sched_ready: MBarrier - mb_sched_done: MBarrier + mb_scheduler_ready: MBarrier + mb_scheduler_done: MBarrier def make_gdn_bars(cfg) -> GdnBwdBars: - """GdnBwdBars factory. MUST be called from inside ``kernel`` (allocates SMEM; - the mbar rings sit ahead of the gate scalar arrays and data buffers).""" + """GdnBwdBars factory.""" ONE_LANE = 1 - MMA_ARRIVERS = len([cfg.mma_warp_id]) + MMA_ARRIVERS = len([cfg.tcgen05_mma_warp_id]) GATE_WARP = cfg.threads_per_warp * len([cfg.load_gate_beta_warp_id]) EPI_WARP = cfg.threads_per_warp * len([cfg.epilogue_warp_id]) CG0_THREADS = cfg.threads_per_warp * len(cfg.compute_group_0_warp_ids) CG1_THREADS = cfg.threads_per_warp * len(cfg.compute_group_1_warp_ids) + CG2_THREADS = cfg.threads_per_warp * len(cfg.compute_group_2_warp_ids) CG0_PLUS_CG1 = CG0_THREADS + CG1_THREADS + ALL_COMPUTE = CG0_THREADS + CG1_THREADS + CG2_THREADS def alloc(n): return cutlass.Array(cutlass.Int64, n, space=cutlass.AddressSpace.smem, alignment=16) @@ -253,7 +244,7 @@ def alloc(n): mb_q_cg1_done=MBarrier(alloc(cfg.smem_q_stages), stages=cfg.smem_q_stages, init_count=CG1_THREADS, producer=Producer.THREAD), mb_k_ready=MBarrier(alloc(cfg.smem_k_stages), stages=cfg.smem_k_stages, init_count=ONE_LANE, producer=Producer.TMA_LOAD), mb_k_mma_done=MBarrier(alloc(cfg.smem_k_stages), stages=cfg.smem_k_stages, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), - mb_k_cg0_done=MBarrier(alloc(cfg.smem_k_stages), stages=cfg.smem_k_stages, init_count=CG0_THREADS, producer=Producer.THREAD), + mb_k_cg2_done=MBarrier(alloc(cfg.smem_k_stages), stages=cfg.smem_k_stages, init_count=CG2_THREADS, producer=Producer.THREAD), mb_v_ready=MBarrier(alloc(cfg.smem_v_stages), stages=cfg.smem_v_stages, init_count=ONE_LANE, producer=Producer.TMA_LOAD), mb_v_mma_done=MBarrier(alloc(cfg.smem_v_stages), stages=cfg.smem_v_stages, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), mb_do_ready=MBarrier(alloc(cfg.smem_do_stages), stages=cfg.smem_do_stages, init_count=ONE_LANE, producer=Producer.TMA_LOAD), @@ -261,31 +252,28 @@ def alloc(n): mb_state_ready=MBarrier(alloc(cfg.smem_state_stages), stages=cfg.smem_state_stages, init_count=ONE_LANE, producer=Producer.TMA_LOAD), mb_state_mma_done=MBarrier(alloc(cfg.smem_state_stages), stages=cfg.smem_state_stages, init_count=1, producer=Producer.MMA_COMMIT), mb_gate_ready=MBarrier(alloc(cfg.smem_gate_stages), stages=cfg.smem_gate_stages, init_count=GATE_WARP, producer=Producer.THREAD), - mb_gate_done=MBarrier(alloc(cfg.smem_gate_stages), stages=cfg.smem_gate_stages, init_count=CG0_PLUS_CG1, producer=Producer.THREAD), + mb_gate_done=MBarrier(alloc(cfg.smem_gate_stages), stages=cfg.smem_gate_stages, init_count=ALL_COMPUTE, producer=Producer.THREAD), mb_beta_ready=MBarrier(alloc(cfg.smem_beta_stages), stages=cfg.smem_beta_stages, init_count=GATE_WARP, producer=Producer.THREAD), mb_beta_done=MBarrier(alloc(cfg.smem_beta_stages), stages=cfg.smem_beta_stages, init_count=CG0_PLUS_CG1, producer=Producer.THREAD), mb_dstate_acc_ready=MBarrier( alloc(cfg.tmem_dstate_acc_stages), stages=cfg.tmem_dstate_acc_stages, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT ), mb_dstate_scale_acc_done=MBarrier( - alloc(cfg.tmem_dstate_acc_stages), stages=cfg.tmem_dstate_acc_stages, init_count=CG1_THREADS, producer=Producer.THREAD + alloc(cfg.tmem_dstate_acc_stages), stages=cfg.tmem_dstate_acc_stages, init_count=CG2_THREADS, producer=Producer.THREAD ), - # five sequential per-chunk productions share the dV/dK TMEM slot mb_du_scale_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), mb_du_scale_acc_done=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), mb_du_total_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), mb_dk_scale_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), - mb_dk_scale_acc_done=MBarrier(alloc(1), stages=1, init_count=CG0_THREADS, producer=Producer.THREAD), + mb_dk_scale_acc_done=MBarrier(alloc(1), stages=1, init_count=CG2_THREADS, producer=Producer.THREAD), mb_dk_attn_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), - mb_dk_attn_acc_done=MBarrier(alloc(1), stages=1, init_count=CG0_THREADS, producer=Producer.THREAD), + mb_dk_attn_acc_done=MBarrier(alloc(1), stages=1, init_count=CG2_THREADS, producer=Producer.THREAD), mb_dk_total_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), - mb_dk_total_acc_done=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_dk_total_acc_done=MBarrier(alloc(1), stages=1, init_count=CG2_THREADS, producer=Producer.THREAD), mb_dq_acc_scale_ready=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), - mb_dq_acc_scale_done=MBarrier(alloc(1), stages=1, init_count=CG0_THREADS, producer=Producer.THREAD), + mb_dq_acc_scale_done=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), mb_dq_acc_total_ready=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), mb_dq_acc_total_done=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), - # shared accumulators at STATIC columns: group A holds - # KK -> k_state -> dY -> dM core, group B holds A -> U -> dA. mb_kk_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), mb_kk_acc_done=MBarrier(alloc(1), stages=1, init_count=CG0_THREADS, producer=Producer.THREAD), mb_a_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), @@ -296,94 +284,84 @@ def alloc(n): mb_dm_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), mb_dm_acc_done=MBarrier(alloc(1), stages=1, init_count=CG0_THREADS, producer=Producer.THREAD), mb_dk_state_path_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), - mb_dstate_inp_ready=MBarrier(alloc(cfg.tmem_dstate_inp_stages), stages=cfg.tmem_dstate_inp_stages, init_count=CG1_THREADS, producer=Producer.THREAD), - mb_dstate_inp_done=MBarrier( - alloc(cfg.tmem_dstate_inp_stages), stages=cfg.tmem_dstate_inp_stages, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT + mb_dstate_input_ready=MBarrier( + alloc(cfg.tmem_dstate_input_stages), stages=cfg.tmem_dstate_input_stages, init_count=CG2_THREADS, producer=Producer.THREAD + ), + mb_dstate_input_done=MBarrier( + alloc(cfg.tmem_dstate_input_stages), stages=cfg.tmem_dstate_input_stages, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT ), - # f16 input restages at STATIC columns: dO' alone; dU and dY' overlap - mb_do_prime_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), - mb_du_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), - mb_dyp_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_do_prime_input_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_du_input_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_dyp_input_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), mb_y_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), mb_t_inv_ready=MBarrier(alloc(cfg.smem_t_inv_stages), stages=cfg.smem_t_inv_stages, init_count=CG0_THREADS, producer=Producer.THREAD), mb_a_ready=MBarrier(alloc(cfg.smem_a_stages), stages=cfg.smem_a_stages, init_count=CG0_THREADS, producer=Producer.THREAD), mb_a_done=MBarrier(alloc(cfg.smem_a_stages), stages=cfg.smem_a_stages, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), mb_u_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), - mb_dstate_smem_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_dstate_smem_ready=MBarrier(alloc(1), stages=1, init_count=CG2_THREADS, producer=Producer.THREAD), mb_state_dot_dstate_done=MBarrier(alloc(1), stages=1, init_count=CG0_THREADS, producer=Producer.THREAD), mb_da_ready=MBarrier(alloc(1), stages=1, init_count=CG0_THREADS, producer=Producer.THREAD), mb_dbeta_cg1_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), mb_dgate_cg1_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_dgate_cg0_ready=MBarrier(alloc(1), stages=1, init_count=CG0_THREADS, producer=Producer.THREAD), + mb_cg2_chunk_ready=MBarrier(alloc(1), stages=1, init_count=CG2_THREADS, producer=Producer.THREAD), mb_dq_tmastg_ready=MBarrier(alloc(cfg.smem_dq_stages), stages=cfg.smem_dq_stages, init_count=CG1_THREADS, producer=Producer.THREAD), mb_dq_tmastg_done=MBarrier(alloc(cfg.smem_dq_stages), stages=cfg.smem_dq_stages, init_count=EPI_WARP, producer=Producer.THREAD), - mb_dk_tmastg_ready=MBarrier(alloc(cfg.smem_dk_stages), stages=cfg.smem_dk_stages, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_dk_tmastg_ready=MBarrier(alloc(cfg.smem_dk_stages), stages=cfg.smem_dk_stages, init_count=CG2_THREADS, producer=Producer.THREAD), mb_dk_tmastg_done=MBarrier(alloc(cfg.smem_dk_stages), stages=cfg.smem_dk_stages, init_count=EPI_WARP, producer=Producer.THREAD), mb_dv_tmastg_ready=MBarrier(alloc(cfg.smem_dv_stages), stages=cfg.smem_dv_stages, init_count=CG1_THREADS, producer=Producer.THREAD), mb_dv_tmastg_done=MBarrier(alloc(cfg.smem_dv_stages), stages=cfg.smem_dv_stages, init_count=EPI_WARP, producer=Producer.THREAD), mb_sdv_done=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), - mb_tmem_done=MBarrier(alloc(1), stages=1, init_count=CG0_THREADS + CG1_THREADS, producer=Producer.THREAD), - mb_sched_ready=MBarrier(alloc(cfg.sched_stages), stages=cfg.sched_stages, init_count=ONE_LANE, producer=Producer.THREAD), - mb_sched_done=MBarrier(alloc(cfg.sched_stages), stages=cfg.sched_stages, init_count=11, producer=Producer.THREAD), + mb_tmem_done=MBarrier(alloc(1), stages=1, init_count=ALL_COMPUTE, producer=Producer.THREAD), + mb_scheduler_ready=MBarrier(alloc(cfg.scheduler_stages), stages=cfg.scheduler_stages, init_count=ONE_LANE, producer=Producer.THREAD), + mb_scheduler_done=MBarrier(alloc(cfg.scheduler_stages), stages=cfg.scheduler_stages, init_count=15, producer=Producer.THREAD), ) -# ---- Dynamic tile scheduler ------------------------------------------------------ +# ---- Dynamic tile scheduler ---------------------------------------------------------- @cute.jit -def sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ctas): +def scheduler_publish_next(cfg, bars, sScheduler, mScheduler, scheduler_state, tile_idx, num_ctas, tail_base, tail_row, elect_one): """TMA-LDG-warp side: pull the next tile off the global ticket, publish it.""" - if cutlass.const_expr(cfg.dyn_sched): - bars.mb_sched_done[sched_state.idx].wait(sched_state.phase) - if nvvm.elect_sync(): - fetched = cutlass.Int32(nvvm.atomicrmw("add", mSched.iterator, cutlass.Int32(1), mem_order="relaxed", syncscope="gpu")) - sSched[sched_state.idx] = num_ctas + fetched + if cutlass.const_expr(cfg.dynamic_scheduling): + sentinel = cutlass.Int32(1 << 28) + pinned = tail_row if tile_idx < tail_base else sentinel + bars.mb_scheduler_done[scheduler_state.idx].wait(scheduler_state.phase) + if elect_one: + fetched = cutlass.Int32(nvvm.atomicrmw("add", mScheduler.iterator, cutlass.Int32(1), mem_order="relaxed", syncscope="gpu")) + granted = num_ctas + fetched + sScheduler[scheduler_state.idx] = granted if granted < tail_base else pinned nvvm.bar_warp_sync(cute.arch.FULL_MASK) - next_tile = sSched[sched_state.idx] - if nvvm.elect_sync(): - bars.mb_sched_ready[sched_state.idx].arrive() - return next_tile, advance(sched_state, cfg.sched_stages) - return tile_idx + num_ctas, sched_state + next_tile = sScheduler[scheduler_state.idx] + if elect_one: + bars.mb_scheduler_ready[scheduler_state.idx].arrive() + return next_tile, advance(scheduler_state, cfg.scheduler_stages) + return tile_idx + num_ctas, scheduler_state @cute.jit -def sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas): +def scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one): """Consumer side: read the TMA-LDG warp's published next tile.""" - if cutlass.const_expr(cfg.dyn_sched): - bars.mb_sched_ready[sched_state.idx].wait(sched_state.phase) - next_tile = sSched[sched_state.idx] - if nvvm.elect_sync(): - bars.mb_sched_done[sched_state.idx].arrive() - return next_tile, advance(sched_state, cfg.sched_stages) - return tile_idx + num_ctas, sched_state + if cutlass.const_expr(cfg.dynamic_scheduling): + bars.mb_scheduler_ready[scheduler_state.idx].wait(scheduler_state.phase) + next_tile = sScheduler[scheduler_state.idx] + if elect_one: + bars.mb_scheduler_done[scheduler_state.idx].arrive() + return next_tile, advance(scheduler_state, cfg.scheduler_stages) + return tile_idx + num_ctas, scheduler_state @cute.jit -def invert_diagonal_NxN(cfg, in_base, out_base, d, tidx, N: int = 8): +def invert_diagonal_NxN(cfg, in_base, out_base, d_idx, tidx, N: int = 8): """Gauss-Jordan inversion of one diagonal NxN block, ``in_base`` -> ``out_base`` (f16 SMEM).""" tidx_in_group = tidx % N BT = cfg.b_t - row_lin_base = (d * N + tidx_in_group) * BT + d * N - row_phys = swizzle_lin_128b(row_lin_base, row_stride_log2=6) - row_ptr_in = ( - cute.make_ptr( - cfg.io_dtype, - in_base, - mem_space=cute.AddressSpace.smem, - assumed_align=cfg.buffer_align_bytes, - ) - + row_phys - ) - row_ptr = ( - cute.make_ptr( - cfg.io_dtype, - out_base, - mem_space=cute.AddressSpace.smem, - assumed_align=cfg.buffer_align_bytes, - ) - + row_phys - ) + row_coord = d_idx * N + tidx_in_group + row_off = row_coord * BT + swizzle_xor_128b(row_coord, d_idx * N) + row_ptr_in = in_base + row_off + row_ptr = out_base + row_off row = [(row_ptr_in + j).load().to(cutlass.Float32) for j in range(N)] for i in cutlass.range_constexpr(N): @@ -400,211 +378,139 @@ def invert_diagonal_NxN(cfg, in_base, out_base, d, tidx, N: int = 8): @cute.jit -def warp_reduce_scatter_frag_16_elems(vals, lane_id): +def warp_reduce_scatter_sum_16_elements(vals, lane_idx): """Reduce-scatter 16 fragment token-partials (tcol = (lane%4)*2 + (j//2)*8 + (j%2)); lane L returns the sums of tokens (L//4)*8 + (L%4)*2 and +1.""" - cur = vals + level = vals for k in cutlass.range_constexpr(3): off = cutlass.const_expr(4 << k) - hi_lane = (lane_id // off) % 2 == 1 - nxt = [] - for i in cutlass.range_constexpr(len(cur)): + hi_lane = (lane_idx // off) % 2 == 1 + next_level = [] + for i in cutlass.range_constexpr(len(level)): if cutlass.const_expr((i & 2) == 0): - lo = cur[i] - hi = cur[i + 2] + lo = level[i] + hi = level[i + 2] send = lo if hi_lane else hi recv = nvvm.shfl_sync(0xFFFFFFFF, send, off, 31, kind=nvvm.Shfl.BFLY) keep = hi if hi_lane else lo - nxt.append(keep + recv) - cur = nxt - return cur[0], cur[1] + next_level.append(keep + recv) + level = next_level + return level[0], level[1] @cute.jit -def blockwise_diagonal_8x8_to_16x16(cfg, base_int, raw_base, d0, lane_id): - """Off-diagonal correction 8x8 -> 16x16 (C <- -D^{-1} C A^{-1}); raw C from ``raw_base``, writes on ``base_int``.""" - bpe = cfg.io_dtype.width // 8 - ldsm_x1_off = (lane_id % 8) * 64 - d = nvvm.ldmatrix( - cutlass.inttoptr( - base_int + swizzle_lin_128b((d0 + 8) * 64 + d0 + 8 + ldsm_x1_off, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16 - ), - 1, - nvvm.MMALayout.ROW, - ) - c = nvvm.ldmatrix( - cutlass.inttoptr(raw_base + swizzle_lin_128b((d0 + 8) * 64 + d0 + ldsm_x1_off, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16), - 1, - nvvm.MMALayout.COL, - ) - - # ---- T = -(D^{-1} @ C) ------------------------------------------------------- +def blockwise_diagonal_8x8_to_16x16(cfg, base, raw_base, d_idx, lane_idx): + """Off-diagonal correction 8x8 -> 16x16 (C <- -D^{-1} C A^{-1}); raw C from ``raw_base``, writes on ``base``.""" + BT = cfg.b_t + row_lo = d_idx + lane_idx % 8 + row_hi = row_lo + 8 + off_d_inv = row_hi * BT + swizzle_xor_128b(row_hi, d_idx + 8) + off_c = row_hi * BT + swizzle_xor_128b(row_hi, d_idx) + off_a_inv = row_lo * BT + swizzle_xor_128b(row_lo, d_idx) + d_inv_frag = nvvm.ldmatrix(base + off_d_inv, 1, nvvm.MMALayout.ROW) + c_frag = nvvm.ldmatrix(raw_base + off_c, 1, nvvm.MMALayout.COL) + + # ---- T = -(D^-1 @ C) ------------------------------------------------------------- c_regs = cutlass.Array(cutlass.Float32, 4, alignment=16, space=cutlass.AddressSpace.rmem) for i in cutlass.range_constexpr(4): c_regs[i] = cutlass.Float32(0.0) - mma_step_k8(c_regs, [d, d], [c], k_step=0, M=16, N=8, ab_dtype=cfg.io_dtype) + mma_step_k8(c_regs, [d_inv_frag, d_inv_frag], [c_frag], k_step=0, M=16, N=8, ab_dtype=cfg.io_dtype) for i in cutlass.range_constexpr(4): c_regs[i] = -c_regs[i] a_pack = [fp32_to_fp16(c_regs[2 * j], c_regs[2 * j + 1], dtype=cfg.io_dtype) for j in range(2)] - # ---- C = T @ A^{-1} ---------------------------------------------------------- - ai = nvvm.ldmatrix( - cutlass.inttoptr(base_int + swizzle_lin_128b(d0 * 64 + d0 + ldsm_x1_off, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16), - 1, - nvvm.MMALayout.COL, - ) + # ---- C = T @ A^-1 ---------------------------------------------------------------- + a_inv_frag = nvvm.ldmatrix(base + off_a_inv, 1, nvvm.MMALayout.COL) o_regs = cutlass.Array(cutlass.Float32, 4, alignment=16, space=cutlass.AddressSpace.rmem) for i in cutlass.range_constexpr(4): o_regs[i] = cutlass.Float32(0.0) - mma_step_k8(o_regs, a_pack, [ai], k_step=0, M=16, N=8, ab_dtype=cfg.io_dtype) + mma_step_k8(o_regs, a_pack, [a_inv_frag], k_step=0, M=16, N=8, ab_dtype=cfg.io_dtype) o_pack = fp32_to_fp16(o_regs[0], o_regs[1], dtype=cfg.io_dtype) - # ---- store corrected C ------------------------------------------------------- - nvvm.stmatrix( - cutlass.inttoptr(base_int + swizzle_lin_128b((d0 + 8) * 64 + d0 + ldsm_x1_off, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16), - o_pack, - nvvm.MMALayout.ROW, - ) + # ---- store corrected C ----------------------------------------------------------- + nvvm.stmatrix(base + off_c, o_pack, nvvm.MMALayout.ROW) @cute.jit -def blockwise_diagonal_16x16_to_32x32(cfg, base_int, raw_base, d0, lane_id): +def blockwise_diagonal_16x16_to_32x32(cfg, base, raw_base, d_idx, lane_idx): """Off-diagonal correction 16x16 -> 32x32 (raw C from ``raw_base``).""" - bpe = cfg.io_dtype.width // 8 - ldsm_x4_off = (lane_id % 16) * 64 + (lane_id // 16) * 8 - d = list( - nvvm.ldmatrix( - cutlass.inttoptr( - base_int + swizzle_lin_128b((d0 + 16) * 64 + d0 + 16 + ldsm_x4_off, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16 - ), - 4, - nvvm.MMALayout.ROW, - ) - ) - c = list( - nvvm.ldmatrix( - cutlass.inttoptr( - raw_base + swizzle_lin_128b((d0 + 16) * 64 + d0 + ldsm_x4_off, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16 - ), - 4, - nvvm.MMALayout.COL, - ) - ) - - # ---- T = -(D^{-1} @ C) ------------------------------------------------------- + BT = cfg.b_t + lane_row = lane_idx % 16 + lane_col = (lane_idx // 16) * 8 + row_lo = d_idx + lane_row + row_hi = row_lo + 16 + off_d_inv = row_hi * BT + swizzle_xor_128b(row_hi, d_idx + 16 + lane_col) + off_c = row_hi * BT + swizzle_xor_128b(row_hi, d_idx + lane_col) + off_a_inv = row_lo * BT + swizzle_xor_128b(row_lo, d_idx + lane_col) + d_inv_frags = list(nvvm.ldmatrix(base + off_d_inv, 4, nvvm.MMALayout.ROW)) + c_frags = list(nvvm.ldmatrix(raw_base + off_c, 4, nvvm.MMALayout.COL)) + + # ---- T = -(D^-1 @ C) ------------------------------------------------------------- c_regs = cutlass.Array(cutlass.Float32, 8, alignment=16, space=cutlass.AddressSpace.rmem) for i in cutlass.range_constexpr(8): c_regs[i] = cutlass.Float32(0.0) - mma_step(c_regs, d, c, k_step=0, M=16, N=16, ab_dtype=cfg.io_dtype) + mma_step(c_regs, d_inv_frags, c_frags, k_step=0, M=16, N=16, ab_dtype=cfg.io_dtype) for i in cutlass.range_constexpr(8): c_regs[i] = -c_regs[i] a_pack = [fp32_to_fp16(c_regs[2 * j], c_regs[2 * j + 1], dtype=cfg.io_dtype) for j in range(4)] - # ---- C = T @ A^{-1} ---------------------------------------------------------- - ai = list( - nvvm.ldmatrix( - cutlass.inttoptr(base_int + swizzle_lin_128b(d0 * 64 + d0 + ldsm_x4_off, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16), - 4, - nvvm.MMALayout.COL, - ) - ) + # ---- C = T @ A^-1 ---------------------------------------------------------------- + a_inv_frags = list(nvvm.ldmatrix(base + off_a_inv, 4, nvvm.MMALayout.COL)) o_regs = cutlass.Array(cutlass.Float32, 8, alignment=16, space=cutlass.AddressSpace.rmem) for i in cutlass.range_constexpr(8): o_regs[i] = cutlass.Float32(0.0) - mma_step(o_regs, a_pack, ai, k_step=0, M=16, N=16, ab_dtype=cfg.io_dtype) + mma_step(o_regs, a_pack, a_inv_frags, k_step=0, M=16, N=16, ab_dtype=cfg.io_dtype) o_pack = [fp32_to_fp16(o_regs[2 * j], o_regs[2 * j + 1], dtype=cfg.io_dtype) for j in range(4)] - # ---- store corrected C ------------------------------------------------------- - nvvm.stmatrix( - cutlass.inttoptr(base_int + swizzle_lin_128b((d0 + 16) * 64 + d0 + ldsm_x4_off, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16), - o_pack, - nvvm.MMALayout.ROW, - ) + # ---- store corrected C ----------------------------------------------------------- + nvvm.stmatrix(base + off_c, o_pack, nvvm.MMALayout.ROW) @cute.jit -def blockwise_diagonal_32x32_to_64x64(cfg, base_int, raw_base, warp_id, lane_id): +def blockwise_diagonal_32x32_to_64x64(cfg, base, raw_base, warp_id, lane_idx): """Off-diagonal correction 32x32 -> 64x64 (2 warps, one 16-row M-band each; raw C from ``raw_base``).""" band = warp_id % 2 - bpe = cfg.io_dtype.width // 8 - ldsm_x4_off = (lane_id % 16) * 64 + (lane_id // 16) * 8 - a_frags = [] + BT = cfg.b_t + lane_row = lane_idx % 16 + lane_col = (lane_idx // 16) * 8 + row_d_inv = 32 + band * 16 + lane_row + d_inv_frags = [] for vs in cutlass.range_constexpr(2): - a_frags += list( - nvvm.ldmatrix( - cutlass.inttoptr( - base_int + swizzle_lin_128b((32 + band * 16) * 64 + 32 + vs * 16 + ldsm_x4_off, row_stride_log2=6) * bpe, - cutlass.AddressSpace.smem, - cutlass.BFloat16, - ), - 4, - nvvm.MMALayout.ROW, - ) - ) - b_frags = [] + d_inv_frags += list(nvvm.ldmatrix(base + row_d_inv * BT + swizzle_xor_128b(row_d_inv, 32 + vs * 16 + lane_col), 4, nvvm.MMALayout.ROW)) + c_frags = [] for vs in cutlass.range_constexpr(4): - b_frags += list( - nvvm.ldmatrix( - cutlass.inttoptr( - raw_base + swizzle_lin_128b((32 + (vs // 2) * 16) * 64 + (vs % 2) * 16 + ldsm_x4_off, row_stride_log2=6) * bpe, - cutlass.AddressSpace.smem, - cutlass.BFloat16, - ), - 4, - nvvm.MMALayout.COL, - ) - ) + row_c = 32 + (vs // 2) * 16 + lane_row + c_frags += list(nvvm.ldmatrix(raw_base + row_c * BT + swizzle_xor_128b(row_c, (vs % 2) * 16 + lane_col), 4, nvvm.MMALayout.COL)) - # ---- T = -(D^{-1} @ C) ------------------------------------------------------- + # ---- T = -(D^-1 @ C) ------------------------------------------------------------- c_regs = cutlass.Array(cutlass.Float32, 16, alignment=16, space=cutlass.AddressSpace.rmem) for i in cutlass.range_constexpr(16): c_regs[i] = cutlass.Float32(0.0) for ks in cutlass.range_constexpr(2): - mma_step(c_regs, a_frags, b_frags[ks * 8 : ks * 8 + 8], k_step=ks, M=16, N=32, ab_dtype=cfg.io_dtype) + mma_step(c_regs, d_inv_frags, c_frags[ks * 8 : ks * 8 + 8], k_step=ks, M=16, N=32, ab_dtype=cfg.io_dtype) for i in cutlass.range_constexpr(16): c_regs[i] = -c_regs[i] a_pack = [fp32_to_fp16(c_regs[2 * j], c_regs[2 * j + 1], dtype=cfg.io_dtype) for j in range(8)] - # ---- C = T @ A^{-1} ---------------------------------------------------------- - ai_frags = [] + # ---- C = T @ A^-1 ---------------------------------------------------------------- + a_inv_frags = [] for vs in cutlass.range_constexpr(4): - ai_frags += list( - nvvm.ldmatrix( - cutlass.inttoptr( - base_int + swizzle_lin_128b(((vs // 2) * 16) * 64 + (vs % 2) * 16 + ldsm_x4_off, row_stride_log2=6) * bpe, - cutlass.AddressSpace.smem, - cutlass.BFloat16, - ), - 4, - nvvm.MMALayout.COL, - ) - ) + row_a_inv = (vs // 2) * 16 + lane_row + a_inv_frags += list(nvvm.ldmatrix(base + row_a_inv * BT + swizzle_xor_128b(row_a_inv, (vs % 2) * 16 + lane_col), 4, nvvm.MMALayout.COL)) o_regs = cutlass.Array(cutlass.Float32, 16, alignment=16, space=cutlass.AddressSpace.rmem) for i in cutlass.range_constexpr(16): o_regs[i] = cutlass.Float32(0.0) for ks in cutlass.range_constexpr(2): - mma_step(o_regs, a_pack, ai_frags[ks * 8 : ks * 8 + 8], k_step=ks, M=16, N=32, ab_dtype=cfg.io_dtype) + mma_step(o_regs, a_pack, a_inv_frags[ks * 8 : ks * 8 + 8], k_step=ks, M=16, N=32, ab_dtype=cfg.io_dtype) o_pack = [fp32_to_fp16(o_regs[2 * j], o_regs[2 * j + 1], dtype=cfg.io_dtype) for j in range(8)] - # ---- store corrected C ------------------------------------------------------- + # ---- store corrected C ----------------------------------------------------------- nvvm.barrier_cta_sync_aligned( cfg.inverse_inner_barrier_id, thread_count=cfg.inverse_inner_barrier_threads, ) - nvvm.stmatrix( - cutlass.inttoptr( - base_int + swizzle_lin_128b((32 + band * 16) * 64 + ldsm_x4_off, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16 - ), - o_pack[0:4], - nvvm.MMALayout.ROW, - ) - nvvm.stmatrix( - cutlass.inttoptr( - base_int + swizzle_lin_128b((32 + band * 16) * 64 + 16 + ldsm_x4_off, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16 - ), - o_pack[4:8], - nvvm.MMALayout.ROW, - ) + nvvm.stmatrix(base + row_d_inv * BT + swizzle_xor_128b(row_d_inv, lane_col), o_pack[0:4], nvvm.MMALayout.ROW) + nvvm.stmatrix(base + row_d_inv * BT + swizzle_xor_128b(row_d_inv, 16 + lane_col), o_pack[4:8], nvvm.MMALayout.ROW) @cute.jit @@ -621,23 +527,23 @@ def tmastg_warp( desc_dq_base, desc_dk_base, desc_dv_base, - sSched, + sScheduler, bars, ): """TMA-STG warp role (warp 11): persistent tile-scheduler loop + per-chunk dQ/dK/dV TMA bulk-stores from the SMEM staging buffers to global memory.""" - elect_one = nvvm.elect_sync() nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) + dq_index = PipelineState.start(phase=0) dk_index = PipelineState.start(phase=0) dv_index = PipelineState.start(phase=0) + scheduler_state = PipelineState.start(phase=0) - sched_state = PipelineState.start(phase=0) + elect_one = nvvm.elect_sync() tile_idx = cutlass.Int32(bidx) FIRST_STATE_CHUNK = 0 if cfg.use_initial_state else 1 - bpe = cfg.io_dtype.width // 8 - granule_elems = 128 // bpe + granule_elements = 128 // bpe sdQ_tma = SmemTile( base=sdQ_raw, elems_per_stage=(cfg.dq_cosize // cfg.smem_dq_stages), @@ -646,7 +552,7 @@ def tmastg_warp( stride_byte_offset=0, layout=0, tma_loads_per_tile=2, - tma_granu_elems=granule_elems, + tma_granu_elems=granule_elements, tma_subtile_stride_elems=4096, ) sdK_tma = SmemTile( @@ -657,7 +563,7 @@ def tmastg_warp( stride_byte_offset=0, layout=0, tma_loads_per_tile=2, - tma_granu_elems=granule_elems, + tma_granu_elems=granule_elements, tma_subtile_stride_elems=4096, ) sdV_tma = SmemTile( @@ -668,14 +574,16 @@ def tmastg_warp( stride_byte_offset=0, layout=0, tma_loads_per_tile=2, - tma_granu_elems=granule_elems, + tma_granu_elems=granule_elements, tma_subtile_stride_elems=4096, ) heads_out = cutlass.Int32(cfg.n_heads_out) desc_qwords = cutlass.Int32(TENSOR_MAP_QWORDS) while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) head_o = head_idx slot = batch_idx * desc_qwords @@ -687,15 +595,15 @@ def tmastg_warp( tma_tensormap_acquire(desc_dk_slot) tma_tensormap_acquire(desc_dv_slot) - for rev_idx in cutlass.range(cend - wstart): - chunk_idx = cend - 1 - rev_idx + for rev_idx in cutlass.range(compute_end - write_start): + chunk_idx = compute_end - 1 - rev_idx tok_coord = chunk_idx * cutlass.Int32(cfg.b_t) dv_idx = dv_index.idx bars.mb_dv_tmastg_ready[dv_idx].wait(dv_index.phase) dv_index = advance(dv_index, cfg.smem_dv_stages) dv_slice = tma_slice_runtime_desc(desc_dv_slot, cutlass.Int32(0), head_o, tok_coord) - if chunk_idx < wend: + if chunk_idx < write_end: tma_store_tile(sdV_tma[dv_idx], dv_slice, acquire=False) tma_store_commit() @@ -703,7 +611,7 @@ def tmastg_warp( bars.mb_dq_tmastg_ready[dq_idx].wait(dq_index.phase) dq_index = advance(dq_index, cfg.smem_dq_stages) dq_slice = tma_slice_runtime_desc(desc_dq_slot, cutlass.Int32(0), head_o, tok_coord) - if chunk_idx < wend: + if chunk_idx < write_end: tma_store_tile(sdQ_tma[dq_idx], dq_slice, acquire=False) tma_store_commit() @@ -711,7 +619,7 @@ def tmastg_warp( bars.mb_dk_tmastg_ready[dk_idx].wait(dk_index.phase) dk_index = advance(dk_index, cfg.smem_dk_stages) dk_slice = tma_slice_runtime_desc(desc_dk_slot, cutlass.Int32(0), head_o, tok_coord) - if chunk_idx < wend: + if chunk_idx < write_end: tma_store_tile(sdK_tma[dk_idx], dk_slice, acquire=False) tma_store_commit() @@ -721,7 +629,7 @@ def tmastg_warp( bars.mb_dq_tmastg_done[dq_idx].arrive() tma_store_wait(0) bars.mb_dk_tmastg_done[dk_idx].arrive() - tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one) @cute.jit @@ -742,68 +650,73 @@ def gate_beta_warp( sCumsumlog, sCumprod, sBeta, - sSched, + sScheduler, bars, ): """Gate/beta LOAD + STORE warp role (warp 10): per-chunk Gate/Beta G->S loads and the dGate/dBeta stores.""" nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) + elect_one = nvvm.elect_sync() + gate_index = PipelineState.start(phase=1) beta_index = PipelineState.start(phase=1) gate_store_index = PipelineState.start(phase=0) beta_store_index = PipelineState.start(phase=0) - lidx = tidx % cfg.threads_per_warp + scheduler_state = PipelineState.start(phase=0) + + lane_idx = tidx % cfg.threads_per_warp + n_cols = cfg.b_t // cfg.threads_per_warp - a_l2 = cutlass.Float32(0.0) + a = cutlass.Float32(0.0) bias = cutlass.Float32(0.0) - sched_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) FIRST_STATE_CHUNK = 0 if cfg.use_initial_state else 1 while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) - num_item_chunks = cend - wstart + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) + num_item_chunks = compute_end - write_start if cutlass.const_expr(cfg.safe_gate): if num_item_chunks > 0: - # per-head transform constants, fixed for the whole tile - a_l2 = -cute.math.exp2(mA_log[head_idx].to(cutlass.Float32) * cutlass.Float32(RCP_LN2), fastmath=True) * cutlass.Float32(RCP_LN2) + a = -cute.math.exp2(mA_log[head_idx].to(cutlass.Float32) * cutlass.Float32(RCP_LN2), fastmath=True) * cutlass.Float32(RCP_LN2) bias = mDt_bias[head_idx].to(cutlass.Float32) - # dGate/dBeta ownership: mask stores past the item's write range - write_end = batch_start + wend * cfg.b_t - write_end = write_end if write_end < batch_end else batch_end + write_end_token = batch_start + write_end * cfg.b_t + write_end_token = write_end_token if write_end_token < batch_end else batch_end - # ---- prefetch: the FIRST backward chunk's Gate/Beta ---------------------- + # ---- prefetch: the FIRST backward chunk's Gate/Beta -------------------------- if num_item_chunks > 0: - chunk_offset = batch_start + (cend - 1) * cfg.b_t + chunk_offset = batch_start + (compute_end - 1) * cfg.b_t gGate = cute.domain_offset((chunk_offset,), mGate[None, head_idx]) gBeta = cute.domain_offset((chunk_offset,), mBeta[None, head_idx]) gate_idx = gate_index.idx gate_index = advance(gate_index, cfg.smem_gate_stages) - pos_valid = [None] * n_cols - for col in cutlass.range_constexpr(n_cols): - pos = lidx + col * cfg.threads_per_warp - pos_valid[col] = cute.elem_less(chunk_offset + pos, batch_end) + pos_valid = [(chunk_offset + lane_idx + col * cfg.threads_per_warp) < batch_end for col in range(n_cols)] - # ---- Gate load: GMEM -> SMEM (OOB neutral: 1.0 -> log2 = 0.0) -------- - gate_vals = [cutlass.Float32(0.0)] * n_cols - for col in cutlass.range_constexpr(n_cols): - pos = lidx + col * cfg.threads_per_warp - oob_neutral = cutlass.Float32(0.0) if cutlass.const_expr(cfg.log_gate) else cutlass.Float32(1.0) - gate_vals[col] = gGate[pos] if pos_valid[col] else oob_neutral + # ---- Gate load: GMEM -> SMEM (OOB neutral: 1.0 -> log2 = 0.0) ------------ + oob_neutral = cutlass.Float32(0.0) if cutlass.const_expr(cfg.log_gate) else cutlass.Float32(1.0) + gate_vals = [gGate[lane_idx + col * cfg.threads_per_warp] if pos_valid[col] else oob_neutral for col in range(n_cols)] if cutlass.const_expr(cfg.safe_gate): - for col in cutlass.range_constexpr(n_cols): - contrib = a_l2 * softplus(gate_vals[col] + bias) - gate_vals[col] = contrib if pos_valid[col] else cutlass.Float32(0.0) + for col in cutlass.range_constexpr(0, n_cols, 2): + biased_lo, biased_hi = fadd2(gate_vals[col], gate_vals[col + 1], bias, bias) + sp_lo, sp_hi = softplus2(biased_lo, biased_hi) + contrib_lo, contrib_hi = fmul2(sp_lo, sp_hi, a, a) + gate_vals[col] = contrib_lo if pos_valid[col] else cutlass.Float32(0.0) + gate_vals[col + 1] = contrib_hi if pos_valid[col + 1] else cutlass.Float32(0.0) elif cutlass.const_expr(cfg.log_gate): - for col in cutlass.range_constexpr(n_cols): - gate_vals[col] = gate_vals[col] * cutlass.Float32(RCP_LN2) + rcp_ln2 = opaque_f32_zero() + cutlass.Float32(RCP_LN2) + for col in cutlass.range_constexpr(0, n_cols, 2): + gate_vals[col], gate_vals[col + 1] = fmul2(gate_vals[col], gate_vals[col + 1], rcp_ln2, rcp_ln2) else: - for col in cutlass.range_constexpr(n_cols): - gate_vals[col] = cute.math.log2(gate_vals[col] + 1e-10, fastmath=True) + floor = cutlass.Float32(1e-10) + for col in cutlass.range_constexpr(0, n_cols, 2): + shifted_lo, shifted_hi = fadd2(gate_vals[col], gate_vals[col + 1], floor, floor) + gate_vals[col] = cute.math.log2(shifted_lo, fastmath=True) + gate_vals[col + 1] = cute.math.log2(shifted_hi, fastmath=True) for offset in [1, 2, 4, 8, 16]: for col in cutlass.range_constexpr(n_cols): n = nvvm.shfl_sync(0xFFFFFFFF, gate_vals[col], offset, 0, kind=nvvm.Shfl.UP) - if lidx >= offset: + if lane_idx >= offset: gate_vals[col] = gate_vals[col] + n for col in cutlass.range_constexpr(1, n_cols): last_v = nvvm.shfl_sync( @@ -816,29 +729,27 @@ def gate_beta_warp( gate_vals[col] += last_v for col in cutlass.range_constexpr(n_cols): - pos = lidx + col * cfg.threads_per_warp + pos = lane_idx + col * cfg.threads_per_warp sCumsumlog[pos, 0, gate_idx] = gate_vals[col] sCumprod[pos, 0, gate_idx] = cute.math.exp2(gate_vals[col], fastmath=True) bars.mb_gate_ready[gate_idx].arrive() - # ---- Beta load: GMEM -> SMEM (per-element cp.async) ------------------ + # ---- Beta load: GMEM -> SMEM (per-element cp.async) ---------------------- beta_idx = beta_index.idx beta_index = advance(beta_index, cfg.smem_beta_stages) if cutlass.const_expr(cfg.beta_sigmoid): - # io-dtype logits -> sigmoid (tanh identity) -> fp32 SMEM for col in cutlass.range_constexpr(n_cols): - pos = lidx + col * cfg.threads_per_warp + pos = lane_idx + col * cfg.threads_per_warp beta_value = cutlass.Float32(0.0) if pos_valid[col]: beta_value = gBeta[pos].to(cutlass.Float32) - half = cutlass.Float32(0.5) - beta_value = (cute.math.tanh(beta_value * half, approx=True) * half + half).to(mBeta.element_type).to(cutlass.Float32) + beta_value = sigmoid(beta_value).to(mBeta.element_type).to(cutlass.Float32) sBeta[pos, 0, beta_idx] = beta_value bars.mb_beta_ready[beta_idx].arrive() else: for col in cutlass.range_constexpr(n_cols): - pos = lidx + col * cfg.threads_per_warp + pos = lane_idx + col * cfg.threads_per_warp src = gBeta.iterator + gBeta.layout((pos,)) dst = sBeta.iterator + sBeta.layout((pos, 0, beta_idx)) cp_size = cutlass.Int32(4) * cutlass.Int32(pos_valid[col]) @@ -846,38 +757,39 @@ def gate_beta_warp( nvvm.cp_async_mbarrier_arrive(bars.mb_beta_ready[beta_idx].smem_ptr, noinc=True) for rev_idx in cutlass.range(num_item_chunks): - # ---- prefetch the NEXT chunk's Gate/Beta ----------------------------- + # ---- prefetch the NEXT chunk's Gate/Beta --------------------------------- if rev_idx + 1 < num_item_chunks: - chunk_offset = batch_start + (cend - 2 - rev_idx) * cfg.b_t + chunk_offset = batch_start + (compute_end - 2 - rev_idx) * cfg.b_t gGate = cute.domain_offset((chunk_offset,), mGate[None, head_idx]) gBeta = cute.domain_offset((chunk_offset,), mBeta[None, head_idx]) gate_idx = gate_index.idx gate_index = advance(gate_index, cfg.smem_gate_stages) - pos_valid = [None] * n_cols - for col in cutlass.range_constexpr(n_cols): - pos = lidx + col * cfg.threads_per_warp - pos_valid[col] = cute.elem_less(chunk_offset + pos, batch_end) + pos_valid = [(chunk_offset + lane_idx + col * cfg.threads_per_warp) < batch_end for col in range(n_cols)] - gate_vals = [cutlass.Float32(0.0)] * n_cols - for col in cutlass.range_constexpr(n_cols): - pos = lidx + col * cfg.threads_per_warp - oob_neutral = cutlass.Float32(0.0) if cutlass.const_expr(cfg.log_gate) else cutlass.Float32(1.0) - gate_vals[col] = gGate[pos] if pos_valid[col] else oob_neutral + oob_neutral = cutlass.Float32(0.0) if cutlass.const_expr(cfg.log_gate) else cutlass.Float32(1.0) + gate_vals = [gGate[lane_idx + col * cfg.threads_per_warp] if pos_valid[col] else oob_neutral for col in range(n_cols)] if cutlass.const_expr(cfg.safe_gate): - for col in cutlass.range_constexpr(n_cols): - contrib = a_l2 * softplus(gate_vals[col] + bias) - gate_vals[col] = contrib if pos_valid[col] else cutlass.Float32(0.0) + for col in cutlass.range_constexpr(0, n_cols, 2): + biased_lo, biased_hi = fadd2(gate_vals[col], gate_vals[col + 1], bias, bias) + sp_lo, sp_hi = softplus2(biased_lo, biased_hi) + contrib_lo, contrib_hi = fmul2(sp_lo, sp_hi, a, a) + gate_vals[col] = contrib_lo if pos_valid[col] else cutlass.Float32(0.0) + gate_vals[col + 1] = contrib_hi if pos_valid[col + 1] else cutlass.Float32(0.0) elif cutlass.const_expr(cfg.log_gate): - for col in cutlass.range_constexpr(n_cols): - gate_vals[col] = gate_vals[col] * cutlass.Float32(RCP_LN2) + rcp_ln2 = opaque_f32_zero() + cutlass.Float32(RCP_LN2) + for col in cutlass.range_constexpr(0, n_cols, 2): + gate_vals[col], gate_vals[col + 1] = fmul2(gate_vals[col], gate_vals[col + 1], rcp_ln2, rcp_ln2) else: - for col in cutlass.range_constexpr(n_cols): - gate_vals[col] = cute.math.log2(gate_vals[col] + 1e-10, fastmath=True) + floor = cutlass.Float32(1e-10) + for col in cutlass.range_constexpr(0, n_cols, 2): + shifted_lo, shifted_hi = fadd2(gate_vals[col], gate_vals[col + 1], floor, floor) + gate_vals[col] = cute.math.log2(shifted_lo, fastmath=True) + gate_vals[col + 1] = cute.math.log2(shifted_hi, fastmath=True) for offset in [1, 2, 4, 8, 16]: for col in cutlass.range_constexpr(n_cols): n = nvvm.shfl_sync(0xFFFFFFFF, gate_vals[col], offset, 0, kind=nvvm.Shfl.UP) - if lidx >= offset: + if lane_idx >= offset: gate_vals[col] = gate_vals[col] + n for col in cutlass.range_constexpr(1, n_cols): last_v = nvvm.shfl_sync( @@ -890,7 +802,7 @@ def gate_beta_warp( gate_vals[col] += last_v for col in cutlass.range_constexpr(n_cols): - pos = lidx + col * cfg.threads_per_warp + pos = lane_idx + col * cfg.threads_per_warp sCumsumlog[pos, 0, gate_idx] = gate_vals[col] sCumprod[pos, 0, gate_idx] = cute.math.exp2(gate_vals[col], fastmath=True) @@ -899,64 +811,63 @@ def gate_beta_warp( beta_idx = beta_index.idx beta_index = advance(beta_index, cfg.smem_beta_stages) if cutlass.const_expr(cfg.beta_sigmoid): - # io-dtype logits -> sigmoid (tanh identity) -> fp32 SMEM for col in cutlass.range_constexpr(n_cols): - pos = lidx + col * cfg.threads_per_warp + pos = lane_idx + col * cfg.threads_per_warp beta_value = cutlass.Float32(0.0) if pos_valid[col]: beta_value = gBeta[pos].to(cutlass.Float32) - half = cutlass.Float32(0.5) - beta_value = (cute.math.tanh(beta_value * half, approx=True) * half + half).to(mBeta.element_type).to(cutlass.Float32) + beta_value = sigmoid(beta_value).to(mBeta.element_type).to(cutlass.Float32) sBeta[pos, 0, beta_idx] = beta_value bars.mb_beta_ready[beta_idx].arrive() else: for col in cutlass.range_constexpr(n_cols): - pos = lidx + col * cfg.threads_per_warp + pos = lane_idx + col * cfg.threads_per_warp src = gBeta.iterator + gBeta.layout((pos,)) dst = sBeta.iterator + sBeta.layout((pos, 0, beta_idx)) cp_size = cutlass.Int32(4) * cutlass.Int32(pos_valid[col]) nvvm.cp_async_shared_global(dst, src, 4, nvvm.LoadCacheModifier.CA, cp_size=cp_size) nvvm.cp_async_mbarrier_arrive(bars.mb_beta_ready[beta_idx].smem_ptr, noinc=True) - # ---- store-ready wait + in-place store back -------------------------- - st_offset = batch_start + (cend - 1 - rev_idx) * cfg.b_t - gGate_st = cute.domain_offset((st_offset,), mDgate[None, head_idx]) - gBeta_st = cute.domain_offset((st_offset,), mDbeta[None, head_idx]) - g_st_idx = gate_store_index.idx - bars.mb_gate_done[g_st_idx].wait(gate_store_index.phase) + # ---- dGate/dBeta store: SMEM -> GMEM (in place) -------------------------- + st_offset = batch_start + (compute_end - 1 - rev_idx) * cfg.b_t + gDgate = cute.domain_offset((st_offset,), mDgate[None, head_idx]) + gDbeta = cute.domain_offset((st_offset,), mDbeta[None, head_idx]) + + # ---- dGate store --------------------------------------------------------- + gate_store_idx = gate_store_index.idx + bars.mb_gate_done[gate_store_idx].wait(gate_store_index.phase) gate_store_index = advance(gate_store_index, cfg.smem_gate_stages) - dgate_vals = [cutlass.Float32(0.0)] * n_cols - for col in cutlass.range_constexpr(n_cols): - pos = lidx + col * cfg.threads_per_warp - dgate_vals[col] = sCumsumlog[pos, 0, g_st_idx] + dgate_vals = [sCumsumlog[lane_idx + col * cfg.threads_per_warp, 0, gate_store_idx] for col in range(n_cols)] for offset in [1, 2, 4, 8, 16]: for col in cutlass.range_constexpr(n_cols): n = nvvm.shfl_sync(0xFFFFFFFF, dgate_vals[col], offset, 31, kind=nvvm.Shfl.DOWN) - if lidx < cfg.threads_per_warp - offset: + if lane_idx < cfg.threads_per_warp - offset: dgate_vals[col] = dgate_vals[col] + n for col in cutlass.range_constexpr(n_cols - 1): rev_col = cutlass.const_expr(n_cols - 2 - col) later_total = nvvm.shfl_sync(0xFFFFFFFF, dgate_vals[rev_col + 1], 0, 0, kind=nvvm.Shfl.IDX) dgate_vals[rev_col] = dgate_vals[rev_col] + later_total for col in cutlass.range_constexpr(n_cols): - pos = lidx + col * cfg.threads_per_warp - if cute.elem_less(st_offset + pos, write_end): - gGate_st[pos] = dgate_vals[col] - b_st_idx = beta_store_index.idx - bars.mb_beta_done[b_st_idx].wait(beta_store_index.phase) + pos = lane_idx + col * cfg.threads_per_warp + if st_offset + pos < write_end_token: + gDgate[pos] = dgate_vals[col] + + # ---- dBeta store --------------------------------------------------------- + beta_store_idx = beta_store_index.idx + bars.mb_beta_done[beta_store_idx].wait(beta_store_index.phase) beta_store_index = advance(beta_store_index, cfg.smem_beta_stages) for col in cutlass.range_constexpr(n_cols): - pos = lidx + col * cfg.threads_per_warp - if cute.elem_less(st_offset + pos, write_end): + pos = lane_idx + col * cfg.threads_per_warp + if st_offset + pos < write_end_token: if cutlass.const_expr(cfg.beta_sigmoid): - gBeta_st[pos] = sBeta[pos, 0, b_st_idx].to(mDbeta.element_type) + gDbeta[pos] = sBeta[pos, 0, beta_store_idx].to(mDbeta.element_type) else: - gBeta_st[pos] = sBeta[pos, 0, b_st_idx] - tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + gDbeta[pos] = sBeta[pos, 0, beta_store_idx] + tile_idx, scheduler_state = scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one) @cute.jit -def mma_warp( +def tcgen05_mma_warp( cfg, total_tiles, bidx, @@ -968,30 +879,29 @@ def mma_warp( sQ_trans, sK, sK_trans, + sV_trans, sV, - sV_kmaj, + sdO_trans, sdO, - sdO_kmaj, + sState_trans, sState, - sState_kmaj, sTinv, sTinv_trans, sA, sA_trans, sDa, sDa_trans, - sDstate, + sDstate_trans, sDm, sDm_trans, - sdV_kmaj, - sSched, + sdV, + sScheduler, bars, ): """MMA issuer role (warp 8): persistent scheduler loop issuing every tcgen05 GEMM.""" - elect_one = nvvm.elect_sync() - nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) + kk_acc_index = PipelineState.start(phase=0) dk_total_index = PipelineState.start(phase=1) du_scale_index = PipelineState.start(phase=0) @@ -1013,19 +923,22 @@ def mma_warp( dq_scale_index = PipelineState.start(phase=0) dq_total_index = PipelineState.start(phase=1) a_index = PipelineState.start(phase=0) - do_prime_inp_ready = PipelineState.start(phase=0) - du_inp_ready = PipelineState.start(phase=0) - dyp_inp_ready = PipelineState.start(phase=0) - dstate_inp_index = PipelineState.start(phase=0) + do_prime_input_ready = PipelineState.start(phase=0) + du_input_ready = PipelineState.start(phase=0) + dyp_input_ready = PipelineState.start(phase=0) + dstate_input_index = PipelineState.start(phase=0) + + elect_one = nvvm.elect_sync() nvvm.tcgen05_alloc(tmem_base_slot, cutlass.Int32(512), group=nvvm.CTAGroup.CTA_1) - nvvm.barrier_cta_sync_aligned( - cfg.tmem_alloc_barrier_id, - thread_count=cfg.tmem_alloc_barrier_threads, - ) + nvvm.barrier_cta_sync_aligned(cfg.tmem_lifecycle_barrier_id, thread_count=cfg.tmem_user_threads) tmem_base = tmem_base_slot.load() + tmem_col = tmem_base & 0xFFFF + tmem_row = tmem_base >> 16 + row_lo_addr = tmem_row << 16 + row_hi_addr = (tmem_row + 16) << 16 - # ---- chunk-invariant GEMM descriptors ---------------------------------------- + # ---- chunk-invariant GEMM descriptors -------------------------------------------- bpe = cfg.io_dtype.width // 8 idesc_qk = nvvm.Tcgen05InstrDesc.build( c_dtype=cutlass.Float32, @@ -1034,7 +947,7 @@ def mma_warp( n_dim=cfg.b_t, m_dim=cfg.b_t, ) - bmm_qk_desc = MmaDesc( + bmm_k_k_desc = MmaDesc( M=cfg.b_t, N=cfg.b_t, K=cfg.d_k, @@ -1046,11 +959,12 @@ def mma_warp( idesc=idesc_qk, kind=nvvm.Tcgen05MMAKind.F16, ) - tmem_shared_acc_col = tmem_base + cfg.tmem_shared_acc_offset - tmem_shared_inp_col = tmem_base + cfg.tmem_shared_inp_offset + bmm_q_k_desc = bmm_k_k_desc + tmem_shared_acc_col = tmem_col + cfg.tmem_shared_acc_offset + tmem_shared_input_col = tmem_col + cfg.tmem_shared_input_offset SHARED_INP_STAGE_COLS = cfg.b_t // 2 - tmem_do_prime_col = tmem_shared_inp_col - tmem_du_col = tmem_shared_inp_col + SHARED_INP_STAGE_COLS + tmem_do_prime_col = tmem_shared_input_col + tmem_du_col = tmem_shared_input_col + SHARED_INP_STAGE_COLS tmem_dyp_col = tmem_du_col ACC_STAGE_COLS = cfg.b_t tmem_acc_a = tmem_shared_acc_col @@ -1062,7 +976,7 @@ def mma_warp( tmem_a_col = tmem_acc_b tmem_u_col = tmem_acc_b tmem_da_col = tmem_acc_b - tmem_dk_state_path_col = tmem_shared_inp_col + tmem_dk_state_path_col = tmem_shared_input_col idesc_dv = nvvm.Tcgen05InstrDesc.build( c_dtype=cutlass.Float32, @@ -1071,7 +985,7 @@ def mma_warp( n_dim=cfg.b_t, m_dim=cfg.d_v, ) - bmm_dv_desc = MmaDesc( + bmm_dstate_k_desc = MmaDesc( M=cfg.d_v, N=cfg.b_t, K=cfg.d_k, @@ -1084,8 +998,8 @@ def mma_warp( idesc=idesc_dv, kind=nvvm.Tcgen05MMAKind.F16, ) - tmem_dstate_inp_col = tmem_base + cfg.tmem_dstate_inp_offset - tmem_dvdk_acc_col = tmem_base + cfg.tmem_dvdk_acc_offset + tmem_dstate_input_col = tmem_col + cfg.tmem_dstate_input_offset + tmem_dvdk_acc_col = tmem_col + cfg.tmem_dvdk_acc_offset DSTATE_INP_STAGE_COLS = cfg.d_k // 2 idesc_k_state = nvvm.Tcgen05InstrDesc.build( @@ -1096,7 +1010,7 @@ def mma_warp( m_dim=cfg.d_v, a_major=1, ) - bmm_k_state_desc = MmaDesc( + bmm_state_k_desc = MmaDesc( M=cfg.d_v, N=cfg.b_t, K=cfg.d_k, @@ -1109,14 +1023,14 @@ def mma_warp( idesc=idesc_k_state, kind=nvvm.Tcgen05MMAKind.F16, ) - idesc_k_state_kmaj = nvvm.Tcgen05InstrDesc.build( + idesc_k_state = nvvm.Tcgen05InstrDesc.build( c_dtype=cutlass.Float32, a_dtype=cfg.io_dtype, b_dtype=cfg.io_dtype, n_dim=cfg.b_t, m_dim=cfg.d_v, ) - bmm_k_state_kmaj_desc = MmaDesc( + bmm_state_k_desc = MmaDesc( M=cfg.d_v, N=cfg.b_t, K=cfg.d_k, @@ -1126,7 +1040,7 @@ def mma_warp( btranspose=False, atranspose=False, cta_group=1, - idesc=idesc_k_state_kmaj, + idesc=idesc_k_state, kind=nvvm.Tcgen05MMAKind.F16, ) @@ -1139,7 +1053,7 @@ def mma_warp( a_major=1, b_major=1, ) - bmm_du_desc = MmaDesc( + bmm_do_a_desc = MmaDesc( M=cfg.d_v, N=cfg.b_t, K=cfg.b_t, @@ -1160,7 +1074,7 @@ def mma_warp( m_dim=cfg.d_k, b_major=1, ) - bmm_dstate_upd_desc = MmaDesc( + bmm_do_prime_q_desc = MmaDesc( M=cfg.d_k, N=cfg.d_v, K=cfg.b_t, @@ -1173,13 +1087,14 @@ def mma_warp( idesc=idesc_dstate_upd, kind=nvvm.Tcgen05MMAKind.F16, ) + bmm_dy_prime_k_desc = bmm_do_prime_q_desc SHARED_INP_STAGE_COLS = cfg.b_t // 2 - tmem_do_prime_col = tmem_shared_inp_col - tmem_du_col = tmem_shared_inp_col + SHARED_INP_STAGE_COLS + tmem_do_prime_col = tmem_shared_input_col + tmem_du_col = tmem_shared_input_col + SHARED_INP_STAGE_COLS tmem_dyp_col = tmem_du_col - tmem_shared_inp_col = tmem_base + cfg.tmem_shared_inp_offset - tmem_dstate_acc_col = tmem_base + cfg.tmem_dstate_acc_offset - tmem_y_col = tmem_base + cfg.tmem_y_offset + tmem_shared_input_col = tmem_col + cfg.tmem_shared_input_offset + tmem_dstate_acc_col = tmem_col + cfg.tmem_dstate_acc_offset + tmem_y_col = tmem_col + cfg.tmem_y_offset idesc_dy = nvvm.Tcgen05InstrDesc.build( c_dtype=cutlass.Float32, @@ -1189,7 +1104,7 @@ def mma_warp( m_dim=cfg.d_v, b_major=1, ) - bmm_dy_desc = MmaDesc( + bmm_du_t_inv_trans_desc = MmaDesc( M=cfg.d_v, N=cfg.b_t, K=cfg.b_t, @@ -1209,7 +1124,7 @@ def mma_warp( n_dim=cfg.b_t, m_dim=cfg.b_t, ) - bmm_da_desc = MmaDesc( + bmm_do_u_desc = MmaDesc( M=cfg.b_t, N=cfg.b_t, K=cfg.d_v, @@ -1221,6 +1136,7 @@ def mma_warp( idesc=idesc_da, kind=nvvm.Tcgen05MMAKind.F16, ) + bmm_dy_u_desc = bmm_do_u_desc idesc_u = nvvm.Tcgen05InstrDesc.build( c_dtype=cutlass.Float32, @@ -1229,7 +1145,7 @@ def mma_warp( n_dim=cfg.b_t, m_dim=cfg.d_v, ) - bmm_u_desc = MmaDesc( + bmm_y_t_inv_desc = MmaDesc( M=cfg.d_v, N=cfg.b_t, K=cfg.b_t, @@ -1250,7 +1166,7 @@ def mma_warp( m_dim=cfg.d_k, a_major=1, ) - bmm_dqdk_inter_at_desc = MmaDesc( + bmm_state_do_desc = MmaDesc( M=cfg.d_k, N=cfg.b_t, K=cfg.d_v, @@ -1263,6 +1179,8 @@ def mma_warp( idesc=idesc_dqdk_inter_at, kind=nvvm.Tcgen05MMAKind.F16, ) + bmm_dstate_u_desc = bmm_state_do_desc + bmm_state_dy_desc = bmm_state_do_desc idesc_dka = nvvm.Tcgen05InstrDesc.build( c_dtype=cutlass.Float32, a_dtype=cfg.io_dtype, @@ -1272,7 +1190,7 @@ def mma_warp( a_major=1, b_major=1, ) - bmm_dka_desc = MmaDesc( + bmm_k_dm_trans_desc = MmaDesc( M=cfg.d_k, N=cfg.b_t, K=cfg.b_t, @@ -1285,6 +1203,7 @@ def mma_warp( idesc=idesc_dka, kind=nvvm.Tcgen05MMAKind.F16, ) + bmm_q_da_trans_desc = bmm_k_dm_trans_desc idesc_dqa = nvvm.Tcgen05InstrDesc.build( c_dtype=cutlass.Float32, a_dtype=cfg.io_dtype, @@ -1293,7 +1212,7 @@ def mma_warp( m_dim=cfg.d_k, a_major=1, ) - bmm_dqa_desc = MmaDesc( + bmm_k_dm_desc = MmaDesc( M=cfg.d_k, N=cfg.b_t, K=cfg.b_t, @@ -1306,11 +1225,12 @@ def mma_warp( idesc=idesc_dqa, kind=nvvm.Tcgen05MMAKind.F16, ) + bmm_k_da_desc = bmm_k_dm_desc - do_prime_inp_ptr = nvvm.make_tmem_ptr(tmem_do_prime_col, cutlass.Int8) - y_inp_ptr = nvvm.make_tmem_ptr(tmem_y_col, cutlass.Int8) - du_inp_ptr = nvvm.make_tmem_ptr(tmem_du_col, cutlass.Int8) - dyp_inp_ptr = nvvm.make_tmem_ptr(tmem_dyp_col, cutlass.Int8) + do_prime_input_ptr = nvvm.make_tmem_ptr(tmem_do_prime_col, cutlass.Int8) + y_input_ptr = nvvm.make_tmem_ptr(tmem_y_col, cutlass.Int8) + du_input_ptr = nvvm.make_tmem_ptr(tmem_du_col, cutlass.Int8) + dyp_input_ptr = nvvm.make_tmem_ptr(tmem_dyp_col, cutlass.Int8) kk_acc_ptr = nvvm.make_tmem_ptr(tmem_kk_col, cutlass.Float32) a_acc_ptr = nvvm.make_tmem_ptr(tmem_a_col, cutlass.Float32) k_state_acc_ptr = nvvm.make_tmem_ptr(tmem_k_state_col, cutlass.Float32) @@ -1320,50 +1240,52 @@ def mma_warp( dm_core_acc_ptr = nvvm.make_tmem_ptr(tmem_dm_core_col, cutlass.Float32) dk_state_path_acc_ptr = nvvm.make_tmem_ptr(tmem_dk_state_path_col, cutlass.Float32) dstate_acc_ptr = nvvm.make_tmem_ptr(tmem_dstate_acc_col, cutlass.Float32) - dq_acc_ptr = nvvm.make_tmem_ptr(tmem_dstate_inp_col, cutlass.Float32) + dq_acc_ptr = nvvm.make_tmem_ptr(tmem_dstate_input_col, cutlass.Float32) dvdk_acc_ptr = nvvm.make_tmem_ptr(tmem_dvdk_acc_col, cutlass.Float32) - # ---- warp-top descriptors (1-stage tiles are loop-constant; K advances) ---- + # ---- warp-top descriptors (1-stage tiles are loop-constant; K advances) ---------- d_q0 = sQ[0].desc() d_k0 = sK[0].desc() d_k_trans0 = sK_trans[0].desc() d_q_trans0 = sQ_trans[0].desc() + d_do_trans0 = sdO_trans[0].desc() d_do0 = sdO[0].desc() - d_do_kmaj0 = sdO_kmaj[0].desc() + d_state_trans0 = sState_trans[0].desc() d_state0 = sState[0].desc() - d_state_kmaj0 = sState_kmaj[0].desc() d_tinv0 = sTinv[0].desc() d_tinv_trans0 = sTinv_trans[0].desc() d_a_trans0 = sA_trans[0].desc() - d_v_kmaj0 = sV_kmaj[0].desc() - d_dv_kmaj0 = sdV_kmaj[0].desc() + d_v0 = sV[0].desc() + d_dv0 = sdV[0].desc() d_da0 = sDa[0].desc() d_da_trans0 = sDa_trans[0].desc() d_dm0 = sDm[0].desc() d_dm_trans0 = sDm_trans[0].desc() - d_dstate0 = sDstate[0].desc() + d_dstate_trans0 = sDstate_trans[0].desc() K_STAGE_BYTES = (cfg.k_cosize // cfg.smem_k_stages) * (cfg.io_dtype.width // 8) - sched_state = PipelineState.start(phase=0) + scheduler_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) FIRST_STATE_CHUNK = 0 if cfg.use_initial_state else 1 while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) - num_item_chunks = cend - wstart + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) + num_item_chunks = compute_end - write_start - # ---- chunks NT-2 .. 0 (backward): full body ------------------------------ + # ---- chunks NT-2 .. 0 (backward): full body ---------------------------------- for rev_idx in cutlass.range(num_item_chunks): - chunk_idx = cend - 1 - rev_idx + chunk_idx = compute_end - 1 - rev_idx have_dstate = cutlass.Boolean(True) if cutlass.const_expr(cfg.use_dstate_in) else rev_idx > 0 - # ---- KK = K(S) @ K^T ------------------------------------------------- + # ---- KK = K(S) @ K^T ----------------------------------------------------- k_idx = k_index.idx bars.mb_k_ready[k_idx].wait(k_index.phase) k_index = advance(k_index, cfg.smem_k_stages) desc_k = d_k0.advance_start_address(k_idx * K_STAGE_BYTES) mma_ss( - bmm_qk_desc, + bmm_k_k_desc, desc_k, desc_k, kk_acc_ptr, @@ -1373,14 +1295,14 @@ def mma_warp( if elect_one: bars.mb_kk_acc_ready[0].arrive(cta_group=1) - # ---- QK = Q(S) @ K^T ------------------------------------------------- + # ---- QK = Q(S) @ K^T ----------------------------------------------------- q_idx = q_index.idx bars.mb_q_ready[q_idx].wait(q_index.phase) q_index = advance(q_index, cfg.smem_q_stages) desc_q = d_q0 mma_ss( - bmm_qk_desc, + bmm_q_k_desc, desc_q, desc_k, a_acc_ptr, @@ -1390,16 +1312,16 @@ def mma_warp( if elect_one: bars.mb_a_acc_ready[0].arrive(cta_group=1) - # ---- k_state = state^T(S) @ K^T ----------------------------------------------- + # ---- k state = state^T(S) @ K^T ------------------------------------------ state_idx = state_index.idx if chunk_idx >= FIRST_STATE_CHUNK: bars.mb_state_ready[state_idx].wait(state_index.phase) state_index = advance(state_index, cfg.smem_state_stages) bars.mb_kk_acc_done[0].wait(kk_acc_index.phase) kk_acc_index = advance(kk_acc_index, 1) - desc_state = d_state_kmaj0 + desc_state = d_state0 mma_ss( - bmm_k_state_kmaj_desc, + bmm_state_k_desc, desc_state, desc_k, k_state_acc_ptr, @@ -1409,28 +1331,28 @@ def mma_warp( bars.mb_k_state_acc_ready[0].arrive(cta_group=1) # ---- dV inter = dstate^T(T) @ K ------------------------------------------ - dstate_inp_idx = dstate_inp_index.idx + dstate_input_idx = dstate_input_index.idx if have_dstate: - bars.mb_dstate_inp_ready[dstate_inp_idx].wait(dstate_inp_index.phase) - dstate_inp_index = advance(dstate_inp_index, cfg.tmem_dstate_inp_stages) + bars.mb_dstate_input_ready[dstate_input_idx].wait(dstate_input_index.phase) + dstate_input_index = advance(dstate_input_index, cfg.tmem_dstate_input_stages) bars.mb_dk_total_acc_done[0].wait(dk_total_index.phase) dk_total_index = advance(dk_total_index, 1) if have_dstate: - dstate_a_ptr = nvvm.make_tmem_ptr(tmem_dstate_inp_col + dstate_inp_idx * DSTATE_INP_STAGE_COLS, cutlass.Int8) - for sub in cutlass.range_constexpr(bmm_dv_desc.num_subtiles_B): - for k in cutlass.range_constexpr(bmm_dv_desc.sps_B): + dstate_a_ptr = nvvm.make_tmem_ptr(tmem_dstate_input_col + dstate_input_idx * DSTATE_INP_STAGE_COLS, cutlass.Int8) + for i in cutlass.range_constexpr(bmm_dstate_k_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_dstate_k_desc.sps_B): mma_ts_step( - bmm_dv_desc, - dstate_a_ptr.subview(sub * bmm_dv_desc.sps_B * bmm_dv_desc.tmem_advance_A), - desc_k + sub * (bmm_dv_desc.smem_subtile_B >> 4), + bmm_dstate_k_desc, + dstate_a_ptr.subview(i * bmm_dstate_k_desc.sps_B * bmm_dstate_k_desc.tmem_advance_A), + desc_k + i * (bmm_dstate_k_desc.smem_subtile_B >> 4), dvdk_acc_ptr, k, - cutlass.Boolean(sub + k > 0), + cutlass.Boolean(i + k > 0), ) if elect_one: bars.mb_du_scale_acc_ready[0].arrive(cta_group=1) - bars.mb_dstate_inp_done[dstate_inp_idx].arrive(cta_group=1) + bars.mb_dstate_input_done[dstate_input_idx].arrive(cta_group=1) # ---- dQ inter = state(S) @ dO^T ------------------------------------------ do_idx = do_index.idx @@ -1439,19 +1361,19 @@ def mma_warp( if chunk_idx >= FIRST_STATE_CHUNK: bars.mb_dq_acc_total_done[0].wait(dq_total_index.phase) dq_total_index = advance(dq_total_index, 1) - desc_state_kmaj_dq_inter = d_state0 - desc_do_kmaj_dq_inter = d_do_kmaj0 + desc_state_dq_inter = d_state_trans0 + desc_do_dq_inter = d_do0 mma_ss( - bmm_dqdk_inter_at_desc, - desc_state_kmaj_dq_inter, - desc_do_kmaj_dq_inter, + bmm_state_do_desc, + desc_state_dq_inter, + desc_do_dq_inter, dq_acc_ptr, accumulate=False, ) if elect_one: bars.mb_dq_acc_scale_ready[0].arrive(cta_group=1) - # ---- dU intra += dO^T(S) @ A ----------------------------------------- + # ---- dU intra += dO^T(S) @ A --------------------------------------------- du_a_idx = a_index.idx bars.mb_a_ready[du_a_idx].wait(a_index.phase) a_index = advance(a_index, cfg.smem_a_stages) @@ -1459,38 +1381,38 @@ def mma_warp( bars.mb_du_scale_acc_done[0].wait(du_scale_index.phase) du_scale_index = advance(du_scale_index, 1) - desc_do_mnmaj = d_do0 - desc_a_t = d_a_trans0 + desc_do_trans = d_do_trans0 + desc_a_trans = d_a_trans0 mma_ss( - bmm_du_desc, - desc_do_mnmaj, - desc_a_t, + bmm_do_a_desc, + desc_do_trans, + desc_a_trans, dvdk_acc_ptr, accumulate=have_dstate, ) if elect_one: bars.mb_du_total_acc_ready[0].arrive(cta_group=1) - # ---- dstate update += dO'^T(T) @ Q --------------------------------------- - bars.mb_do_prime_inp_ready[0].wait(do_prime_inp_ready.phase) - do_prime_inp_ready = advance(do_prime_inp_ready, 1) + # ---- dstate Q-term += dO'^T(T) @ Q --------------------------------------- + bars.mb_do_prime_input_ready[0].wait(do_prime_input_ready.phase) + do_prime_input_ready = advance(do_prime_input_ready, 1) dstate_idx = dstate_acc_index.idx bars.mb_dstate_scale_acc_done[dstate_idx].wait(dstate_acc_index.phase) dstate_acc_index = advance(dstate_acc_index, cfg.tmem_dstate_acc_stages) - desc_q_t = d_q_trans0 - for sub in cutlass.range_constexpr(bmm_dstate_upd_desc.num_subtiles_B): - for k in cutlass.range_constexpr(bmm_dstate_upd_desc.sps_B): + desc_q_trans = d_q_trans0 + for i in cutlass.range_constexpr(bmm_do_prime_q_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_do_prime_q_desc.sps_B): mma_ts_step( - bmm_dstate_upd_desc, - do_prime_inp_ptr.subview(sub * bmm_dstate_upd_desc.sps_B * bmm_dstate_upd_desc.tmem_advance_A), - desc_q_t + sub * (bmm_dstate_upd_desc.smem_subtile_B >> 4), + bmm_do_prime_q_desc, + do_prime_input_ptr.subview(i * bmm_do_prime_q_desc.sps_B * bmm_do_prime_q_desc.tmem_advance_A), + desc_q_trans + i * (bmm_do_prime_q_desc.smem_subtile_B >> 4), dstate_acc_ptr, k, - cutlass.Boolean(True) if cutlass.const_expr(sub + k > 0) else have_dstate, + cutlass.Boolean(True) if cutlass.const_expr(i + k > 0) else have_dstate, ) - # ---- U^T recompute = Y^T(T) @ T^T ------------------------------------ + # ---- U^T recompute = Y^T(T) @ T^T ---------------------------------------- tinv_idx = tinv_index.idx bars.mb_t_inv_ready[tinv_idx].wait(tinv_index.phase) tinv_index = advance(tinv_index, cfg.smem_t_inv_stages) @@ -1500,63 +1422,63 @@ def mma_warp( v_index = advance(v_index, cfg.smem_v_stages) desc_tinv = d_tinv0 - for sub in cutlass.range_constexpr(bmm_u_desc.num_subtiles_B): - for k in cutlass.range_constexpr(bmm_u_desc.sps_B): + for i in cutlass.range_constexpr(bmm_y_t_inv_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_y_t_inv_desc.sps_B): mma_ts_step( - bmm_u_desc, - y_inp_ptr.subview(sub * bmm_u_desc.sps_B * bmm_u_desc.tmem_advance_A), - desc_tinv + sub * (bmm_u_desc.smem_subtile_B >> 4), + bmm_y_t_inv_desc, + y_input_ptr.subview(i * bmm_y_t_inv_desc.sps_B * bmm_y_t_inv_desc.tmem_advance_A), + desc_tinv + i * (bmm_y_t_inv_desc.smem_subtile_B >> 4), u_acc_ptr, k, - cutlass.Boolean(sub + k > 0), + cutlass.Boolean(i + k > 0), ) if elect_one: bars.mb_u_acc_ready[0].arrive(cta_group=1) - # ---- dY = dU^T(T) @ T ------------------------------------------------ - bars.mb_du_inp_ready[0].wait(du_inp_ready.phase) - du_inp_ready = advance(du_inp_ready, 1) + # ---- dY = dU^T(T) @ T ---------------------------------------------------- + bars.mb_du_input_ready[0].wait(du_input_ready.phase) + du_input_ready = advance(du_input_ready, 1) - desc_tinv_t = d_tinv_trans0 - for sub in cutlass.range_constexpr(bmm_dy_desc.num_subtiles_B): - for k in cutlass.range_constexpr(bmm_dy_desc.sps_B): + desc_tinv_trans = d_tinv_trans0 + for i in cutlass.range_constexpr(bmm_du_t_inv_trans_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_du_t_inv_trans_desc.sps_B): mma_ts_step( - bmm_dy_desc, - du_inp_ptr.subview(sub * bmm_dy_desc.sps_B * bmm_dy_desc.tmem_advance_A), - desc_tinv_t + sub * (bmm_dy_desc.smem_subtile_B >> 4), + bmm_du_t_inv_trans_desc, + du_input_ptr.subview(i * bmm_du_t_inv_trans_desc.sps_B * bmm_du_t_inv_trans_desc.tmem_advance_A), + desc_tinv_trans + i * (bmm_du_t_inv_trans_desc.smem_subtile_B >> 4), dy_acc_ptr, k, - cutlass.Boolean(sub + k > 0), + cutlass.Boolean(i + k > 0), ) if elect_one: bars.mb_dy_acc_ready[0].arrive(cta_group=1) - # ---- dK_inter = dstate_entry(S) @ U^T ------------------------------------ + # ---- dK inter = dstate entry(S) @ U^T ------------------------------------ bars.mb_u_ready[0].wait(u_index.phase) u_index = advance(u_index, 1) if have_dstate: bars.mb_dstate_smem_ready[0].wait(dstate_smem_index.phase) dstate_smem_index = advance(dstate_smem_index, 1) - desc_dstate = d_dstate0 - desc_u_kmaj_dk_inter = d_v_kmaj0 + desc_dstate = d_dstate_trans0 + desc_u_dk_inter = d_v0 mma_ss( - bmm_dqdk_inter_at_desc, + bmm_dstate_u_desc, desc_dstate, - desc_u_kmaj_dk_inter, + desc_u_dk_inter, dvdk_acc_ptr, accumulate=False, ) if elect_one: bars.mb_dk_scale_acc_ready[0].arrive(cta_group=1) - # ---- dA_eff = dO(S) @ U^T -------------------------------------------- - desc_do_kmaj_da = d_do_kmaj0 - desc_u_kmaj_da = d_v_kmaj0 + # ---- dA eff = dO(S) @ U^T ------------------------------------------------ + desc_do_da = d_do0 + desc_u_da = d_v0 mma_ss( - bmm_da_desc, - desc_do_kmaj_da, - desc_u_kmaj_da, + bmm_do_u_desc, + desc_do_da, + desc_u_da, da_acc_ptr, accumulate=False, ) @@ -1564,17 +1486,17 @@ def mma_warp( bars.mb_da_acc_ready[0].arrive(cta_group=1) bars.mb_do_mma_done[do_idx].arrive(cta_group=1) - # ---- dM core = dY(S) @ U^T ------------------------------------------- + # ---- dM core = dY(S) @ U^T ----------------------------------------------- dv_ready_idx = dv_ready_index.idx bars.mb_dv_tmastg_ready[dv_ready_idx].wait(dv_ready_index.phase) dv_ready_index = advance(dv_ready_index, cfg.smem_dv_stages) - desc_dy_kmaj_dm = d_dv_kmaj0 - desc_u_kmaj_dm = d_v_kmaj0 + desc_dy_dm = d_dv0 + desc_u_dm = d_v0 mma_ss( - bmm_da_desc, - desc_dy_kmaj_dm, - desc_u_kmaj_dm, + bmm_dy_u_desc, + desc_dy_dm, + desc_u_dm, dm_core_acc_ptr, accumulate=False, ) @@ -1582,17 +1504,17 @@ def mma_warp( bars.mb_dm_acc_ready[0].arrive(cta_group=1) bars.mb_v_mma_done[v_idx].arrive(cta_group=1) - # ---- dstate update-term += dY'^T(T) @ K -------------------------------------- - bars.mb_dyp_inp_ready[0].wait(dyp_inp_ready.phase) - dyp_inp_ready = advance(dyp_inp_ready, 1) + # ---- dstate K-term += dY'^T(T) @ K --------------------------------------- + bars.mb_dyp_input_ready[0].wait(dyp_input_ready.phase) + dyp_input_ready = advance(dyp_input_ready, 1) - desc_k_t = d_k_trans0.advance_start_address(k_idx * K_STAGE_BYTES) - for sub in cutlass.range_constexpr(bmm_dstate_upd_desc.num_subtiles_B): - for k in cutlass.range_constexpr(bmm_dstate_upd_desc.sps_B): + desc_k_trans = d_k_trans0.advance_start_address(k_idx * K_STAGE_BYTES) + for i in cutlass.range_constexpr(bmm_dy_prime_k_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_dy_prime_k_desc.sps_B): mma_ts_step( - bmm_dstate_upd_desc, - dyp_inp_ptr.subview(sub * bmm_dstate_upd_desc.sps_B * bmm_dstate_upd_desc.tmem_advance_A), - desc_k_t + sub * (bmm_dstate_upd_desc.smem_subtile_B >> 4), + bmm_dy_prime_k_desc, + dyp_input_ptr.subview(i * bmm_dy_prime_k_desc.sps_B * bmm_dy_prime_k_desc.tmem_advance_A), + desc_k_trans + i * (bmm_dy_prime_k_desc.smem_subtile_B >> 4), dstate_acc_ptr, k, cutlass.Boolean(True), @@ -1600,19 +1522,19 @@ def mma_warp( if elect_one: bars.mb_dstate_acc_ready[dstate_idx].arrive(cta_group=1) - # ---- dK attn += Q^T(S) @ dA ------------------------------------------ + # ---- dK attn += Q^T(S) @ dA ---------------------------------------------- if have_dstate: bars.mb_dk_scale_acc_done[0].wait(dk_scale_index.phase) dk_scale_index = advance(dk_scale_index, 1) bars.mb_da_ready[0].wait(da_ready_index.phase) da_ready_index = advance(da_ready_index, 1) - desc_q_mnmaj_dk_attn = d_q_trans0 - desc_da_t = d_da_trans0 + desc_q_trans_dk_attn = d_q_trans0 + desc_da_trans = d_da_trans0 mma_ss( - bmm_dka_desc, - desc_q_mnmaj_dk_attn, - desc_da_t, + bmm_q_da_trans_desc, + desc_q_trans_dk_attn, + desc_da_trans, dvdk_acc_ptr, accumulate=have_dstate, ) @@ -1620,27 +1542,27 @@ def mma_warp( bars.mb_dk_attn_acc_ready[0].arrive(cta_group=1) bars.mb_q_mma_done[q_idx].arrive(cta_group=1) - # ---- dQ attn += K^T(S) @ dA ------------------------------------------ + # ---- dQ attn += K^T(S) @ dA ---------------------------------------------- if chunk_idx >= FIRST_STATE_CHUNK: bars.mb_dq_acc_scale_done[0].wait(dq_scale_index.phase) dq_scale_index = advance(dq_scale_index, 1) if chunk_idx < FIRST_STATE_CHUNK: bars.mb_dq_acc_total_done[0].wait(dq_total_index.phase) dq_total_index = advance(dq_total_index, 1) - desc_k_mnmaj_dq_attn = d_k_trans0.advance_start_address(k_idx * K_STAGE_BYTES) + desc_k_trans_dq_attn = d_k_trans0.advance_start_address(k_idx * K_STAGE_BYTES) desc_da = d_da0 if chunk_idx >= FIRST_STATE_CHUNK: mma_ss( - bmm_dqa_desc, - desc_k_mnmaj_dq_attn, + bmm_k_da_desc, + desc_k_trans_dq_attn, desc_da, dq_acc_ptr, accumulate=True, ) if chunk_idx < FIRST_STATE_CHUNK: mma_ss( - bmm_dqa_desc, - desc_k_mnmaj_dq_attn, + bmm_k_da_desc, + desc_k_trans_dq_attn, desc_da, dq_acc_ptr, accumulate=False, @@ -1649,14 +1571,14 @@ def mma_warp( bars.mb_dq_acc_total_ready[0].arrive(cta_group=1) bars.mb_a_done[du_a_idx].arrive(cta_group=1) - # ---- dK state-path = state(S) @ dY^T ----------------------------------------- + # ---- dK state-path = state(S) @ dY^T ------------------------------------- if chunk_idx >= FIRST_STATE_CHUNK: - desc_state_kmaj_spath = d_state0 - desc_dy_kmaj_spath = d_dv_kmaj0 + desc_state_spath = d_state_trans0 + desc_dy_spath = d_dv0 mma_ss( - bmm_dqdk_inter_at_desc, - desc_state_kmaj_spath, - desc_dy_kmaj_spath, + bmm_state_dy_desc, + desc_state_spath, + desc_dy_spath, dk_state_path_acc_ptr, accumulate=False, ) @@ -1666,24 +1588,24 @@ def mma_warp( if elect_one: bars.mb_sdv_done[0].arrive(cta_group=1) - # ---- dK dM-terms += K^T(S) @ dM^T + K^T(S) @ dM ---------------------- + # ---- dK dM-terms += K^T(S) @ dM^T + K^T(S) @ dM -------------------------- bars.mb_dm_acc_done[0].wait(dm_index.phase) dm_index = advance(dm_index, 1) bars.mb_dk_attn_acc_done[0].wait(dk_attn_index.phase) dk_attn_index = advance(dk_attn_index, 1) - desc_k_mnmaj_dm_terms = d_k_trans0.advance_start_address(k_idx * K_STAGE_BYTES) - desc_dm_t = d_dm_trans0 + desc_k_trans_dm_terms = d_k_trans0.advance_start_address(k_idx * K_STAGE_BYTES) + desc_dm_trans = d_dm_trans0 mma_ss( - bmm_dka_desc, - desc_k_mnmaj_dm_terms, - desc_dm_t, + bmm_k_dm_trans_desc, + desc_k_trans_dm_terms, + desc_dm_trans, dvdk_acc_ptr, accumulate=True, ) desc_dm = d_dm0 mma_ss( - bmm_dqa_desc, - desc_k_mnmaj_dm_terms, + bmm_k_dm_desc, + desc_k_trans_dm_terms, desc_dm, dvdk_acc_ptr, accumulate=True, @@ -1692,14 +1614,13 @@ def mma_warp( bars.mb_dk_total_acc_ready[0].arrive(cta_group=1) bars.mb_k_mma_done[k_idx].arrive(cta_group=1) - tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one) bars.mb_dk_total_acc_done[0].wait(dk_total_index.phase) - bars.mb_tmem_done[0].wait(0) nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_1) nvvm.tcgen05_dealloc( - nvvm.make_tmem_ptr(tmem_base, cutlass.Int8), + nvvm.make_tmem_ptr(tmem_col, cutlass.Int8), cutlass.Int32(512), group=nvvm.CTAGroup.CTA_1, ) @@ -1723,49 +1644,54 @@ def tmaldg_warp( desc_v_base, desc_do_base, desc_checkpoint_base, - sSched, - mSched, + sScheduler, + mScheduler, bars, ): """TMA-LDG warp role (warp 9): persistent tile loop issuing every Q/K/V/dO/state TMA load.""" - elect_one = nvvm.elect_sync() nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) + q_index = PipelineState.start(phase=1) k_index = PipelineState.start(phase=1) v_index = PipelineState.start(phase=1) do_index = PipelineState.start(phase=1) state_index = PipelineState.start(phase=1) - sched_state = PipelineState.start(phase=1) + scheduler_state = PipelineState.start(phase=1) + tail_count = ((total_tiles - cutlass.Int32(1)) % num_ctas) + cutlass.Int32(1) + tail_base = (total_tiles - tail_count) if tail_count * 2 >= num_ctas else total_tiles + tail_row = tail_base + cute.arch.smid() + tail_row = tail_row if tail_row < total_tiles else cutlass.Int32(1 << 28) + + elect_one = nvvm.elect_sync() tile_idx = cutlass.Int32(bidx) FIRST_STATE_CHUNK = 0 if cfg.use_initial_state else 1 - bpe = cfg.io_dtype.width // 8 - granule_elems = 128 // bpe + granule_elements = 128 // bpe bt = cfg.b_t - q_stage_elems = cfg.q_cosize // cfg.smem_q_stages - k_stage_elems = cfg.k_cosize // cfg.smem_k_stages + q_stage_elements = cfg.q_cosize // cfg.smem_q_stages + k_stage_elements = cfg.k_cosize // cfg.smem_k_stages sQ_tma = SmemTile( base=sQ_raw, - elems_per_stage=q_stage_elems, + elems_per_stage=q_stage_elements, stages=cfg.smem_q_stages, leading_byte_offset=0, stride_byte_offset=0, layout=0, tma_loads_per_tile=2, - tma_granu_elems=granule_elems, - tma_subtile_stride_elems=bt * granule_elems, + tma_granu_elems=granule_elements, + tma_subtile_stride_elems=bt * granule_elements, ) sK_tma = SmemTile( base=sK_raw, - elems_per_stage=k_stage_elems, + elems_per_stage=k_stage_elements, stages=cfg.smem_k_stages, leading_byte_offset=0, stride_byte_offset=0, layout=0, tma_loads_per_tile=2, - tma_granu_elems=granule_elems, - tma_subtile_stride_elems=bt * granule_elems, + tma_granu_elems=granule_elements, + tma_subtile_stride_elems=bt * granule_elements, ) sV_tma = SmemTile( base=sV_raw, @@ -1775,7 +1701,7 @@ def tmaldg_warp( stride_byte_offset=0, layout=0, tma_loads_per_tile=2, - tma_granu_elems=granule_elems, + tma_granu_elems=granule_elements, tma_subtile_stride_elems=4096, ) sdO_tma = SmemTile( @@ -1786,7 +1712,7 @@ def tmaldg_warp( stride_byte_offset=0, layout=0, tma_loads_per_tile=2, - tma_granu_elems=granule_elems, + tma_granu_elems=granule_elements, tma_subtile_stride_elems=4096, ) sCheckpoint_tma = SmemTile( @@ -1804,8 +1730,9 @@ def tmaldg_warp( desc_qwords = cutlass.Int32(TENSOR_MAP_QWORDS) while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) - + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) head_o = head_idx head_q = head_idx if cfg.q_ratio == 1 else head_idx // cutlass.Int32(cfg.q_ratio) head_k = head_idx if cfg.k_ratio == 1 else head_idx // cutlass.Int32(cfg.k_ratio) @@ -1823,21 +1750,21 @@ def tmaldg_warp( tma_tensormap_acquire(desc_do_slot) tma_tensormap_acquire(desc_checkpoint_slot) - for rev_idx in cutlass.range(cend - wstart): - chunk_idx = cend - 1 - rev_idx + for rev_idx in cutlass.range(compute_end - write_start): + chunk_idx = compute_end - 1 - rev_idx tok_coord = chunk_idx * cutlass.Int32(cfg.b_t) - # ---- K load ---------------------------------------------------------- + # ---- K load -------------------------------------------------------------- k_idx = k_index.idx bars.mb_k_mma_done[k_idx].wait(k_index.phase) - bars.mb_k_cg0_done[k_idx].wait(k_index.phase) + bars.mb_k_cg2_done[k_idx].wait(k_index.phase) k_index = advance(k_index, cfg.smem_k_stages) if elect_one: bars.mb_k_ready[k_idx].arrive(n_bytes=cfg.tma_k_bytes) k_slice = tma_slice_runtime_desc(desc_k_slot, cutlass.Int32(0), head_k, tok_coord) tma_load_tile(sK_tma[k_idx], k_slice, bars.mb_k_ready[k_idx].smem_ptr, acquire=False) - # ---- Q load ---------------------------------------------------------- + # ---- Q load -------------------------------------------------------------- q_idx = q_index.idx bars.mb_q_mma_done[q_idx].wait(q_index.phase) bars.mb_q_cg1_done[q_idx].wait(q_index.phase) @@ -1847,7 +1774,7 @@ def tmaldg_warp( q_slice = tma_slice_runtime_desc(desc_q_slot, cutlass.Int32(0), head_q, tok_coord) tma_load_tile(sQ_tma[q_idx], q_slice, bars.mb_q_ready[q_idx].smem_ptr, acquire=False) - # ---- V load ---------------------------------------------------------- + # ---- V load -------------------------------------------------------------- v_idx = v_index.idx bars.mb_v_mma_done[v_idx].wait(v_index.phase) v_index = advance(v_index, cfg.smem_v_stages) @@ -1856,7 +1783,7 @@ def tmaldg_warp( v_slice = tma_slice_runtime_desc(desc_v_slot, cutlass.Int32(0), head_v, tok_coord) tma_load_tile(sV_tma[v_idx], v_slice, bars.mb_v_ready[v_idx].smem_ptr, acquire=False) - # ---- dO load --------------------------------------------------------- + # ---- dO load ------------------------------------------------------------- do_idx = do_index.idx bars.mb_do_mma_done[do_idx].wait(do_index.phase) do_index = advance(do_index, cfg.smem_do_stages) @@ -1865,7 +1792,7 @@ def tmaldg_warp( do_slice = tma_slice_runtime_desc(desc_do_slot, cutlass.Int32(0), head_o, tok_coord) tma_load_tile(sdO_tma[do_idx], do_slice, bars.mb_do_ready[do_idx].smem_ptr, acquire=False) - # ---- entering state ---------- + # ---- entering state ------------------------------------------------------ if chunk_idx >= FIRST_STATE_CHUNK: state_idx = state_index.idx bars.mb_state_mma_done[state_idx].wait(state_index.phase) @@ -1875,7 +1802,10 @@ def tmaldg_warp( checkpoint_slice = tma_slice_runtime_desc(desc_checkpoint_slot, cutlass.Int32(0), cutlass.Int32(0), chunk_idx, head_o) tma_load_tile(sCheckpoint_tma[state_idx], checkpoint_slice, bars.mb_state_ready[state_idx].smem_ptr, acquire=False) - tile_idx, sched_state = sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ctas) + next_tile, scheduler_state = scheduler_publish_next( + cfg, bars, sScheduler, mScheduler, scheduler_state, tile_idx, num_ctas, tail_base, tail_row, elect_one + ) + tile_idx = next_tile for _ in range(cfg.smem_q_stages): bars.mb_q_mma_done[q_index.idx].wait(q_index.phase) @@ -1883,7 +1813,7 @@ def tmaldg_warp( q_index = advance(q_index, cfg.smem_q_stages) for _ in range(cfg.smem_k_stages): bars.mb_k_mma_done[k_index.idx].wait(k_index.phase) - bars.mb_k_cg0_done[k_index.idx].wait(k_index.phase) + bars.mb_k_cg2_done[k_index.idx].wait(k_index.phase) k_index = advance(k_index, cfg.smem_k_stages) for _ in range(cfg.smem_v_stages): bars.mb_v_mma_done[v_index.idx].wait(v_index.phase) @@ -1912,18 +1842,17 @@ def compute0_warp_group( sA, sDa, sDm, - sK, - sdQ, - sDstate, sstate_flat, sdstate_flat, - sSched, + sScheduler, bars, ): """Compute warp-group 0 role (warps 0-3): persistent scheduler loop building each chunk's blockwise-inverse T and attention matrices.""" nvvm.setmaxregister(cfg.num_regs_compute_group_0, nvvm.SetMaxRegisterAction.INCREASE) + elect_one = nvvm.elect_sync() + gate_index = PipelineState.start(phase=0) beta_index = PipelineState.start(phase=0) a_index = PipelineState.start(phase=1) @@ -1931,31 +1860,39 @@ def compute0_warp_group( da_acc_index = PipelineState.start(phase=0) dm_ready_index = PipelineState.start(phase=0) cg0_dbeta_index = PipelineState.start(phase=0) - - nvvm.barrier_cta_sync_aligned( - cfg.tmem_alloc_barrier_id, - thread_count=cfg.tmem_alloc_barrier_threads, - ) - tmem_base = tmem_base_slot.load() + cg0_kk_ready = PipelineState.start(phase=0) + cg0_a_ready = PipelineState.start(phase=0) + cg0_dstate_smem_ready = PipelineState.start(phase=0) + cg0_state_ready_index = PipelineState.start(phase=0) num_threads_cg0 = cfg.threads_per_warp * len(cfg.compute_group_0_warp_ids) cg0_tidx = tidx % num_threads_cg0 warp_id = cg0_tidx // cfg.threads_per_warp - lane_id = cg0_tidx % cfg.threads_per_warp + lane_idx = cg0_tidx % cfg.threads_per_warp + store_row = warp_id * 16 + lane_idx % 16 + store_col = (lane_idx // 16) * 8 + bpe = cfg.io_dtype.width // 8 num_vals = 32 FRAG_COLS = 16 ACC_N_FRAGS = cfg.b_t // FRAG_COLS - store_row = warp_id * 16 + lane_id % 16 - store_col = (lane_id // 16) * 8 - tmem_warp_row = warp_id * cfg.threads_per_warp - tmem_shared_acc_col = tmem_base + cfg.tmem_shared_acc_offset - tmem_shared_inp_col = tmem_base + cfg.tmem_shared_inp_offset SHARED_INP_STAGE_COLS = cfg.b_t // 2 - tmem_do_prime_col = tmem_shared_inp_col - tmem_du_col = tmem_shared_inp_col + SHARED_INP_STAGE_COLS - tmem_dyp_col = tmem_du_col ACC_STAGE_COLS = cfg.b_t + acc_zero = cfg.acc_dtype(0.0) + mask_zero = opaque_f32_zero() + DSTATE_IN0 = 1 if cfg.use_dstate_in else 0 + + nvvm.barrier_cta_sync_aligned(cfg.tmem_lifecycle_barrier_id, thread_count=cfg.tmem_user_threads) + tmem_base = tmem_base_slot.load() + tmem_col = tmem_base & 0xFFFF + tmem_row = tmem_base >> 16 + row_lo_addr = tmem_row << 16 + row_hi_addr = (tmem_row + 16) << 16 + tmem_shared_acc_col = tmem_col + cfg.tmem_shared_acc_offset + tmem_shared_input_col = tmem_col + cfg.tmem_shared_input_offset + tmem_do_prime_col = tmem_shared_input_col + tmem_du_col = tmem_shared_input_col + SHARED_INP_STAGE_COLS + tmem_dyp_col = tmem_du_col tmem_acc_a = tmem_shared_acc_col tmem_acc_b = tmem_shared_acc_col + ACC_STAGE_COLS tmem_kk_col = tmem_acc_a @@ -1965,66 +1902,45 @@ def compute0_warp_group( tmem_a_col = tmem_acc_b tmem_u_col = tmem_acc_b tmem_da_col = tmem_acc_b - tmem_dk_state_path_col = tmem_shared_inp_col - acc_zero = cfg.acc_dtype(0.0) - mask_zero = opaque_f32_zero() - frag_row = cg0_tidx % 8 + (cg0_tidx // 16 % 2) * 8 - frag_col = (cg0_tidx // 8 % 2) * 8 + (cg0_tidx // 32 % 2) * 32 - frag_slab_off = (cg0_tidx // 64) * 4096 - sK_base_p = cute.make_ptr(cfg.io_dtype, sK[0].base, mem_space=cute.AddressSpace.smem, assumed_align=cfg.buffer_align_bytes) - k_stage_elems_cg0 = cfg.k_cosize // cfg.smem_k_stages - sdQ_base = cute.make_ptr(cfg.io_dtype, sdQ[0].base, mem_space=cute.AddressSpace.smem, assumed_align=cfg.buffer_align_bytes) - sdH_parts_p = cute.make_ptr(cfg.io_dtype, sDstate[0].base, mem_space=cute.AddressSpace.smem, assumed_align=cfg.buffer_align_bytes) - # dGate fold scratch in sKK - skk_red = cute.make_ptr(cutlass.Float32, sKK[0].base, mem_space=cute.AddressSpace.smem, assumed_align=cfg.buffer_align_bytes) - cg0_k_index = PipelineState.start(phase=0) - cg0_dgate_index = PipelineState.start(phase=0) - tmem_dvdk_acc_col = tmem_base + cfg.tmem_dvdk_acc_offset - tmem_dstate_inp_col = tmem_base + cfg.tmem_dstate_inp_offset - cg0_kk_ready = PipelineState.start(phase=0) - cg0_a_ready = PipelineState.start(phase=0) - cg0_dk_scale_ready = PipelineState.start(phase=0) - cg0_dq_scale_ready = PipelineState.start(phase=0) - cg0_dstate_smem_ready = PipelineState.start(phase=0) - cg0_dk_attn_ready = PipelineState.start(phase=0) - DSTATE_IN0 = 1 if cfg.use_dstate_in else 0 + tmem_dk_state_path_col = tmem_shared_input_col + tmem_dstate_input_col = tmem_col + cfg.tmem_dstate_input_offset - tinv_zero_ptr = cute.make_ptr(cutlass.Int32, sTinv[0].base, mem_space=cute.AddressSpace.smem, assumed_align=cfg.buffer_align_bytes) + skk_red = cute.make_ptr(cutlass.Float32, sKK[0].base.toint(), mem_space=cute.AddressSpace.smem, assumed_align=cfg.buffer_align_bytes) + tinv_zero_ptr = cute.make_ptr(cutlass.Int32, sTinv[0].base.toint(), mem_space=cute.AddressSpace.smem, assumed_align=cfg.buffer_align_bytes) for z in cutlass.range_constexpr(cfg.t_inv_cosize * bpe // 4 // num_threads_cg0): (tinv_zero_ptr + cg0_tidx + z * num_threads_cg0).store(cutlass.Int32(0)) - sched_state = PipelineState.start(phase=0) + scheduler_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) FIRST_STATE_CHUNK = 0 if cfg.use_initial_state else 1 while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) - num_item_chunks = cend - wstart + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) + num_item_chunks = compute_end - write_start for chunk_idx in cutlass.range(num_item_chunks): - # ---- T-pairwise ------------------------------------------------------ + # ---- T-pairwise ---------------------------------------------------------- gate_idx = gate_index.idx bars.mb_gate_ready[gate_idx].wait(gate_index.phase) gate_index = advance(gate_index, cfg.smem_gate_stages) - row_cs = [] + row_cumsums = [] for r in cutlass.range_constexpr(2): - row_cs.append(sCumsumlog[warp_id * 16 + lane_id // 4 + r * 8, 0, gate_idx]) - col_cs = [] + row_cumsums.append(sCumsumlog[warp_id * 16 + lane_idx // 4 + r * 8, 0, gate_idx]) + col_cumsums = [] for g in cutlass.range_constexpr(8): for b in cutlass.range_constexpr(2): - col_cs.append(sCumsumlog[(lane_id % 4) * 2 + g * 8 + b, 0, gate_idx]) + col_cumsums.append(sCumsumlog[(lane_idx % 4) * 2 + g * 8 + b, 0, gate_idx]) decay_t = [] decay_t_strict = [] for k in cutlass.range_constexpr(num_vals): - crow = warp_id * 16 + lane_id // 4 + ((k // 2) % 2) * 8 - ccol = (lane_id % 4) * 2 + ((k // 4) * 8 + k % 2) - decay_t.append(cute.math.exp2(row_cs[(k // 2) % 2] - col_cs[(k // 4) * 2 + (k % 2)], fastmath=True) if crow >= ccol else mask_zero) - decay_t_strict.append(mask_zero if crow == ccol else decay_t[k]) - last_cs = sCumsumlog[cfg.b_t - 1, 0, gate_idx] - decay_scale_fp32 = [] - for i in cutlass.range_constexpr(16): - decay_scale_fp32.append(cute.math.exp2(last_cs - col_cs[i], fastmath=True)) - decay_scale_vals = [decay_scale_fp32[(k // 4) * 2 + (k % 2)] for k in range(num_vals)] + chunk_row = warp_id * 16 + lane_idx // 4 + ((k // 2) % 2) * 8 + chunk_col = (lane_idx % 4) * 2 + ((k // 4) * 8 + k % 2) + decay_t.append( + cute.math.exp2(row_cumsums[(k // 2) % 2] - col_cumsums[(k // 4) * 2 + (k % 2)], fastmath=True) if chunk_row >= chunk_col else mask_zero + ) + decay_t_strict.append(mask_zero if chunk_row == chunk_col else decay_t[k]) cumprod_total = sCumprod[sCumprod.shape[0] - 1, 0, gate_idx] beta_idx = beta_index.idx @@ -2033,10 +1949,10 @@ def compute0_warp_group( gBeta = [] for k in cutlass.range_constexpr(num_vals): - crow = warp_id * 16 + lane_id // 4 + ((k // 2) % 2) * 8 - gBeta.append(sBeta[crow, 0, beta_idx]) + chunk_row = warp_id * 16 + lane_idx // 4 + ((k // 2) % 2) * 8 + gBeta.append(sBeta[chunk_row, 0, beta_idx]) - # ---- KK epi: M_kk[i,j] = W_kk[i,j] * T_strict[i,j] * Beta[i] -------- + # ---- KK epilogue: M kk[i,j] = W kk[i,j] * T strict[i,j] * Beta[i] -------- tinv_idx = tinv_index.idx tinv_index = advance(tinv_index, cfg.smem_t_inv_stages) bars.mb_kk_acc_ready[0].wait(cg0_kk_ready.phase) @@ -2044,7 +1960,7 @@ def compute0_warp_group( tinv_base = sTinv[tinv_idx].base kk_base = sKK[tinv_idx].base - kk_vec = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_kk_col, cutlass.Float32), num=8) + kk_vec = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_lo_addr + tmem_kk_col, cutlass.Float32), num=8) kk_pack = [] for k in cutlass.range_constexpr(num_vals // 2): p0, p1 = fmul2(kk_vec[2 * k], kk_vec[2 * k + 1], decay_t_strict[2 * k], decay_t_strict[2 * k + 1]) @@ -2052,18 +1968,14 @@ def compute0_warp_group( kk_pack.append(fp32_to_fp16(v0, v1, dtype=cfg.io_dtype)) for c in cutlass.range_constexpr(ACC_N_FRAGS): nvvm.stmatrix( - cutlass.inttoptr( - kk_base + (store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS)) * bpe, - cutlass.AddressSpace.smem, - cutlass.BFloat16, - ), + kk_base + store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS), [kk_pack[c * 4 + 0], kk_pack[c * 4 + 1], kk_pack[c * 4 + 2], kk_pack[c * 4 + 3]], nvvm.MMALayout.ROW, ) - if chunk_idx < cend - FIRST_STATE_CHUNK: + if chunk_idx < compute_end - FIRST_STATE_CHUNK: bars.mb_kk_acc_done[0].arrive() - # ---- A epi: A[i,j] = W_qk[i,j] * T[i,j] * scale --------------------- + # ---- A epilogue: A[i,j] = W qk[i,j] * T[i,j] * scale --------------------- a_idx = a_index.idx a_phase = a_index.phase a_index = advance(a_index, cfg.smem_a_stages) @@ -2071,7 +1983,7 @@ def compute0_warp_group( cg0_a_ready = advance(cg0_a_ready, 1) a_base = sA[a_idx].base - a_vec = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_a_col, cutlass.Float32), num=8) + a_vec = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_lo_addr + tmem_a_col, cutlass.Float32), num=8) a_pack = [] for k in cutlass.range_constexpr(num_vals // 2): p0, p1 = fmul2(a_vec[2 * k], a_vec[2 * k + 1], decay_t[2 * k], decay_t[2 * k + 1]) @@ -2080,18 +1992,14 @@ def compute0_warp_group( bars.mb_a_done[a_idx].wait(a_phase) for c in cutlass.range_constexpr(ACC_N_FRAGS): nvvm.stmatrix( - cutlass.inttoptr( - a_base + (store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS)) * bpe, - cutlass.AddressSpace.smem, - cutlass.BFloat16, - ), + a_base + store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS), [a_pack[c * 4 + 0], a_pack[c * 4 + 1], a_pack[c * 4 + 2], a_pack[c * 4 + 3]], nvvm.MMALayout.ROW, ) nvvm.fence_proxy("async.shared", space="cta") bars.mb_a_ready[a_idx].arrive() - # ---- blockwise inverse: T_inv = ------------------------------------- + # ---- blockwise inverse --------------------------------------------------- nvvm.barrier_cta_sync_aligned( cfg.inverse_barrier_id, thread_count=cfg.inverse_barrier_threads, @@ -2103,39 +2011,35 @@ def compute0_warp_group( thread_count=cfg.inverse_barrier_threads, ) - blockwise_diagonal_8x8_to_16x16(cfg, tinv_base, kk_base, warp_id * 16, lane_id) + blockwise_diagonal_8x8_to_16x16(cfg, tinv_base, kk_base, warp_id * 16, lane_idx) nvvm.barrier_cta_sync_aligned( cfg.inverse_barrier_id, thread_count=cfg.inverse_barrier_threads, ) if warp_id < 2: - blockwise_diagonal_16x16_to_32x32(cfg, tinv_base, kk_base, warp_id * 32, lane_id) + blockwise_diagonal_16x16_to_32x32(cfg, tinv_base, kk_base, warp_id * 32, lane_idx) nvvm.barrier_cta_sync_aligned( cfg.inverse_barrier_id, thread_count=cfg.inverse_barrier_threads, ) if warp_id < 2: - blockwise_diagonal_32x32_to_64x64(cfg, tinv_base, kk_base, warp_id, lane_id) + blockwise_diagonal_32x32_to_64x64(cfg, tinv_base, kk_base, warp_id, lane_idx) nvvm.barrier_cta_sync_aligned( cfg.inverse_barrier_id, thread_count=cfg.inverse_barrier_threads, ) - # ---- Beta column-scaling in place: T_inv[i,j] *= Beta[j] ------------ + # ---- Beta column-scaling in place: T^-1[i,j] *= Beta[j] ------------------ beta_col = [] for k in cutlass.range_constexpr(num_vals): - beta_col.append(sBeta[(lane_id % 4) * 2 + ((k // 4) * 8 + k % 2), 0, beta_idx]) + beta_col.append(sBeta[(lane_idx % 4) * 2 + ((k // 4) * 8 + k % 2), 0, beta_idx]) tinv_frags = [] for c in cutlass.range_constexpr(ACC_N_FRAGS): tinv_frags += list( nvvm.ldmatrix( - cutlass.inttoptr( - tinv_base + (store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS)) * bpe, - cutlass.AddressSpace.smem, - cutlass.BFloat16, - ), + tinv_base + store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS), 4, nvvm.MMALayout.ROW, ) @@ -2147,11 +2051,7 @@ def compute0_warp_group( tinv_pack.append(fp32_to_fp16(s0, s1, dtype=cfg.io_dtype)) for c in cutlass.range_constexpr(ACC_N_FRAGS): nvvm.stmatrix( - cutlass.inttoptr( - tinv_base + (store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS)) * bpe, - cutlass.AddressSpace.smem, - cutlass.BFloat16, - ), + tinv_base + store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS), [tinv_pack[c * 4 + 0], tinv_pack[c * 4 + 1], tinv_pack[c * 4 + 2], tinv_pack[c * 4 + 3]], nvvm.MMALayout.ROW, ) @@ -2159,34 +2059,13 @@ def compute0_warp_group( nvvm.fence_proxy("async.shared", space="cta") bars.mb_t_inv_ready[tinv_idx].arrive() - # ---- dQ inter rescale ------------------------------------------------ - if chunk_idx < cend - FIRST_STATE_CHUNK: - cumprod_fp32 = [] - for g in cutlass.range_constexpr(8): - for b in cutlass.range_constexpr(2): - cumprod_fp32.append(sCumprod[(lane_id % 4) * 2 + g * 8 + b, 0, gate_idx]) - cumprod_vals = [cumprod_fp32[(k // 4) * 2 + (k % 2)] for k in range(num_vals)] - bars.mb_dq_acc_scale_ready[0].wait(cg0_dq_scale_ready.phase) - cg0_dq_scale_ready = advance(cg0_dq_scale_ready, 1) - dqi_ptrs = [] - dqi_vecs = [] - for sub in cutlass.range_constexpr(2): - dqi_ptrs.append(nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_dstate_inp_col, cutlass.Float32)) - dqi_vecs.append(nvvm.tcgen05_ld("16x256b", dqi_ptrs[sub], num=8)) - for sub in cutlass.range_constexpr(2): - dqi_scaled = [] - for j in cutlass.range_constexpr(16): - p0, p1 = fmul2(dqi_vecs[sub][2 * j], dqi_vecs[sub][2 * j + 1], cumprod_vals[2 * j], cumprod_vals[2 * j + 1]) - s0, s1 = fmul2(p0, p1, scale, scale) - dqi_scaled += [s0, s1] - nvvm.tcgen05_st("16x256b", dqi_ptrs[sub], cutlass.Vector.from_elements(tuple(dqi_scaled), cutlass.Float32)) - nvvm.tcgen05_wait("store") - bars.mb_dq_acc_scale_done[0].arrive() - - # ---- dGate_last state dot dstate term ---------------------------------------- + # ---- dGate last state dot dstate term ------------------------------------ if chunk_idx + DSTATE_IN0 >= 1: bars.mb_dstate_smem_ready[0].wait(cg0_dstate_smem_ready.phase) cg0_dstate_smem_ready = advance(cg0_dstate_smem_ready, 1) + if chunk_idx < compute_end - FIRST_STATE_CHUNK: + bars.mb_state_ready[cg0_state_ready_index.idx].wait(cg0_state_ready_index.phase) + cg0_state_ready_index = advance(cg0_state_ready_index, cfg.smem_state_stages) sdstate_base = sdstate_flat.iterator.toint() sstate_base = sstate_flat.iterator.toint() state_dot_dstate_lo = [opaque_f32_zero(), opaque_f32_zero(), opaque_f32_zero(), opaque_f32_zero()] @@ -2212,50 +2091,15 @@ def compute0_warp_group( ) for off in [1, 2, 4, 8, 16]: state_dot_dstate += nvvm.shfl_sync(0xFFFFFFFF, state_dot_dstate, off, 31, kind=nvvm.Shfl.BFLY) + hdh_last_w = cumprod_total * state_dot_dstate if chunk_idx + DSTATE_IN0 >= 1 and chunk_idx < compute_end - FIRST_STATE_CHUNK else acc_zero bars.mb_state_dot_dstate_done[0].arrive() - # ---- dK inter rescale ------------------------------------------------ - cg0_k_idx = cg0_k_index.idx - cg0_k_index = advance(cg0_k_index, cfg.smem_k_stages) - k_dot_dk_inter = cutlass.Float32(0.0) - if chunk_idx + DSTATE_IN0 >= 1: - bars.mb_dk_scale_acc_ready[0].wait(cg0_dk_scale_ready.phase) - cg0_dk_scale_ready = advance(cg0_dk_scale_ready, 1) - k_dot_dk_inter_lo = [opaque_f32_zero(), opaque_f32_zero(), opaque_f32_zero(), opaque_f32_zero()] - k_dot_dk_inter_hi = [opaque_f32_zero(), opaque_f32_zero(), opaque_f32_zero(), opaque_f32_zero()] - dki_ptrs = [] - dki_vecs = [] - for sub in cutlass.range_constexpr(2): - dki_ptrs.append(nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_dvdk_acc_col, cutlass.Float32)) - dki_vecs.append(nvvm.tcgen05_ld("16x256b", dki_ptrs[sub], num=8)) - for sub in cutlass.range_constexpr(2): - dki_scaled = [] - for j in cutlass.range_constexpr(16): - s0, s1 = fmul2(dki_vecs[sub][2 * j], dki_vecs[sub][2 * j + 1], decay_scale_vals[2 * j], decay_scale_vals[2 * j + 1]) - dki_scaled += [s0, s1] - nvvm.tcgen05_st("16x256b", dki_ptrs[sub], cutlass.Vector.from_elements(tuple(dki_scaled), cutlass.Float32)) - for m0 in cutlass.range_constexpr(4): - frag_addr = frag_slab_off + (frag_row + m0 * 16) * 64 + swizzle_xor_128b(frag_row + m0 * 16, frag_col + sub * 16) - k_frag = nvvm.ldmatrix((sK_base_p + cg0_k_idx * k_stage_elems_cg0 + frag_addr).raw_ptr(), 4, nvvm.MMALayout.COL) - for i in cutlass.range_constexpr(4): - k_lo, k_hi = f16x2_to_f32(k_frag[i], dtype=cfg.io_dtype) - k_dot_dk_inter_lo[i], k_dot_dk_inter_hi[i] = ffma2( - dki_scaled[8 * m0 + 2 * i], dki_scaled[8 * m0 + 2 * i + 1], k_lo, k_hi, k_dot_dk_inter_lo[i], k_dot_dk_inter_hi[i] - ) - nvvm.tcgen05_wait("store") - bars.mb_dk_scale_acc_done[0].arrive() - k_dot_dk_inter = ((k_dot_dk_inter_lo[0] + k_dot_dk_inter_lo[1]) + (k_dot_dk_inter_lo[2] + k_dot_dk_inter_lo[3])) + ( - (k_dot_dk_inter_hi[0] + k_dot_dk_inter_hi[1]) + (k_dot_dk_inter_hi[2] + k_dot_dk_inter_hi[3]) - ) - for off in [1, 2, 4, 8, 16]: - k_dot_dk_inter += nvvm.shfl_sync(0xFFFFFFFF, k_dot_dk_inter, off, 31, kind=nvvm.Shfl.BFLY) - - # ---- dA epilogue ----------------------------------------------------- + # ---- dA epilogue --------------------------------------------------------- bars.mb_da_acc_ready[0].wait(da_acc_index.phase) da_acc_index = advance(da_acc_index, 1) da_base = sDa[0].base - da_vec = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_da_col, cutlass.Float32), num=8) + da_vec = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_lo_addr + tmem_da_col, cutlass.Float32), num=8) da_pack = [] for k in cutlass.range_constexpr(num_vals // 2): p0, p1 = fmul2(da_vec[2 * k], da_vec[2 * k + 1], decay_t[2 * k], decay_t[2 * k + 1]) @@ -2263,22 +2107,18 @@ def compute0_warp_group( da_pack.append(fp32_to_fp16(v0, v1, dtype=cfg.io_dtype)) for c in cutlass.range_constexpr(ACC_N_FRAGS): nvvm.stmatrix( - cutlass.inttoptr( - da_base + (store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS)) * bpe, - cutlass.AddressSpace.smem, - cutlass.BFloat16, - ), + da_base + store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS), [da_pack[c * 4 + 0], da_pack[c * 4 + 1], da_pack[c * 4 + 2], da_pack[c * 4 + 3]], nvvm.MMALayout.ROW, ) nvvm.fence_proxy("async.shared", space="cta") bars.mb_da_ready[0].arrive() - # ---- dM epilogue ----------------------------------------------------- + # ---- dM epilogue --------------------------------------------------------- bars.mb_dm_acc_ready[0].wait(dm_ready_index.phase) dm_ready_index = advance(dm_ready_index, 1) dm_base = sDm[0].base - dm_vec = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_dm_core_col, cutlass.Float32), num=8) + dm_vec = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_lo_addr + tmem_dm_core_col, cutlass.Float32), num=8) dm_pack = [] for k in cutlass.range_constexpr(num_vals // 2): @@ -2288,41 +2128,19 @@ def compute0_warp_group( dm_pack.append(fp32_to_fp16(v0, v1, dtype=cfg.io_dtype)) for c in cutlass.range_constexpr(ACC_N_FRAGS): nvvm.stmatrix( - cutlass.inttoptr( - dm_base + (store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS)) * bpe, - cutlass.AddressSpace.smem, - cutlass.BFloat16, - ), + dm_base + store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS), [dm_pack[c * 4 + 0], dm_pack[c * 4 + 1], dm_pack[c * 4 + 2], dm_pack[c * 4 + 3]], nvvm.MMALayout.ROW, ) nvvm.fence_proxy("async.shared", space="cta") bars.mb_dm_acc_done[0].arrive() - # ---- dK attn read ---------------------------------------------------- - bars.mb_dk_attn_acc_ready[0].wait(cg0_dk_attn_ready.phase) - cg0_dk_attn_ready = advance(cg0_dk_attn_ready, 1) - dks_regs = [] - for sub in cutlass.range_constexpr(2): - dks_vec = nvvm.tcgen05_ld( - "16x256b", - nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_dvdk_acc_col, cutlass.Float32), - num=8, - ) - dks_regs.append([dks_vec[k] for k in range(32)]) - nvvm.tcgen05_wait("load") - bars.mb_dk_attn_acc_done[0].arrive() - - # ---- dBeta/dGate M-terms: E = dM_core ⊙ M_kk(sKK, strict-masked). ---- + # ---- dBeta/dGate M-terms: E = dM core ⊙ M kk(sKK, strict-masked) --------- kk_frag = [] for c in cutlass.range_constexpr(ACC_N_FRAGS): kk_frag += list( nvvm.ldmatrix( - cutlass.inttoptr( - kk_base + (store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS)) * bpe, - cutlass.AddressSpace.smem, - cutlass.BFloat16, - ), + kk_base + store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS), 4, nvvm.MMALayout.ROW, ) @@ -2331,16 +2149,23 @@ def compute0_warp_group( cute.math.rcp(gBeta[0] + cutlass.Float32(1e-10), approx=True, ftz=True), cute.math.rcp(gBeta[2] + cutlass.Float32(1e-10), approx=True, ftz=True), ] - row_acc = [cutlass.Float32(0.0)] * 8 - col_part = [opaque_f32_zero() for _ in range(16)] + row_acc = [None] * 8 + col_part = [None] * 16 for j in cutlass.range_constexpr(num_vals // 2): - klo, khi = f16x2_to_f32(kk_frag[j], dtype=cfg.io_dtype) + k_lo, k_hi = f16x2_to_f32(kk_frag[j], dtype=cfg.io_dtype) binv_j = binv_row[j % 2] - p_lo, p_hi = fmul2(dm_vec[2 * j], dm_vec[2 * j + 1], klo, khi) - e_lo, e_hi = fmul2(p_lo, p_hi, binv_j, binv_j) - row_acc[(j % 2) * 4 + (j // 2) % 4] += e_lo + e_hi + lo, hi = fmul2(dm_vec[2 * j], dm_vec[2 * j + 1], k_lo, k_hi) + lo, hi = fmul2(lo, hi, binv_j, binv_j) + r = cutlass.const_expr((j % 2) * 4 + (j // 2) % 4) + if cutlass.const_expr(j < 8): + row_acc[r] = lo + hi + else: + row_acc[r] += lo + hi c0 = cutlass.const_expr((j // 2) * 2) - col_part[c0], col_part[c0 + 1] = fadd2(col_part[c0], col_part[c0 + 1], e_lo, e_hi) + if cutlass.const_expr(j % 2 == 0): + col_part[c0], col_part[c0 + 1] = lo, hi + else: + col_part[c0], col_part[c0 + 1] = fadd2(col_part[c0], col_part[c0 + 1], lo, hi) row_part = [ (row_acc[0] + row_acc[1]) + (row_acc[2] + row_acc[3]), (row_acc[4] + row_acc[5]) + (row_acc[6] + row_acc[7]), @@ -2348,100 +2173,38 @@ def compute0_warp_group( for off in [1, 2]: for rp in cutlass.range_constexpr(2): row_part[rp] += nvvm.shfl_sync(0xFFFFFFFF, row_part[rp], off, 31, kind=nvvm.Shfl.BFLY) + # ---- dGate column fold --------------------------------------------------- + nvvm.barrier_cta_sync_aligned(cfg.inverse_barrier_id, thread_count=cfg.inverse_barrier_threads) + col_lo, col_hi = warp_reduce_scatter_sum_16_elements(col_part, lane_idx) + tok0 = (lane_idx // 4) * 8 + (lane_idx % 4) * 2 + (skk_red + warp_id * 64 + tok0).store(col_lo) + (skk_red + warp_id * 64 + tok0 + 1).store(col_hi + hdh_last_w if lane_idx == 31 else col_hi) + + if lane_idx % 4 == 0: + for rp in cutlass.range_constexpr(2): + crow_r = warp_id * 16 + lane_idx // 4 + rp * 8 + sCumsumlog[crow_r, 0, gate_idx] = acc_zero - row_part[rp] + nvvm.barrier_cta_sync_aligned(cfg.inverse_barrier_id, thread_count=cfg.inverse_barrier_threads) + if cg0_tidx < 64: + col_sum = (skk_red + cg0_tidx).load() + (skk_red + 64 + cg0_tidx).load() + (skk_red + 128 + cg0_tidx).load() + (skk_red + 192 + cg0_tidx).load() + sCumsumlog[cg0_tidx, 0, gate_idx] = sCumsumlog[cg0_tidx, 0, gate_idx] + col_sum + bars.mb_dgate_cg0_ready[0].arrive() + bars.mb_gate_done[gate_idx].arrive() + + # ---- dBeta M-term -------------------------------------------------------- bars.mb_dbeta_cg1_ready[0].wait(cg0_dbeta_index.phase) cg0_dbeta_index = advance(cg0_dbeta_index, 1) - if lane_id % 4 == 0: + if lane_idx % 4 == 0: for rp in cutlass.range_constexpr(2): - crow_r = warp_id * 16 + lane_id // 4 + rp * 8 + crow_r = warp_id * 16 + lane_idx // 4 + rp * 8 db = sBeta[crow_r, 0, beta_idx] - row_part[rp] * binv_row[rp] if cutlass.const_expr(cfg.beta_sigmoid): b = gBeta[2 * rp] db = db * (b - b * b) sBeta[crow_r, 0, beta_idx] = db - # ---- part reductions ------------------------------------------------- - if chunk_idx + DSTATE_IN0 >= FIRST_STATE_CHUNK: - part_k = [acc_zero] * 16 - for sub in cutlass.range_constexpr(2): - for m0 in cutlass.range_constexpr(4): - frag_addr = frag_slab_off + (frag_row + m0 * 16) * 64 + swizzle_xor_128b(frag_row + m0 * 16, frag_col + sub * 16) - k_frag = nvvm.ldmatrix((sK_base_p + cg0_k_idx * k_stage_elems_cg0 + frag_addr).raw_ptr(), 4, nvvm.MMALayout.COL) - for i in cutlass.range_constexpr(4): - k_lo, k_hi = f16x2_to_f32(k_frag[i], dtype=cfg.io_dtype) - frag_e0 = cutlass.const_expr(8 * m0 + 2 * i) - part_e0 = cutlass.const_expr((frag_e0 // 4) * 2 + (frag_e0 % 2)) - if cutlass.const_expr(sub == 0 and i % 2 == 0): - part_k[part_e0], part_k[part_e0 + 1] = fmul2(dks_regs[sub][frag_e0], dks_regs[sub][frag_e0 + 1], k_lo, k_hi) - else: - part_k[part_e0], part_k[part_e0 + 1] = ffma2( - dks_regs[sub][frag_e0], dks_regs[sub][frag_e0 + 1], k_lo, k_hi, part_k[part_e0], part_k[part_e0 + 1] - ) - nvvm.fence_proxy("async.shared", space="cta") - bars.mb_k_cg0_done[cg0_k_idx].arrive() - - dgate_part = [col_part[j] - part_k[j] for j in range(16)] - dgate_part_lo, dgate_part_hi = warp_reduce_scatter_frag_16_elems(dgate_part, lane_id) - dgate_last_w = k_dot_dk_inter + ( - (cumprod_total * state_dot_dstate if chunk_idx + DSTATE_IN0 >= 1 else acc_zero) if chunk_idx < cend - FIRST_STATE_CHUNK else acc_zero - ) - tok0 = (lane_id // 4) * 8 + (lane_id % 4) * 2 - (skk_red + warp_id * 64 + tok0).store(dgate_part_lo) - (skk_red + warp_id * 64 + tok0 + 1).store(dgate_part_hi + dgate_last_w if lane_id == 31 else dgate_part_hi) - - bars.mb_dgate_cg1_ready[0].wait(cg0_dgate_index.phase) - cg0_dgate_index = advance(cg0_dgate_index, 1) - if lane_id % 4 == 0: - for rp in cutlass.range_constexpr(2): - crow_r = warp_id * 16 + lane_id // 4 + rp * 8 - sCumsumlog[crow_r, 0, gate_idx] = sCumsumlog[crow_r, 0, gate_idx] - row_part[rp] - nvvm.barrier_cta_sync_aligned(cfg.inverse_barrier_id, thread_count=cfg.inverse_barrier_threads) - if cg0_tidx < 64: - dgate_sum = ( - (skk_red + cg0_tidx).load() + (skk_red + 64 + cg0_tidx).load() + (skk_red + 128 + cg0_tidx).load() + (skk_red + 192 + cg0_tidx).load() - ) - sCumsumlog[cg0_tidx, 0, gate_idx] = sCumsumlog[cg0_tidx, 0, gate_idx] + dgate_sum - - if chunk_idx + DSTATE_IN0 < FIRST_STATE_CHUNK: - part_k = [acc_zero] * 16 - for sub in cutlass.range_constexpr(2): - for m0 in cutlass.range_constexpr(4): - frag_addr = frag_slab_off + (frag_row + m0 * 16) * 64 + swizzle_xor_128b(frag_row + m0 * 16, frag_col + sub * 16) - k_frag = nvvm.ldmatrix((sK_base_p + cg0_k_idx * k_stage_elems_cg0 + frag_addr).raw_ptr(), 4, nvvm.MMALayout.COL) - for i in cutlass.range_constexpr(4): - k_lo, k_hi = f16x2_to_f32(k_frag[i], dtype=cfg.io_dtype) - frag_e0 = cutlass.const_expr(8 * m0 + 2 * i) - part_e0 = cutlass.const_expr((frag_e0 // 4) * 2 + (frag_e0 % 2)) - if cutlass.const_expr(sub == 0 and i % 2 == 0): - part_k[part_e0], part_k[part_e0 + 1] = fmul2(dks_regs[sub][frag_e0], dks_regs[sub][frag_e0 + 1], k_lo, k_hi) - else: - part_k[part_e0], part_k[part_e0 + 1] = ffma2( - dks_regs[sub][frag_e0], dks_regs[sub][frag_e0 + 1], k_lo, k_hi, part_k[part_e0], part_k[part_e0 + 1] - ) - nvvm.fence_proxy("async.shared", space="cta") - bars.mb_k_cg0_done[cg0_k_idx].arrive() - - dgate_part = [col_part[j] - part_k[j] for j in range(16)] - dgate_part_lo, dgate_part_hi = warp_reduce_scatter_frag_16_elems(dgate_part, lane_id) - tok0 = (lane_id // 4) * 8 + (lane_id % 4) * 2 - (skk_red + warp_id * 64 + tok0).store(dgate_part_lo) - (skk_red + warp_id * 64 + tok0 + 1).store(dgate_part_hi) - - bars.mb_dgate_cg1_ready[0].wait(cg0_dgate_index.phase) - cg0_dgate_index = advance(cg0_dgate_index, 1) - if lane_id % 4 == 0: - for rp in cutlass.range_constexpr(2): - crow_r = warp_id * 16 + lane_id // 4 + rp * 8 - sCumsumlog[crow_r, 0, gate_idx] = sCumsumlog[crow_r, 0, gate_idx] - row_part[rp] - nvvm.barrier_cta_sync_aligned(cfg.inverse_barrier_id, thread_count=cfg.inverse_barrier_threads) - if cg0_tidx < 64: - dgate_sum = ( - (skk_red + cg0_tidx).load() + (skk_red + 64 + cg0_tidx).load() + (skk_red + 128 + cg0_tidx).load() + (skk_red + 192 + cg0_tidx).load() - ) - sCumsumlog[cg0_tidx, 0, gate_idx] = sCumsumlog[cg0_tidx, 0, gate_idx] + dgate_sum - - bars.mb_gate_done[gate_idx].arrive() bars.mb_beta_done[beta_idx].arrive() - tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one) for _ in range(cfg.smem_a_stages): bars.mb_a_done[a_index.idx].wait(a_index.phase) a_index = advance(a_index, cfg.smem_a_stages) @@ -2457,77 +2220,87 @@ def compute1_warp_group( num_ctas, cu_seqlens, mWorkItems, - mDstate0_out, - mDstate_in, tidx, warp_idx, tmem_base_slot, scale, sQ, - sK, sV, sdO, sCumsumlog, sCumprod, sBeta, sdQ, - sdK, sdV, - sDstate, + sdK, sDa, sDm, - sSched, + sScheduler, bars, ): """Compute warp-group 1 role (warps 4-7): persistent scheduler loop running each chunk's gradient epilogues and stagings.""" + nvvm.setmaxregister(cfg.num_regs_compute_group_1, nvvm.SetMaxRegisterAction.INCREASE) + elect_one = nvvm.elect_sync() + v_index = PipelineState.start(phase=0) do_index = PipelineState.start(phase=0) gate_index = PipelineState.start(phase=0) cg1_k_state_ready = PipelineState.start(phase=0) cg1_u_ready = PipelineState.start(phase=0) + cg1_dq_scale_ready = PipelineState.start(phase=0) + cg1_dstate_smem_index = PipelineState.start(phase=0) + cg1_dstate_scale_index = PipelineState.start(phase=0) cg1_dy_ready = PipelineState.start(phase=0) cg1_du_scale_ready = PipelineState.start(phase=0) cg1_du_total_ready = PipelineState.start(phase=0) - cg1_dk_total_ready = PipelineState.start(phase=0) - dstate_acc_index = PipelineState.start(phase=0) - cg1_state_dot_dstate_index = PipelineState.start(phase=0) + cg1_dgate_cg0_index = PipelineState.start(phase=0) + cg1_cg2_chunk_index = PipelineState.start(phase=0) dq_index = PipelineState.start(phase=1) cg1_beta_index = PipelineState.start(phase=0) sdv_done_index = PipelineState.start(phase=1) - cg1_dk_state_path_ready = PipelineState.start(phase=0) dq_total_ready_index = PipelineState.start(phase=0) - dstate_inp_index = PipelineState.start(phase=1) - dk_index = PipelineState.start(phase=1) dv_index = PipelineState.start(phase=1) - - nvvm.setmaxregister(cfg.num_regs_compute_group_1, nvvm.SetMaxRegisterAction.INCREASE) - nvvm.barrier_cta_sync_aligned( - cfg.tmem_alloc_barrier_id, - thread_count=cfg.tmem_alloc_barrier_threads, - ) - tmem_base = tmem_base_slot.load() + scheduler_state = PipelineState.start(phase=0) num_threads_cg1 = cfg.threads_per_warp * len(cfg.compute_group_1_warp_ids) cg1_tidx = tidx % num_threads_cg1 - lane_id = cg1_tidx % cfg.threads_per_warp - tmem_warp_row = (cg1_tidx // cfg.threads_per_warp) * cfg.threads_per_warp - ldtm_width = 32 - sttm_width = ldtm_width // 2 - num_state_subs = cutlass.const_expr(cfg.d_v // ldtm_width) - tmem_dstate_acc_col = tmem_base + cfg.tmem_dstate_acc_offset - tmem_dstate_inp_col = tmem_base + cfg.tmem_dstate_inp_offset - tmem_dvdk_acc_col = tmem_base + cfg.tmem_dvdk_acc_offset - tmem_shared_acc_col = tmem_base + cfg.tmem_shared_acc_offset - tmem_shared_inp_col = tmem_base + cfg.tmem_shared_inp_offset + lane_idx = cg1_tidx % cfg.threads_per_warp + SHARED_INP_STAGE_COLS = cfg.b_t // 2 - tmem_do_prime_col = tmem_shared_inp_col - tmem_du_col = tmem_shared_inp_col + SHARED_INP_STAGE_COLS + ACC_STAGE_COLS = cfg.b_t + frag_row = cg1_tidx % 8 + (cg1_tidx // 16 % 2) * 8 + frag_col = (cg1_tidx // 8 % 2) * 8 + (cg1_tidx // 32 % 2) * 32 + frag_segment = (cg1_tidx // 64) * 4096 + dv_stage_elements = cfg.dv_cosize // cfg.smem_dv_stages + sdV_base = sdV[0].base + v_stage_elements = cfg.v_cosize // cfg.smem_v_stages + sV_base = sV[0].base + sQ_base = sQ[0].base + do_stage_elements = cfg.do_cosize // cfg.smem_do_stages + sdO_base = sdO[0].base + dq_stage_elements = cfg.dq_cosize // cfg.smem_dq_stages + sdQ_base = sdQ[0].base + cg1_warp_id = cg1_tidx // cfg.threads_per_warp + tile_idx = cutlass.Int32(bidx) + FIRST_STATE_CHUNK = 0 if cfg.use_initial_state else 1 + + nvvm.barrier_cta_sync_aligned(cfg.tmem_lifecycle_barrier_id, thread_count=cfg.tmem_user_threads) + tmem_base = tmem_base_slot.load() + tmem_col = tmem_base & 0xFFFF + tmem_row = tmem_base >> 16 + row_lo_addr = tmem_row << 16 + row_hi_addr = (tmem_row + 16) << 16 + tmem_dstate_input_col = tmem_col + cfg.tmem_dstate_input_offset + tmem_dvdk_acc_col = tmem_col + cfg.tmem_dvdk_acc_offset + tmem_shared_acc_col = tmem_col + cfg.tmem_shared_acc_offset + tmem_shared_input_col = tmem_col + cfg.tmem_shared_input_offset + tmem_do_prime_col = tmem_shared_input_col + tmem_du_col = tmem_shared_input_col + SHARED_INP_STAGE_COLS tmem_dyp_col = tmem_du_col - tmem_y_col = tmem_base + cfg.tmem_y_offset + tmem_y_col = tmem_col + cfg.tmem_y_offset tmem_g_k_state_col = tmem_y_col + SHARED_INP_STAGE_COLS - ACC_STAGE_COLS = cfg.b_t tmem_acc_a = tmem_shared_acc_col tmem_acc_b = tmem_shared_acc_col + ACC_STAGE_COLS tmem_kk_col = tmem_acc_a @@ -2537,83 +2310,23 @@ def compute1_warp_group( tmem_a_col = tmem_acc_b tmem_u_col = tmem_acc_b tmem_da_col = tmem_acc_b - tmem_dk_state_path_col = tmem_shared_inp_col - frag_row = cg1_tidx % 8 + (cg1_tidx // 16 % 2) * 8 - frag_col = (cg1_tidx // 8 % 2) * 8 + (cg1_tidx // 32 % 2) * 32 - frag_slab_off = (cg1_tidx // 64) * 4096 - dv_stage_elems = cfg.dv_cosize // cfg.smem_dv_stages - sdV_base = cute.make_ptr(cfg.io_dtype, sdV[0].base, mem_space=cute.AddressSpace.smem, assumed_align=cfg.buffer_align_bytes) - v_stage_elems = cfg.v_cosize // cfg.smem_v_stages - sV_base = cute.make_ptr(cfg.io_dtype, sV[0].base, mem_space=cute.AddressSpace.smem, assumed_align=cfg.buffer_align_bytes) - sQ_base = cute.make_ptr(cfg.io_dtype, sQ[0].base, mem_space=cute.AddressSpace.smem, assumed_align=cfg.buffer_align_bytes) - do_stage_elems = cfg.do_cosize // cfg.smem_do_stages - sdO_base = cute.make_ptr(cfg.io_dtype, sdO[0].base, mem_space=cute.AddressSpace.smem, assumed_align=cfg.buffer_align_bytes) - dk_stage_elems = cfg.dk_cosize // cfg.smem_dk_stages - dq_stage_elems = cfg.dq_cosize // cfg.smem_dq_stages - sdQ_base = cute.make_ptr(cfg.io_dtype, sdQ[0].base, mem_space=cute.AddressSpace.smem, assumed_align=cfg.buffer_align_bytes) - sdK_base = cute.make_ptr(cfg.io_dtype, sdK[0].base, mem_space=cute.AddressSpace.smem, assumed_align=cfg.buffer_align_bytes) - sDstate_base_int = sDstate[0].base - sred_base = cute.make_ptr(cutlass.Float32, sdQ[0].base, mem_space=cute.AddressSpace.smem, assumed_align=cfg.buffer_align_bytes) - sdstate_red = cute.make_ptr(cutlass.Float32, sDstate[0].base, mem_space=cute.AddressSpace.smem, assumed_align=cfg.buffer_align_bytes) - dstate_done_idx = cutlass.Int32(0) - cg1_warp_id = cg1_tidx // cfg.threads_per_warp - sched_state = PipelineState.start(phase=0) - tile_idx = cutlass.Int32(bidx) - FIRST_STATE_CHUNK = 0 if cfg.use_initial_state else 1 + sred_base = cute.make_ptr(cutlass.Float32, sdQ[0].base.toint(), mem_space=cute.AddressSpace.smem, assumed_align=cfg.buffer_align_bytes) + sdk_red = cute.make_ptr(cutlass.Float32, sdK[0].base.toint(), mem_space=cute.AddressSpace.smem, assumed_align=cfg.buffer_align_bytes) + PART_Q_SCRATCH_OFF = 4 * 64 + tmem_dstate_acc_col = tmem_col + cfg.tmem_dstate_acc_offset + ldtm_width = 32 + num_ldtms = cutlass.const_expr(cfg.d_v // ldtm_width) + cg1_dk_index = PipelineState.start(phase=1) while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) - num_item_chunks = cend - wstart - - # ---- d_final_state prologue ---------------------------------------------- - if cutlass.const_expr(cfg.use_dstate_in): - if num_item_chunks > 0: - gDstate_in = mDstate_in[None, None, head_idx, batch_idx] - seed_from_dstate_in = cend == num_chunks_b - dstate_inp_idx = dstate_inp_index.idx - bars.mb_dstate_inp_done[dstate_inp_idx].wait(dstate_inp_index.phase) - dstate_inp_index = advance(dstate_inp_index, cfg.tmem_dstate_inp_stages) - for sub in cutlass.range_constexpr(num_state_subs): - dstate_in_vals = [] - for kk in cutlass.range_constexpr(ldtm_width): - v = gDstate_in[cg1_tidx, sub * ldtm_width + kk] - v = v if seed_from_dstate_in else cutlass.Float32(0.0) - dstate_in_vals.append(v) - nvvm.tcgen05_st( - "32x32b", - nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_dstate_acc_col + sub * ldtm_width, cutlass.Float32), - cutlass.Vector.from_elements(tuple(dstate_in_vals), cutlass.Float32), - ) - dstate_in_pack = [fp32_to_fp16(dstate_in_vals[2 * j], dstate_in_vals[2 * j + 1], dtype=cfg.io_dtype) for j in range(16)] - nvvm.tcgen05_st( - "32x32b", - nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_dstate_inp_col + sub * sttm_width, cutlass.Int32), - cutlass.Vector.from_elements(tuple(dstate_in_pack), cutlass.Int32), - ) - nvvm.tcgen05_wait("store") - bars.mb_dstate_inp_ready[dstate_inp_idx].arrive() - - for sub in cutlass.range_constexpr(num_state_subs): - dstate_smem_vec = nvvm.tcgen05_ld( - "32x32b", - nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_dstate_acc_col + sub * ldtm_width, cutlass.Float32), - num=32, - ) - for g in cutlass.range_constexpr(ldtm_width // 8): - dstate_smem_pack = tuple( - fp32_to_fp16(dstate_smem_vec[g * 8 + 2 * t], dstate_smem_vec[g * 8 + 2 * t + 1], dtype=cfg.io_dtype) for t in range(4) - ) - dstate_dk = sub * ldtm_width + g * 8 - dstate_addr = (dstate_dk // 64) * (cfg.d_v * 64) + cg1_tidx * 64 + swizzle_xor_128b(cg1_tidx, dstate_dk % 64) - cutlass.inttoptr(sDstate_base_int + dstate_addr * 2, cutlass.AddressSpace.smem, cfg.io_dtype).store( - cutlass.Vector.from_elements(dstate_smem_pack, cutlass.Int32).bitcast(cfg.io_dtype), alignment=16 - ) - nvvm.fence_proxy("async.shared", space="cta") - bars.mb_dstate_smem_ready[0].arrive() + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) + num_item_chunks = compute_end - write_start - # ---- chunks NT-1 .. 0 (backward) ------------------------------------------ + # ---- chunks NT-1 .. 0 (backward) --------------------------------------------- for rev_idx in cutlass.range(num_item_chunks): - chunk_idx = cend - 1 - rev_idx + chunk_idx = compute_end - 1 - rev_idx have_dstate = cutlass.Boolean(True) if cutlass.const_expr(cfg.use_dstate_in) else rev_idx > 0 gate_idx = gate_index.idx bars.mb_gate_ready[gate_idx].wait(gate_index.phase) @@ -2623,12 +2336,21 @@ def compute1_warp_group( bars.mb_beta_ready[beta_idx].wait(cg1_beta_index.phase) cg1_beta_index = advance(cg1_beta_index, cfg.smem_beta_stages) - # ---- dstate rescale: dstate *= this chunk's cumprod -------------------------- + num_vals = 32 + cumprod_fp32 = [] + for g in cutlass.range_constexpr(8): + for b in cutlass.range_constexpr(2): + cumprod_fp32.append(sCumprod[(lane_idx % 4) * 2 + g * 8 + b, 0, gate_idx]) + cumprod_vals = [cumprod_fp32[(k // 4) * 2 + (k % 2)] for k in range(num_vals)] + + # ---- dstate rescale: dstate *= this chunk's cumprod ---------------------- if have_dstate: + bars.mb_dstate_smem_ready[0].wait(cg1_dstate_smem_index.phase) + cg1_dstate_smem_index = advance(cg1_dstate_smem_index, 1) cumprod_top = sCumprod[sCumprod.shape[0] - 1, 0, gate_idx] - for sub in cutlass.range_constexpr(num_state_subs): + for i in cutlass.range_constexpr(num_ldtms): dstate_rescale_vec = nvvm.tcgen05_ld( - "32x32b", nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_dstate_acc_col + sub * ldtm_width, cutlass.Float32), num=32 + "32x32b", nvvm.make_tmem_ptr(row_lo_addr + tmem_dstate_acc_col + i * ldtm_width, cutlass.Float32), num=32 ) dstate_rescaled = [] for j in cutlass.range_constexpr(16): @@ -2636,107 +2358,67 @@ def compute1_warp_group( dstate_rescaled += [h0, h1] nvvm.tcgen05_st( "32x32b", - nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_dstate_acc_col + sub * ldtm_width, cutlass.Float32), + nvvm.make_tmem_ptr(row_lo_addr + tmem_dstate_acc_col + i * ldtm_width, cutlass.Float32), cutlass.Vector.from_elements(tuple(dstate_rescaled), cutlass.Float32), ) nvvm.tcgen05_wait("store") - bars.mb_dstate_scale_acc_done[dstate_done_idx].arrive() + bars.mb_dstate_scale_acc_done[cg1_dstate_scale_index.idx].arrive() + cg1_dstate_scale_index = advance(cg1_dstate_scale_index, cfg.tmem_dstate_acc_stages) - num_vals = 32 - cumprod_fp32 = [] - for g in cutlass.range_constexpr(8): - for b in cutlass.range_constexpr(2): - cumprod_fp32.append(sCumprod[(lane_id % 4) * 2 + g * 8 + b, 0, gate_idx]) - cumprod_vals = [cumprod_fp32[(k // 4) * 2 + (k % 2)] for k in range(num_vals)] - - # ---- dO' restage: dO * cumprod_vals * scale -> shared_inp TMEM ----------- - do_idx = do_index.idx - bars.mb_do_ready[do_idx].wait(do_index.phase) - do_index = advance(do_index, cfg.smem_do_stages) - do_regs = [[cutlass.Float32(0.0), cutlass.Float32(0.0)] for _ in range(32)] - for c in cutlass.range_constexpr(8): - m0 = cutlass.const_expr(c % 4) - sub = cutlass.const_expr(c // 4) - do_frag = nvvm.ldmatrix( - ( - sdO_base - + do_idx * do_stage_elems - + frag_slab_off - + (frag_row + m0 * 16) * 64 - + swizzle_xor_128b(frag_row + m0 * 16, frag_col + sub * 16) - ).raw_ptr(), - 4, - nvvm.MMALayout.COL, - ) - for i in cutlass.range_constexpr(4): - lo, hi = f16x2_to_f32(do_frag[i], dtype=cfg.io_dtype) - p0, p1 = fmul2(lo, hi, cumprod_vals[8 * m0 + 2 * i], cumprod_vals[8 * m0 + 2 * i + 1]) - do_regs[8 * m0 + 2 * i][sub], do_regs[8 * m0 + 2 * i + 1][sub] = fmul2(p0, p1, scale, scale) - for sub in cutlass.range_constexpr(2): - do_pack = [fp32_to_fp16(do_regs[2 * j][sub], do_regs[2 * j + 1][sub], dtype=cfg.io_dtype) for j in range(16)] - nvvm.tcgen05_st( - "16x128b", - nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_do_prime_col, cutlass.Int32), - cutlass.Vector.from_elements(tuple(do_pack), cutlass.Int32), - ) - nvvm.tcgen05_wait("store") - bars.mb_do_prime_inp_ready[0].arrive() - - # ---- dV inter: in-place decay scale ---------------------------------- - if have_dstate: - last_cumsumlog = sCumsumlog[cfg.b_t - 1, 0, gate_idx] - col_cs_fp32 = [] - for g in cutlass.range_constexpr(8): - for b in cutlass.range_constexpr(2): - col_cs_fp32.append(sCumsumlog[(lane_id % 4) * 2 + g * 8 + b, 0, gate_idx]) - decay_scale_fp32 = [] - for i in cutlass.range_constexpr(16): - decay_scale_fp32.append(cute.math.exp2(last_cumsumlog - col_cs_fp32[i], fastmath=True)) - decay_scale_vals = [decay_scale_fp32[(k // 4) * 2 + (k % 2)] for k in range(num_vals)] - bars.mb_du_scale_acc_ready[0].wait(cg1_du_scale_ready.phase) - cg1_du_scale_ready = advance(cg1_du_scale_ready, 1) - dv_ptrs = [] - dv_vecs = [] - for sub in cutlass.range_constexpr(2): - dv_ptrs.append(nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_dvdk_acc_col, cutlass.Float32)) - dv_vecs.append(nvvm.tcgen05_ld("16x256b", dv_ptrs[sub], num=8)) - for sub in cutlass.range_constexpr(2): - dv_scaled = [] - for j in cutlass.range_constexpr(16): - s0, s1 = fmul2(dv_vecs[sub][2 * j], dv_vecs[sub][2 * j + 1], decay_scale_vals[2 * j], decay_scale_vals[2 * j + 1]) - dv_scaled += [s0, s1] - nvvm.tcgen05_st("16x256b", dv_ptrs[sub], cutlass.Vector.from_elements(tuple(dv_scaled), cutlass.Float32)) - nvvm.tcgen05_wait("store") - bars.mb_du_scale_acc_done[0].arrive() - - # ---- Y staging: Y = V - cumprod*(K @ state) -> f16 TMEM slots ----------- + # ---- V fragment preload -------------------------------------------------- v_idx = v_index.idx bars.mb_v_ready[v_idx].wait(v_index.phase) v_index = advance(v_index, cfg.smem_v_stages) - if chunk_idx >= FIRST_STATE_CHUNK: - v_frags = [[cutlass.Int32(0), cutlass.Int32(0)] for _ in range(16)] - for c in cutlass.range_constexpr(8): - m0 = cutlass.const_expr(c % 4) - sub = cutlass.const_expr(c // 4) - v_frag = nvvm.ldmatrix( + v_frag = [] + for half in cutlass.range_constexpr(2): + v_words = [] + for block in cutlass.range_constexpr(4): + v_raw = nvvm.ldmatrix( ( sV_base - + v_idx * v_stage_elems - + frag_slab_off - + (frag_row + m0 * 16) * 64 - + swizzle_xor_128b(frag_row + m0 * 16, frag_col + sub * 16) - ).raw_ptr(), + + v_idx * v_stage_elements + + frag_segment + + (frag_row + block * 16) * 64 + + swizzle_xor_128b(frag_row + block * 16, frag_col + half * 16) + ), + 4, + nvvm.MMALayout.COL, + ) + for i in cutlass.range_constexpr(4): + v_words.append(v_raw[i]) + v_frag.append(v_words) + + # ---- dO fragment preload ------------------------------------------------- + do_idx = do_index.idx + bars.mb_do_ready[do_idx].wait(do_index.phase) + do_index = advance(do_index, cfg.smem_do_stages) + do_raw_frag = [] + for half in cutlass.range_constexpr(2): + do_words_raw = [] + for block in cutlass.range_constexpr(4): + do_raw = nvvm.ldmatrix( + ( + sdO_base + + do_idx * do_stage_elements + + frag_segment + + (frag_row + block * 16) * 64 + + swizzle_xor_128b(frag_row + block * 16, frag_col + half * 16) + ), 4, nvvm.MMALayout.COL, ) for i in cutlass.range_constexpr(4): - v_frags[4 * m0 + i][sub] = v_frag[i] + do_words_raw.append(do_raw[i]) + do_raw_frag.append(do_words_raw) + + # ---- Y stage: Y = V - cumprod * k state -> TMEM f16 ---------------------- + if chunk_idx >= FIRST_STATE_CHUNK: bars.mb_k_state_acc_ready[0].wait(cg1_k_state_ready.phase) cg1_k_state_ready = advance(cg1_k_state_ready, 1) - for sub in cutlass.range_constexpr(2): + for half in cutlass.range_constexpr(2): k_state_vec = nvvm.tcgen05_ld( "16x256b", - nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_k_state_col, cutlass.Float32), + nvvm.make_tmem_ptr(((tmem_row + half * 16) << 16) + tmem_k_state_col, cutlass.Float32), num=8, ) g_k_state_pack = [] @@ -2745,126 +2427,175 @@ def compute1_warp_group( g0, g1 = fmul2(k_state_vec[2 * j], k_state_vec[2 * j + 1], cumprod_vals[2 * j], cumprod_vals[2 * j + 1]) g_k_state_word = fp32_to_fp16(g0, g1, dtype=cfg.io_dtype) g_k_state_pack.append(g_k_state_word) - y_pack.append(sub_f16x2(v_frags[j][sub], g_k_state_word, cfg.io_dtype)) + y_pack.append(sub_f16x2(v_frag[half][j], g_k_state_word, cfg.io_dtype)) nvvm.tcgen05_st( "16x128b", - nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_y_col, cutlass.Int32), + nvvm.make_tmem_ptr(((tmem_row + half * 16) << 16) + tmem_y_col, cutlass.Int32), cutlass.Vector.from_elements(tuple(y_pack), cutlass.Int32), ) nvvm.tcgen05_st( "16x128b", - nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_g_k_state_col, cutlass.Int32), + nvvm.make_tmem_ptr(((tmem_row + half * 16) << 16) + tmem_g_k_state_col, cutlass.Int32), cutlass.Vector.from_elements(tuple(g_k_state_pack), cutlass.Int32), ) if chunk_idx < FIRST_STATE_CHUNK: - for sub in cutlass.range_constexpr(2): - v_pack = [] - for m0 in cutlass.range_constexpr(4): - v_frag = nvvm.ldmatrix( - ( - sV_base - + v_idx * v_stage_elems - + frag_slab_off - + (frag_row + m0 * 16) * 64 - + swizzle_xor_128b(frag_row + m0 * 16, frag_col + sub * 16) - ).raw_ptr(), - 4, - nvvm.MMALayout.COL, - ) - for i in cutlass.range_constexpr(4): - v_pack.append(v_frag[i]) + for half in cutlass.range_constexpr(2): nvvm.tcgen05_st( "16x128b", - nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_y_col, cutlass.Int32), - cutlass.Vector.from_elements(tuple(v_pack), cutlass.Int32), + nvvm.make_tmem_ptr(((tmem_row + half * 16) << 16) + tmem_y_col, cutlass.Int32), + cutlass.Vector.from_elements(tuple(v_frag[half]), cutlass.Int32), ) nvvm.tcgen05_wait("store") bars.mb_y_ready[0].arrive() - # ---- dU restage: dV acc ---------------------------------------------- + # ---- dV inter ------------------------------------------------------------ + if have_dstate: + last_cumsumlog = sCumsumlog[cfg.b_t - 1, 0, gate_idx] + col_cumsumlog_fp32 = [] + for g in cutlass.range_constexpr(8): + for b in cutlass.range_constexpr(2): + col_cumsumlog_fp32.append(sCumsumlog[(lane_idx % 4) * 2 + g * 8 + b, 0, gate_idx]) + decay_scale_fp32 = [] + for i in cutlass.range_constexpr(16): + decay_scale_fp32.append(cute.math.exp2(last_cumsumlog - col_cumsumlog_fp32[i], fastmath=True)) + decay_scale_vals = [decay_scale_fp32[(k // 4) * 2 + (k % 2)] for k in range(num_vals)] + bars.mb_du_scale_acc_ready[0].wait(cg1_du_scale_ready.phase) + cg1_du_scale_ready = advance(cg1_du_scale_ready, 1) + dv_ptrs = [] + dv_vecs = [] + for half in cutlass.range_constexpr(2): + dv_ptrs.append(nvvm.make_tmem_ptr(((tmem_row + half * 16) << 16) + tmem_dvdk_acc_col, cutlass.Float32)) + dv_vecs.append(nvvm.tcgen05_ld("16x256b", dv_ptrs[half], num=8)) + for half in cutlass.range_constexpr(2): + dv_scaled = [] + for j in cutlass.range_constexpr(16): + s0, s1 = fmul2(dv_vecs[half][2 * j], dv_vecs[half][2 * j + 1], decay_scale_vals[2 * j], decay_scale_vals[2 * j + 1]) + dv_scaled += [s0, s1] + nvvm.tcgen05_st("16x256b", dv_ptrs[half], cutlass.Vector.from_elements(tuple(dv_scaled), cutlass.Float32)) + nvvm.tcgen05_wait("store") + bars.mb_du_scale_acc_done[0].arrive() + + # ---- dO' stage: dO * cumprod * scale -> shared input TMEM ---------------- + bars.mb_cg2_chunk_ready[0].wait(cg1_cg2_chunk_index.phase) + cg1_cg2_chunk_index = advance(cg1_cg2_chunk_index, 1) + for half in cutlass.range_constexpr(2): + do_pack = [] + for j in cutlass.range_constexpr(16): + lo, hi = f16x2_to_f32(do_raw_frag[half][j], dtype=cfg.io_dtype) + p0, p1 = fmul2(lo, hi, cumprod_vals[2 * j], cumprod_vals[2 * j + 1]) + q0, q1 = fmul2(p0, p1, scale, scale) + do_pack.append(fp32_to_fp16(q0, q1, dtype=cfg.io_dtype)) + nvvm.tcgen05_st( + "16x128b", + nvvm.make_tmem_ptr(((tmem_row + half * 16) << 16) + tmem_do_prime_col, cutlass.Int32), + cutlass.Vector.from_elements(tuple(do_pack), cutlass.Int32), + ) + nvvm.tcgen05_wait("store") + bars.mb_do_prime_input_ready[0].arrive() + + # ---- dU stage: dV acc -> TMEM f16 A operand ------------------------------ bars.mb_du_total_acc_ready[0].wait(cg1_du_total_ready.phase) cg1_du_total_ready = advance(cg1_du_total_ready, 1) - for sub in cutlass.range_constexpr(2): - du_vec = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_dvdk_acc_col, cutlass.Float32), num=8) + for half in cutlass.range_constexpr(2): + du_vec = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(((tmem_row + half * 16) << 16) + tmem_dvdk_acc_col, cutlass.Float32), num=8) du_pack = [fp32_to_fp16(du_vec[2 * j], du_vec[2 * j + 1], dtype=cfg.io_dtype) for j in range(16)] nvvm.tcgen05_st( "16x128b", - nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_du_col, cutlass.Int32), + nvvm.make_tmem_ptr(((tmem_row + half * 16) << 16) + tmem_du_col, cutlass.Int32), cutlass.Vector.from_elements(tuple(du_pack), cutlass.Int32), ) nvvm.tcgen05_wait("store") - bars.mb_du_inp_ready[0].arrive() + bars.mb_du_input_ready[0].arrive() - # ---- U readout ------------------------------------------------------- + # ---- dQ inter rescale ---------------------------------------------------- + if chunk_idx >= FIRST_STATE_CHUNK: + bars.mb_dq_acc_scale_ready[0].wait(cg1_dq_scale_ready.phase) + cg1_dq_scale_ready = advance(cg1_dq_scale_ready, 1) + dqi_ptrs = [] + dqi_vecs = [] + for half in cutlass.range_constexpr(2): + dqi_ptrs.append(nvvm.make_tmem_ptr(((tmem_row + half * 16) << 16) + tmem_dstate_input_col, cutlass.Float32)) + dqi_vecs.append(nvvm.tcgen05_ld("16x256b", dqi_ptrs[half], num=8)) + for half in cutlass.range_constexpr(2): + dqi_scaled = [] + for j in cutlass.range_constexpr(16): + p0, p1 = fmul2(dqi_vecs[half][2 * j], dqi_vecs[half][2 * j + 1], cumprod_vals[2 * j], cumprod_vals[2 * j + 1]) + s0, s1 = fmul2(p0, p1, scale, scale) + dqi_scaled += [s0, s1] + nvvm.tcgen05_st("16x256b", dqi_ptrs[half], cutlass.Vector.from_elements(tuple(dqi_scaled), cutlass.Float32)) + nvvm.tcgen05_wait("store") + bars.mb_dq_acc_scale_done[0].arrive() + + # ---- U read: TMEM -> sU -------------------------------------------------- bars.mb_u_acc_ready[0].wait(cg1_u_ready.phase) cg1_u_ready = advance(cg1_u_ready, 1) u_regs = [] - for sub in cutlass.range_constexpr(2): - u_vec = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_u_col, cutlass.Float32), num=8) + for half in cutlass.range_constexpr(2): + u_vec = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(((tmem_row + half * 16) << 16) + tmem_u_col, cutlass.Float32), num=8) u_regs.append([u_vec[k] for k in range(32)]) - for sub in cutlass.range_constexpr(2): - for m0 in cutlass.range_constexpr(4): - u_pack = [fp32_to_fp16(u_regs[sub][8 * m0 + 2 * j], u_regs[sub][8 * m0 + 2 * j + 1], dtype=cfg.io_dtype) for j in range(4)] + for half in cutlass.range_constexpr(2): + for block in cutlass.range_constexpr(4): + u_pack = [fp32_to_fp16(u_regs[half][8 * block + 2 * j], u_regs[half][8 * block + 2 * j + 1], dtype=cfg.io_dtype) for j in range(4)] nvvm.stmatrix( ( sV_base - + v_idx * v_stage_elems - + frag_slab_off - + (frag_row + m0 * 16) * 64 - + swizzle_xor_128b(frag_row + m0 * 16, frag_col + sub * 16) - ).raw_ptr(), + + v_idx * v_stage_elements + + frag_segment + + (frag_row + block * 16) * 64 + + swizzle_xor_128b(frag_row + block * 16, frag_col + half * 16) + ), u_pack, nvvm.MMALayout.COL, ) nvvm.fence_proxy("async.shared", space="cta") bars.mb_u_ready[0].arrive() - # ---- dY -------------------------------------------------------------- + # ---- dY ------------------------------------------------------------------ bars.mb_dy_acc_ready[0].wait(cg1_dy_ready.phase) cg1_dy_ready = advance(cg1_dy_ready, 1) dy_regs = [] - for sub in cutlass.range_constexpr(2): + for half in cutlass.range_constexpr(2): dy_vec = nvvm.tcgen05_ld( "16x256b", - nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_dy_col, cutlass.Float32), + nvvm.make_tmem_ptr(((tmem_row + half * 16) << 16) + tmem_dy_col, cutlass.Float32), num=8, ) dy_regs.append([dy_vec[k] for k in range(32)]) - # ---- dY' = -cumprod_vals * dY -> f16 shared_inp -------------------------- + # ---- dY' = -cumprod * dY -> f16 shared input ----------------------------- neg_one = cutlass.Float32(-1.0) cumprod_neg_vals = [cumprod_vals[k] * neg_one for k in range(32)] - for sub in cutlass.range_constexpr(2): + for half in cutlass.range_constexpr(2): dyp = [] for j in cutlass.range_constexpr(16): - n0, n1 = fmul2(dy_regs[sub][2 * j], dy_regs[sub][2 * j + 1], cumprod_neg_vals[2 * j], cumprod_neg_vals[2 * j + 1]) + n0, n1 = fmul2(dy_regs[half][2 * j], dy_regs[half][2 * j + 1], cumprod_neg_vals[2 * j], cumprod_neg_vals[2 * j + 1]) dyp += [n0, n1] dyp_pack = [fp32_to_fp16(dyp[2 * j], dyp[2 * j + 1], dtype=cfg.io_dtype) for j in range(16)] nvvm.tcgen05_st( "16x128b", - nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_dyp_col, cutlass.Int32), + nvvm.make_tmem_ptr(((tmem_row + half * 16) << 16) + tmem_dyp_col, cutlass.Int32), cutlass.Vector.from_elements(tuple(dyp_pack), cutlass.Int32), ) nvvm.tcgen05_wait("store") - bars.mb_dyp_inp_ready[0].arrive() + bars.mb_dyp_input_ready[0].arrive() - # ---- dV staging: dV = dY -> the sdV slot ----------------------------- + # ---- dV stage: dV = dY -> sdV -------------------------------------------- dv_stg_idx = dv_index.idx bars.mb_dv_tmastg_done[dv_stg_idx].wait(dv_index.phase) dv_index = advance(dv_index, cfg.smem_dv_stages) bars.mb_sdv_done[0].wait(sdv_done_index.phase) sdv_done_index = advance(sdv_done_index, 1) - for sub in cutlass.range_constexpr(2): - for m0 in cutlass.range_constexpr(4): - dv_pack = [fp32_to_fp16(dy_regs[sub][8 * m0 + 2 * j], dy_regs[sub][8 * m0 + 2 * j + 1], dtype=cfg.io_dtype) for j in range(4)] + for half in cutlass.range_constexpr(2): + for block in cutlass.range_constexpr(4): + dv_pack = [fp32_to_fp16(dy_regs[half][8 * block + 2 * j], dy_regs[half][8 * block + 2 * j + 1], dtype=cfg.io_dtype) for j in range(4)] nvvm.stmatrix( ( sdV_base - + dv_stg_idx * dv_stage_elems - + frag_slab_off - + (frag_row + m0 * 16) * 64 - + swizzle_xor_128b(frag_row + m0 * 16, frag_col + sub * 16) - ).raw_ptr(), + + dv_stg_idx * dv_stage_elements + + frag_segment + + (frag_row + block * 16) * 64 + + swizzle_xor_128b(frag_row + block * 16, frag_col + half * 16) + ), dv_pack, nvvm.MMALayout.COL, ) @@ -2873,41 +2604,40 @@ def compute1_warp_group( dq_stg_idx = dq_index.idx bars.mb_dq_tmastg_done[dq_stg_idx].wait(dq_index.phase) dq_index = advance(dq_index, cfg.smem_dq_stages) - sred = sred_base + dq_stg_idx * (dq_stage_elems // 2) + sred = sred_base + dq_stg_idx * (dq_stage_elements // 2) - # ---- dBeta/dGate V-terms: dBeta_t += rowsum(dV ⊙ Y)_t / Beta_t ------- - part_y = [cutlass.Float32(0.0)] * 16 - for sub in cutlass.range_constexpr(2): - y_vec = nvvm.tcgen05_ld("16x128b", nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_y_col, cutlass.Int32), num=8) + # ---- dBeta/dGate V-terms: dBeta += rowsum(dV ⊙ Y) / Beta ----------------- + gsum = cutlass.Float32(0.0) + part_y = [None] * 16 + for half in cutlass.range_constexpr(2): + y_vec = nvvm.tcgen05_ld("16x128b", nvvm.make_tmem_ptr(((tmem_row + half * 16) << 16) + tmem_y_col, cutlass.Int32), num=8) for j in cutlass.range_constexpr(16): lo, hi = f16x2_to_f32(y_vec[j], dtype=cfg.io_dtype) frag_e0 = cutlass.const_expr(2 * j) part_e0 = cutlass.const_expr((frag_e0 // 4) * 2 + (frag_e0 % 2)) - if cutlass.const_expr(sub == 0 and j % 2 == 0): - part_y[part_e0], part_y[part_e0 + 1] = fmul2(dy_regs[sub][frag_e0], dy_regs[sub][frag_e0 + 1], lo, hi) + if cutlass.const_expr(half == 0 and j % 2 == 0): + part_y[part_e0], part_y[part_e0 + 1] = fmul2(dy_regs[half][frag_e0], dy_regs[half][frag_e0 + 1], lo, hi) else: part_y[part_e0], part_y[part_e0 + 1] = ffma2( - dy_regs[sub][frag_e0], dy_regs[sub][frag_e0 + 1], lo, hi, part_y[part_e0], part_y[part_e0 + 1] + dy_regs[half][frag_e0], dy_regs[half][frag_e0 + 1], lo, hi, part_y[part_e0], part_y[part_e0 + 1] ) - py_lo, py_hi = warp_reduce_scatter_frag_16_elems(part_y, lane_id) - vt_tok0 = (lane_id // 4) * 8 + (lane_id % 4) * 2 + py_lo, py_hi = warp_reduce_scatter_sum_16_elements(part_y, lane_idx) + vt_tok0 = (lane_idx // 4) * 8 + (lane_idx % 4) * 2 if chunk_idx >= FIRST_STATE_CHUNK: - part_g = [cutlass.Float32(0.0)] * 16 - for sub in cutlass.range_constexpr(2): - g_k_state_vec = nvvm.tcgen05_ld( - "16x128b", nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_g_k_state_col, cutlass.Int32), num=8 - ) + part_g = [None] * 16 + for half in cutlass.range_constexpr(2): + g_k_state_vec = nvvm.tcgen05_ld("16x128b", nvvm.make_tmem_ptr(((tmem_row + half * 16) << 16) + tmem_g_k_state_col, cutlass.Int32), num=8) for j in cutlass.range_constexpr(16): lo, hi = f16x2_to_f32(g_k_state_vec[j], dtype=cfg.io_dtype) frag_e0 = cutlass.const_expr(2 * j) part_e0 = cutlass.const_expr((frag_e0 // 4) * 2 + (frag_e0 % 2)) - if cutlass.const_expr(sub == 0 and j % 2 == 0): - part_g[part_e0], part_g[part_e0 + 1] = fmul2(dy_regs[sub][frag_e0], dy_regs[sub][frag_e0 + 1], lo, hi) + if cutlass.const_expr(half == 0 and j % 2 == 0): + part_g[part_e0], part_g[part_e0 + 1] = fmul2(dy_regs[half][frag_e0], dy_regs[half][frag_e0 + 1], lo, hi) else: part_g[part_e0], part_g[part_e0 + 1] = ffma2( - dy_regs[sub][frag_e0], dy_regs[sub][frag_e0 + 1], lo, hi, part_g[part_e0], part_g[part_e0 + 1] + dy_regs[half][frag_e0], dy_regs[half][frag_e0 + 1], lo, hi, part_g[part_e0], part_g[part_e0 + 1] ) - pg_lo, pg_hi = warp_reduce_scatter_frag_16_elems(part_g, lane_id) + pg_lo, pg_hi = warp_reduce_scatter_sum_16_elements(part_g, lane_idx) (sred + cg1_warp_id * 64 + vt_tok0).store(py_lo) (sred + cg1_warp_id * 64 + vt_tok0 + 1).store(py_hi) (sred + 256 + cg1_warp_id * 64 + vt_tok0).store(pg_lo) @@ -2918,7 +2648,6 @@ def compute1_warp_group( ysum = (sred + cg1_tidx).load() + (sred + 64 + cg1_tidx).load() + (sred + 128 + cg1_tidx).load() + (sred + 192 + cg1_tidx).load() gsum = (sred + 256 + cg1_tidx).load() + (sred + 320 + cg1_tidx).load() + (sred + 384 + cg1_tidx).load() + (sred + 448 + cg1_tidx).load() sBeta[cg1_tidx, 0, beta_idx] = ysum * binv_t - sCumsumlog[cg1_tidx, 0, gate_idx] = cutlass.Float32(0.0) - gsum if chunk_idx < FIRST_STATE_CHUNK: (sred + cg1_warp_id * 64 + vt_tok0).store(py_lo) (sred + cg1_warp_id * 64 + vt_tok0 + 1).store(py_hi) @@ -2927,250 +2656,468 @@ def compute1_warp_group( binv_t = cute.math.rcp(sBeta[cg1_tidx, 0, beta_idx] + cutlass.Float32(1e-10), approx=True, ftz=True) ysum = (sred + cg1_tidx).load() + (sred + 64 + cg1_tidx).load() + (sred + 128 + cg1_tidx).load() + (sred + 192 + cg1_tidx).load() sBeta[cg1_tidx, 0, beta_idx] = ysum * binv_t - sCumsumlog[cg1_tidx, 0, gate_idx] = cutlass.Float32(0.0) bars.mb_beta_done[beta_idx].arrive() bars.mb_dbeta_cg1_ready[0].arrive() nvvm.barrier_cta_sync_aligned(cfg.cg1_barrier_id, thread_count=cfg.cg1_barrier_threads) - # ---- Q fragments held in registers ----------------------------------- + # ---- Q fragments held in registers --------------------------------------- q_frag = [] - for sub in cutlass.range_constexpr(2): + for half in cutlass.range_constexpr(2): q_words = [] - for m0 in cutlass.range_constexpr(4): - q_f16 = nvvm.ldmatrix( - (sQ_base + frag_slab_off + (frag_row + m0 * 16) * 64 + swizzle_xor_128b(frag_row + m0 * 16, frag_col + sub * 16)).raw_ptr(), + for block in cutlass.range_constexpr(4): + q_raw = nvvm.ldmatrix( + sQ_base + frag_segment + (frag_row + block * 16) * 64 + swizzle_xor_128b(frag_row + block * 16, frag_col + half * 16), 4, nvvm.MMALayout.COL, ) for i in cutlass.range_constexpr(4): - q_words.append(q_f16[i]) + q_words.append(q_raw[i]) q_frag.append(q_words) nvvm.tcgen05_st( "16x128b", - nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_y_col, cutlass.Int32), + nvvm.make_tmem_ptr(((tmem_row + half * 16) << 16) + tmem_y_col, cutlass.Int32), cutlass.Vector.from_elements(tuple(q_words), cutlass.Int32), ) nvvm.tcgen05_wait("store") bars.mb_q_cg1_done[0].arrive() - # ---- dQ final read -> sdQ -------------------------------------------- + # ---- dQ final read -> sdQ ------------------------------------------------ bars.mb_dq_acc_total_ready[0].wait(dq_total_ready_index.phase) dq_total_ready_index = advance(dq_total_ready_index, 1) dq_regs = [] - for sub in cutlass.range_constexpr(2): - dq_vec = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_dstate_inp_col, cutlass.Float32), num=8) + for half in cutlass.range_constexpr(2): + dq_vec = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(((tmem_row + half * 16) << 16) + tmem_dstate_input_col, cutlass.Float32), num=8) dq_regs.append([dq_vec[k] for k in range(32)]) - for m0 in cutlass.range_constexpr(4): - frag_addr = frag_slab_off + (frag_row + m0 * 16) * 64 + swizzle_xor_128b(frag_row + m0 * 16, frag_col + sub * 16) - dq_pack = [fp32_to_fp16(dq_vec[8 * m0 + 2 * j], dq_vec[8 * m0 + 2 * j + 1], dtype=cfg.io_dtype) for j in range(4)] - nvvm.stmatrix((sdQ_base + dq_stg_idx * dq_stage_elems + frag_addr).raw_ptr(), dq_pack, nvvm.MMALayout.COL) - nvvm.fence_proxy("async.shared", space="cta") + nvvm.tcgen05_wait("load") bars.mb_dq_acc_total_done[0].arrive() + for half in cutlass.range_constexpr(2): + for block in cutlass.range_constexpr(4): + frag_addr = frag_segment + (frag_row + block * 16) * 64 + swizzle_xor_128b(frag_row + block * 16, frag_col + half * 16) + dq_pack = [fp32_to_fp16(dq_regs[half][8 * block + 2 * j], dq_regs[half][8 * block + 2 * j + 1], dtype=cfg.io_dtype) for j in range(4)] + nvvm.stmatrix(sdQ_base + dq_stg_idx * dq_stage_elements + frag_addr, dq_pack, nvvm.MMALayout.COL) + nvvm.fence_proxy("async.shared", space="cta") bars.mb_dq_tmastg_ready[dq_stg_idx].arrive() - # ---- dQ dot (part_q) ------------------------------------------------- - part_q = [cutlass.Float32(0.0)] * 16 - for sub in cutlass.range_constexpr(2): - for m0 in cutlass.range_constexpr(4): + # ---- dQ dot (part q) ----------------------------------------------------- + part_q = [None] * 16 + for half in cutlass.range_constexpr(2): + for block in cutlass.range_constexpr(4): for i in cutlass.range_constexpr(4): - q_lo, q_hi = f16x2_to_f32(q_frag[sub][4 * m0 + i], dtype=cfg.io_dtype) - frag_e0 = cutlass.const_expr(8 * m0 + 2 * i) + q_lo, q_hi = f16x2_to_f32(q_frag[half][4 * block + i], dtype=cfg.io_dtype) + frag_e0 = cutlass.const_expr(8 * block + 2 * i) part_e0 = cutlass.const_expr((frag_e0 // 4) * 2 + (frag_e0 % 2)) - if cutlass.const_expr(sub == 0 and i % 2 == 0): - part_q[part_e0], part_q[part_e0 + 1] = fmul2(dq_regs[sub][frag_e0], dq_regs[sub][frag_e0 + 1], q_lo, q_hi) + if cutlass.const_expr(half == 0 and i % 2 == 0): + part_q[part_e0], part_q[part_e0 + 1] = fmul2(dq_regs[half][frag_e0], dq_regs[half][frag_e0 + 1], q_lo, q_hi) else: part_q[part_e0], part_q[part_e0 + 1] = ffma2( - dq_regs[sub][frag_e0], dq_regs[sub][frag_e0 + 1], q_lo, q_hi, part_q[part_e0], part_q[part_e0 + 1] + dq_regs[half][frag_e0], dq_regs[half][frag_e0 + 1], q_lo, q_hi, part_q[part_e0], part_q[part_e0 + 1] ) - part_q_lo, part_q_hi = warp_reduce_scatter_frag_16_elems(part_q, lane_id) - tok0 = (lane_id // 4) * 8 + (lane_id % 4) * 2 - (sdstate_red + (cg1_tidx // 32) * 64 + tok0).store(part_q_lo) - (sdstate_red + (cg1_tidx // 32) * 64 + tok0 + 1).store(part_q_hi) + bars.mb_dk_tmastg_done[cg1_dk_index.idx].wait(cg1_dk_index.phase) + cg1_dk_index = advance(cg1_dk_index, cfg.smem_dk_stages) + part_q_lo, part_q_hi = warp_reduce_scatter_sum_16_elements(part_q, lane_idx) + tok0 = (lane_idx // 4) * 8 + (lane_idx % 4) * 2 + (sdk_red + PART_Q_SCRATCH_OFF + (cg1_tidx // 32) * 64 + tok0).store(part_q_lo) + (sdk_red + PART_Q_SCRATCH_OFF + (cg1_tidx // 32) * 64 + tok0 + 1).store(part_q_hi) nvvm.barrier_cta_sync_aligned(cfg.cg1_barrier_id, thread_count=cfg.cg1_barrier_threads) + bars.mb_dgate_cg0_ready[0].wait(cg1_dgate_cg0_index.phase) + cg1_dgate_cg0_index = advance(cg1_dgate_cg0_index, 1) if cg1_tidx < 64: pq_sum = ( - (sdstate_red + cg1_tidx).load() - + (sdstate_red + 64 + cg1_tidx).load() - + (sdstate_red + 128 + cg1_tidx).load() - + (sdstate_red + 192 + cg1_tidx).load() + (sdk_red + PART_Q_SCRATCH_OFF + cg1_tidx).load() + + (sdk_red + PART_Q_SCRATCH_OFF + 64 + cg1_tidx).load() + + (sdk_red + PART_Q_SCRATCH_OFF + 128 + cg1_tidx).load() + + (sdk_red + PART_Q_SCRATCH_OFF + 192 + cg1_tidx).load() ) - sCumsumlog[cg1_tidx, 0, gate_idx] = sCumsumlog[cg1_tidx, 0, gate_idx] + pq_sum + sCumsumlog[cg1_tidx, 0, gate_idx] = sCumsumlog[cg1_tidx, 0, gate_idx] + pq_sum - gsum bars.mb_gate_done[gate_idx].arrive() bars.mb_dgate_cg1_ready[0].arrive() + tile_idx, scheduler_state = scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one) + + bars.mb_tmem_done[0].arrive() + + for _ in range(cfg.smem_dv_stages): + bars.mb_dv_tmastg_done[dv_index.idx].wait(dv_index.phase) + dv_index = advance(dv_index, cfg.smem_dv_stages) + + +@cute.jit +def compute2_warp_group( + cfg, + total_tiles, + bidx, + num_ctas, + mWorkItems, + mDstate0_out, + mDstate_in, + tidx, + tmem_base_slot, + sCumsumlog, + sCumprod, + sK, + sdK, + sDstate_trans, + sScheduler, + bars, +): + """Compute warp-group 2 role (warps 8-11): the whole read side of the + dV/dK accumulator slot. One chunk walks the slot's last three + productions in place -- inter rescale, attn read, dM-terms fold -- so the + dGate K part and the dK store come off the same registers.""" + + nvvm.setmaxregister(cfg.num_regs_compute_group_2, nvvm.SetMaxRegisterAction.DECREASE) + elect_one = nvvm.elect_sync() + + gate_index = PipelineState.start(phase=0) + cg2_k_index = PipelineState.start(phase=0) + cg2_kk_scale_ready = PipelineState.start(phase=0) + cg2_dk_attn_ready = PipelineState.start(phase=0) + cg2_dgate_index = PipelineState.start(phase=0) + cg2_dk_state_path_ready = PipelineState.start(phase=0) + cg2_dk_total_ready = PipelineState.start(phase=0) + dstate_acc_index = PipelineState.start(phase=0) + dstate_input_index = PipelineState.start(phase=1) + cg2_state_dot_dstate_index = PipelineState.start(phase=0) + cg2_dq_total_index = PipelineState.start(phase=0) + dk_index = PipelineState.start(phase=1) + scheduler_state = PipelineState.start(phase=0) + + num_threads_cg2 = cfg.threads_per_warp * len(cfg.compute_group_2_warp_ids) + cg2_tidx = tidx % num_threads_cg2 + warp_id = cg2_tidx // cfg.threads_per_warp + lane_idx = cg2_tidx % cfg.threads_per_warp + + num_vals = 32 + acc_zero = cfg.acc_dtype(0.0) + neg_one = cutlass.Float32(-1.0) + frag_row = cg2_tidx % 8 + (cg2_tidx // 16 % 2) * 8 + frag_col = (cg2_tidx // 8 % 2) * 8 + (cg2_tidx // 32 % 2) * 32 + frag_segment = (cg2_tidx // 64) * 4096 + sK_base = sK[0].base + k_stage_elements_cg2 = cfg.k_cosize // cfg.smem_k_stages + sdK_base = sdK[0].base + dk_stage_elements = cfg.dk_cosize // cfg.smem_dk_stages + ldtm_width = 32 + sttm_width = 16 + num_ldtms = cutlass.const_expr(cfg.d_v // ldtm_width) + dstate_done_idx = cutlass.Int32(0) + sDstate_base = sDstate_trans[0].base + DSTATE_IN0 = 1 if cfg.use_dstate_in else 0 + tile_idx = cutlass.Int32(bidx) + FIRST_STATE_CHUNK = 0 if cfg.use_initial_state else 1 + + nvvm.barrier_cta_sync_aligned(cfg.tmem_lifecycle_barrier_id, thread_count=cfg.tmem_user_threads) + tmem_base = tmem_base_slot.load() + tmem_col = tmem_base & 0xFFFF + tmem_row = tmem_base >> 16 + row_lo_addr = tmem_row << 16 + row_hi_addr = (tmem_row + 16) << 16 + tmem_dvdk_acc_col = tmem_col + cfg.tmem_dvdk_acc_offset + tmem_dk_state_path_col = tmem_col + cfg.tmem_shared_input_offset + tmem_dstate_acc_col = tmem_col + cfg.tmem_dstate_acc_offset + tmem_dstate_input_col = tmem_col + cfg.tmem_dstate_input_offset + + # dGate fold scratch in the dK staging tile + sdk_red = cute.make_ptr(cutlass.Float32, sdK[0].base.toint(), mem_space=cute.AddressSpace.smem, assumed_align=cfg.buffer_align_bytes) + while tile_idx < total_tiles: + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) + num_item_chunks = compute_end - write_start + + # ---- dstate seed prologue ---------------------------------------------------- + if cutlass.const_expr(cfg.use_dstate_in): + if num_item_chunks > 0: + gDstate_in = mDstate_in[None, None, head_idx, batch_idx] + seed_from_dstate_in = compute_end == batch_num_chunks + dstate_input_idx = dstate_input_index.idx + bars.mb_dstate_input_done[dstate_input_idx].wait(dstate_input_index.phase) + dstate_input_index = advance(dstate_input_index, cfg.tmem_dstate_input_stages) + for i in cutlass.range_constexpr(num_ldtms): + dstate_in_vals = [] + for kk in cutlass.range_constexpr(ldtm_width): + v = gDstate_in[cg2_tidx, i * ldtm_width + kk] + v = v if seed_from_dstate_in else cutlass.Float32(0.0) + dstate_in_vals.append(v) + nvvm.tcgen05_st( + "32x32b", + nvvm.make_tmem_ptr(row_lo_addr + tmem_dstate_acc_col + i * ldtm_width, cutlass.Float32), + cutlass.Vector.from_elements(tuple(dstate_in_vals), cutlass.Float32), + ) + dstate_in_pack = [fp32_to_fp16(dstate_in_vals[2 * j], dstate_in_vals[2 * j + 1], dtype=cfg.io_dtype) for j in range(16)] + nvvm.tcgen05_st( + "32x32b", + nvvm.make_tmem_ptr(row_lo_addr + tmem_dstate_input_col + i * sttm_width, cutlass.Int32), + cutlass.Vector.from_elements(tuple(dstate_in_pack), cutlass.Int32), + ) + nvvm.tcgen05_wait("store") + bars.mb_dstate_input_ready[dstate_input_idx].arrive() + + for i in cutlass.range_constexpr(num_ldtms): + dstate_smem_vec = nvvm.tcgen05_ld( + "32x32b", + nvvm.make_tmem_ptr(row_lo_addr + tmem_dstate_acc_col + i * ldtm_width, cutlass.Float32), + num=32, + ) + for g in cutlass.range_constexpr(ldtm_width // 8): + dstate_smem_pack = tuple( + fp32_to_fp16(dstate_smem_vec[g * 8 + 2 * t], dstate_smem_vec[g * 8 + 2 * t + 1], dtype=cfg.io_dtype) for t in range(4) + ) + dstate_dk = i * ldtm_width + g * 8 + dstate_addr = (dstate_dk // 64) * (cfg.d_v * 64) + cg2_tidx * 64 + swizzle_xor_128b(cg2_tidx, dstate_dk % 64) + (sDstate_base + dstate_addr).store(cutlass.Vector.from_elements(dstate_smem_pack, cutlass.Int32).bitcast(cfg.io_dtype), alignment=16) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_dstate_smem_ready[0].arrive() + + for rev_idx in cutlass.range(num_item_chunks): + abs_chunk = compute_end - 1 - rev_idx + have_dstate = cutlass.Boolean(True) if cutlass.const_expr(cfg.use_dstate_in) else rev_idx > 0 + + # ---- decay-scale latch --------------------------------------------------- + gate_idx = gate_index.idx + bars.mb_gate_ready[gate_idx].wait(gate_index.phase) + gate_index = advance(gate_index, cfg.smem_gate_stages) + last_cumsumlog = sCumsumlog[cfg.b_t - 1, 0, gate_idx] + decay_scale_fp32 = [] + for i in cutlass.range_constexpr(16): + decay_scale_fp32.append(cute.math.exp2(last_cumsumlog - sCumsumlog[(lane_idx % 4) * 2 + (i // 2) * 8 + i % 2, 0, gate_idx], fastmath=True)) + decay_scale_vals = [decay_scale_fp32[(k // 4) * 2 + (k % 2)] for k in range(num_vals)] + bars.mb_cg2_chunk_ready[0].arrive() + + # ---- dK inter rescale ---------------------------------------------------- + cg2_k_idx = cg2_k_index.idx + cg2_k_index = advance(cg2_k_index, cfg.smem_k_stages) + k_dot_dk_inter = cutlass.Float32(0.0) + if rev_idx + DSTATE_IN0 >= 1: + bars.mb_dk_scale_acc_ready[0].wait(cg2_kk_scale_ready.phase) + cg2_kk_scale_ready = advance(cg2_kk_scale_ready, 1) + k_dot_dk_inter_lo = [opaque_f32_zero(), opaque_f32_zero(), opaque_f32_zero(), opaque_f32_zero()] + k_dot_dk_inter_hi = [opaque_f32_zero(), opaque_f32_zero(), opaque_f32_zero(), opaque_f32_zero()] + for half in cutlass.range_constexpr(2): + dki_ptr = nvvm.make_tmem_ptr(((tmem_row + half * 16) << 16) + tmem_dvdk_acc_col, cutlass.Float32) + dki_vec = nvvm.tcgen05_ld("16x256b", dki_ptr, num=8) + dki_scaled = [] + for j in cutlass.range_constexpr(16): + s0, s1 = fmul2(dki_vec[2 * j], dki_vec[2 * j + 1], decay_scale_vals[2 * j], decay_scale_vals[2 * j + 1]) + dki_scaled += [s0, s1] + nvvm.tcgen05_st("16x256b", dki_ptr, cutlass.Vector.from_elements(tuple(dki_scaled), cutlass.Float32)) + for block in cutlass.range_constexpr(4): + frag_addr = frag_segment + (frag_row + block * 16) * 64 + swizzle_xor_128b(frag_row + block * 16, frag_col + half * 16) + k_frag = nvvm.ldmatrix(sK_base + cg2_k_idx * k_stage_elements_cg2 + frag_addr, 4, nvvm.MMALayout.COL) + for i in cutlass.range_constexpr(4): + k_lo, k_hi = f16x2_to_f32(k_frag[i], dtype=cfg.io_dtype) + k_dot_dk_inter_lo[i], k_dot_dk_inter_hi[i] = ffma2( + dki_scaled[8 * block + 2 * i], dki_scaled[8 * block + 2 * i + 1], k_lo, k_hi, k_dot_dk_inter_lo[i], k_dot_dk_inter_hi[i] + ) + nvvm.tcgen05_wait("store") + bars.mb_dk_scale_acc_done[0].arrive() + k_dot_dk_inter = ((k_dot_dk_inter_lo[0] + k_dot_dk_inter_lo[1]) + (k_dot_dk_inter_lo[2] + k_dot_dk_inter_lo[3])) + ( + (k_dot_dk_inter_hi[0] + k_dot_dk_inter_hi[1]) + (k_dot_dk_inter_hi[2] + k_dot_dk_inter_hi[3]) + ) + for off in [1, 2, 4, 8, 16]: + k_dot_dk_inter += nvvm.shfl_sync(0xFFFFFFFF, k_dot_dk_inter, off, 31, kind=nvvm.Shfl.BFLY) + + # ---- dK attn read: TMEM -> registers ------------------------------------- + bars.mb_dk_attn_acc_ready[0].wait(cg2_dk_attn_ready.phase) + cg2_dk_attn_ready = advance(cg2_dk_attn_ready, 1) + dks_regs = [] + for half in cutlass.range_constexpr(2): + dks_vec = nvvm.tcgen05_ld( + "16x256b", + nvvm.make_tmem_ptr(((tmem_row + half * 16) << 16) + tmem_dvdk_acc_col, cutlass.Float32), + num=8, + ) + dks_regs.append([dks_vec[k] for k in range(32)]) + nvvm.tcgen05_wait("load") + bars.mb_dk_attn_acc_done[0].arrive() + + # ---- dK stage slot ------------------------------------------------------- + dk_stg_idx = dk_index.idx + bars.mb_dk_tmastg_done[dk_stg_idx].wait(dk_index.phase) + dk_index = advance(dk_index, cfg.smem_dk_stages) + + # ---- part reductions ----------------------------------------------------- + part_k = [acc_zero] * 16 + for half in cutlass.range_constexpr(2): + for block in cutlass.range_constexpr(4): + frag_addr = frag_segment + (frag_row + block * 16) * 64 + swizzle_xor_128b(frag_row + block * 16, frag_col + half * 16) + k_frag = nvvm.ldmatrix(sK_base + cg2_k_idx * k_stage_elements_cg2 + frag_addr, 4, nvvm.MMALayout.COL) + for i in cutlass.range_constexpr(4): + k_lo, k_hi = f16x2_to_f32(k_frag[i], dtype=cfg.io_dtype) + frag_e0 = cutlass.const_expr(8 * block + 2 * i) + part_e0 = cutlass.const_expr((frag_e0 // 4) * 2 + (frag_e0 % 2)) + if cutlass.const_expr(half == 0 and i % 2 == 0): + part_k[part_e0], part_k[part_e0 + 1] = fmul2(dks_regs[half][frag_e0], dks_regs[half][frag_e0 + 1], k_lo, k_hi) + else: + part_k[part_e0], part_k[part_e0 + 1] = ffma2( + dks_regs[half][frag_e0], dks_regs[half][frag_e0 + 1], k_lo, k_hi, part_k[part_e0], part_k[part_e0 + 1] + ) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_k_cg2_done[cg2_k_idx].arrive() + + part_k_lo, part_k_hi = warp_reduce_scatter_sum_16_elements(part_k, lane_idx) + tok0 = (lane_idx // 4) * 8 + (lane_idx % 4) * 2 + (sdk_red + warp_id * 64 + tok0).store(part_k_lo) + (sdk_red + warp_id * 64 + tok0 + 1).store(part_k_hi - k_dot_dk_inter if lane_idx == 31 else part_k_hi) + + bars.mb_dq_acc_total_done[0].wait(cg2_dq_total_index.phase) + cg2_dq_total_index = advance(cg2_dq_total_index, 1) + # ---- NEXT-CHUNK dstate prep ---------------------------------------------- - if chunk_idx >= wstart + 1: + if abs_chunk >= write_start + 1: dstate_idx = dstate_acc_index.idx bars.mb_dstate_acc_ready[dstate_idx].wait(dstate_acc_index.phase) dstate_acc_index = advance(dstate_acc_index, cfg.tmem_dstate_acc_stages) dstate_done_idx = dstate_idx - dstate_inp_idx = dstate_inp_index.idx - bars.mb_dstate_inp_done[dstate_inp_idx].wait(dstate_inp_index.phase) - dstate_inp_index = advance(dstate_inp_index, cfg.tmem_dstate_inp_stages) - dstate_regs = [[cutlass.Float32(0.0) for _ in range(num_state_subs)] for _ in range(32)] - for sub in cutlass.range_constexpr(num_state_subs): - dstate_vec = nvvm.tcgen05_ld( - "32x32b", nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_dstate_acc_col + sub * ldtm_width, cutlass.Float32), num=32 - ) - for k in cutlass.range_constexpr(32): - dstate_regs[k][sub] = dstate_vec[k] - - dstate_pack = [fp32_to_fp16(dstate_regs[2 * j][sub], dstate_regs[2 * j + 1][sub], dtype=cfg.io_dtype) for j in range(16)] + dstate_input_idx = dstate_input_index.idx + bars.mb_dstate_input_done[dstate_input_idx].wait(dstate_input_index.phase) + dstate_input_index = advance(dstate_input_index, cfg.tmem_dstate_input_stages) + for i in cutlass.range_constexpr(num_ldtms): + dstate_vec = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_lo_addr + tmem_dstate_acc_col + i * ldtm_width, cutlass.Float32), num=32) + dstate_pack = [fp32_to_fp16(dstate_vec[2 * j], dstate_vec[2 * j + 1], dtype=cfg.io_dtype) for j in range(16)] nvvm.tcgen05_st( "32x32b", - nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_dstate_inp_col + sub * sttm_width, cutlass.Int32), + nvvm.make_tmem_ptr(row_lo_addr + tmem_dstate_input_col + i * sttm_width, cutlass.Int32), cutlass.Vector.from_elements(tuple(dstate_pack), cutlass.Int32), ) nvvm.tcgen05_wait("store") - bars.mb_dstate_inp_ready[dstate_inp_idx].arrive() + bars.mb_dstate_input_ready[dstate_input_idx].arrive() - # ---- dK fold --------------------------------------------------------- - dk_stg_idx = dk_index.idx - bars.mb_dk_tmastg_done[dk_stg_idx].wait(dk_index.phase) - dk_index = advance(dk_index, cfg.smem_dk_stages) - if chunk_idx >= FIRST_STATE_CHUNK: - bars.mb_dk_state_path_acc_ready[0].wait(cg1_dk_state_path_ready.phase) - cg1_dk_state_path_ready = advance(cg1_dk_state_path_ready, 1) - dk_state_path_vecs = [] - for sub in cutlass.range_constexpr(2): - dk_state_path_vecs.append( - nvvm.tcgen05_ld( - "16x256b", - nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_dk_state_path_col, cutlass.Float32), - num=8, - ) + # ---- dstate prep: TMEM -> SMEM ------------------------------------------- + if abs_chunk >= write_start + 1: + bars.mb_state_dot_dstate_done[0].wait(cg2_state_dot_dstate_index.phase) + cg2_state_dot_dstate_index = advance(cg2_state_dot_dstate_index, 1) + for i in cutlass.range_constexpr(num_ldtms): + dstate_smem_vec = nvvm.tcgen05_ld( + "32x32b", + nvvm.make_tmem_ptr(row_lo_addr + tmem_dstate_acc_col + i * ldtm_width, cutlass.Float32), + num=32, ) - nvvm.tcgen05_wait("load") - dk_state_path_regs = [] - for sub in cutlass.range_constexpr(2): - dk_state_path_row = [] - for j in cutlass.range_constexpr(16): - n0, n1 = fmul2(dk_state_path_vecs[sub][2 * j], dk_state_path_vecs[sub][2 * j + 1], cumprod_neg_vals[2 * j], cumprod_neg_vals[2 * j + 1]) - dk_state_path_row += [n0, n1] - dk_state_path_regs.append(dk_state_path_row) - - bars.mb_dk_total_acc_ready[0].wait(cg1_dk_total_ready.phase) - cg1_dk_total_ready = advance(cg1_dk_total_ready, 1) - dmr_vecs = [] - for sub in cutlass.range_constexpr(2): - dmr_vecs.append( - nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_dvdk_acc_col, cutlass.Float32), num=8) + for g in cutlass.range_constexpr(ldtm_width // 8): + dstate_smem_pack = tuple( + fp32_to_fp16(dstate_smem_vec[g * 8 + 2 * t], dstate_smem_vec[g * 8 + 2 * t + 1], dtype=cfg.io_dtype) for t in range(4) + ) + dstate_dk = i * ldtm_width + g * 8 + dstate_addr = (dstate_dk // 64) * (cfg.d_v * 64) + cg2_tidx * 64 + swizzle_xor_128b(cg2_tidx, dstate_dk % 64) + (sDstate_base + dstate_addr).store(cutlass.Vector.from_elements(dstate_smem_pack, cutlass.Int32).bitcast(cfg.io_dtype), alignment=16) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_dstate_smem_ready[0].arrive() + + if abs_chunk < write_start + 1: + cg2_state_dot_dstate_index = advance(cg2_state_dot_dstate_index, 1) + + bars.mb_dgate_cg1_ready[0].wait(cg2_dgate_index.phase) + cg2_dgate_index = advance(cg2_dgate_index, 1) + nvvm.barrier_cta_sync_aligned(cfg.cg2_barrier_id, thread_count=cfg.cg2_barrier_threads) + if cg2_tidx < 64: + dgate_sum = ( + (sdk_red + cg2_tidx).load() + (sdk_red + 64 + cg2_tidx).load() + (sdk_red + 128 + cg2_tidx).load() + (sdk_red + 192 + cg2_tidx).load() + ) + sCumsumlog[cg2_tidx, 0, gate_idx] = sCumsumlog[cg2_tidx, 0, gate_idx] - dgate_sum + nvvm.barrier_cta_sync_aligned(cfg.cg2_barrier_id, thread_count=cfg.cg2_barrier_threads) + + # ---- dK fold ------------------------------------------------------------- + if abs_chunk >= FIRST_STATE_CHUNK: + cumprod_neg = [] + for i in cutlass.range_constexpr(16): + cumprod_neg.append(sCumprod[(lane_idx % 4) * 2 + (i // 2) * 8 + i % 2, 0, gate_idx] * neg_one) + cumprod_neg_vals = [cumprod_neg[(k // 4) * 2 + (k % 2)] for k in range(num_vals)] + bars.mb_dk_state_path_acc_ready[0].wait(cg2_dk_state_path_ready.phase) + cg2_dk_state_path_ready = advance(cg2_dk_state_path_ready, 1) + bars.mb_dk_total_acc_ready[0].wait(cg2_dk_total_ready.phase) + cg2_dk_total_ready = advance(cg2_dk_total_ready, 1) + for half in cutlass.range_constexpr(2): + dk_state_path_vec = nvvm.tcgen05_ld( + "16x256b", + nvvm.make_tmem_ptr(((tmem_row + half * 16) << 16) + tmem_dk_state_path_col, cutlass.Float32), + num=8, ) - nvvm.tcgen05_wait("load") - bars.mb_dk_total_acc_done[0].arrive() - for sub in cutlass.range_constexpr(2): - dk_sum = [dmr_vecs[sub][k] + dk_state_path_regs[sub][k] for k in range(32)] - for m0 in cutlass.range_constexpr(4): - dk_pack = [fp32_to_fp16(dk_sum[8 * m0 + 2 * j], dk_sum[8 * m0 + 2 * j + 1], dtype=cfg.io_dtype) for j in range(4)] + dmr_vec = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(((tmem_row + half * 16) << 16) + tmem_dvdk_acc_col, cutlass.Float32), num=8) + nvvm.tcgen05_wait("load") + if cutlass.const_expr(half == 1): + bars.mb_dk_total_acc_done[0].arrive() + for block in cutlass.range_constexpr(4): + dk_row = [] + for j in cutlass.range_constexpr(4): + e = cutlass.const_expr(8 * block + 2 * j) + n0, n1 = fmul2(dk_state_path_vec[e], dk_state_path_vec[e + 1], cumprod_neg_vals[e], cumprod_neg_vals[e + 1]) + dk_row += [dmr_vec[e] + n0, dmr_vec[e + 1] + n1] + dk_pack = [fp32_to_fp16(dk_row[2 * j], dk_row[2 * j + 1], dtype=cfg.io_dtype) for j in range(4)] nvvm.stmatrix( ( sdK_base - + dk_stg_idx * dk_stage_elems - + frag_slab_off - + (frag_row + m0 * 16) * 64 - + swizzle_xor_128b(frag_row + m0 * 16, frag_col + sub * 16) - ).raw_ptr(), + + dk_stg_idx * dk_stage_elements + + frag_segment + + (frag_row + block * 16) * 64 + + swizzle_xor_128b(frag_row + block * 16, frag_col + half * 16) + ), dk_pack, nvvm.MMALayout.COL, ) - nvvm.fence_proxy("async.shared", space="cta") - bars.mb_dk_tmastg_ready[dk_stg_idx].arrive() - if chunk_idx < FIRST_STATE_CHUNK: - bars.mb_dk_total_acc_ready[0].wait(cg1_dk_total_ready.phase) - cg1_dk_total_ready = advance(cg1_dk_total_ready, 1) - dmr_vecs = [] - for sub in cutlass.range_constexpr(2): - dmr_vecs.append( - nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_dvdk_acc_col, cutlass.Float32), num=8) - ) - nvvm.tcgen05_wait("load") - bars.mb_dk_total_acc_done[0].arrive() - for sub in cutlass.range_constexpr(2): - for m0 in cutlass.range_constexpr(4): - dk_pack = [fp32_to_fp16(dmr_vecs[sub][8 * m0 + 2 * j], dmr_vecs[sub][8 * m0 + 2 * j + 1], dtype=cfg.io_dtype) for j in range(4)] + if abs_chunk < FIRST_STATE_CHUNK: + bars.mb_dk_total_acc_ready[0].wait(cg2_dk_total_ready.phase) + cg2_dk_total_ready = advance(cg2_dk_total_ready, 1) + for half in cutlass.range_constexpr(2): + dmr_vec = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(((tmem_row + half * 16) << 16) + tmem_dvdk_acc_col, cutlass.Float32), num=8) + nvvm.tcgen05_wait("load") + if cutlass.const_expr(half == 1): + bars.mb_dk_total_acc_done[0].arrive() + for block in cutlass.range_constexpr(4): + dk_pack = [fp32_to_fp16(dmr_vec[8 * block + 2 * j], dmr_vec[8 * block + 2 * j + 1], dtype=cfg.io_dtype) for j in range(4)] nvvm.stmatrix( ( sdK_base - + dk_stg_idx * dk_stage_elems - + frag_slab_off - + (frag_row + m0 * 16) * 64 - + swizzle_xor_128b(frag_row + m0 * 16, frag_col + sub * 16) - ).raw_ptr(), + + dk_stg_idx * dk_stage_elements + + frag_segment + + (frag_row + block * 16) * 64 + + swizzle_xor_128b(frag_row + block * 16, frag_col + half * 16) + ), dk_pack, nvvm.MMALayout.COL, ) - nvvm.fence_proxy("async.shared", space="cta") - bars.mb_dk_tmastg_ready[dk_stg_idx].arrive() - - # ---- dstate prep --------------------------------------------------------- - if chunk_idx >= wstart + 1: - bars.mb_state_dot_dstate_done[0].wait(cg1_state_dot_dstate_index.phase) - cg1_state_dot_dstate_index = advance(cg1_state_dot_dstate_index, 1) - for sub in cutlass.range_constexpr(num_state_subs): - dstate_smem_vec = nvvm.tcgen05_ld( - "32x32b", - nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_dstate_acc_col + sub * ldtm_width, cutlass.Float32), - num=32, - ) - for g in cutlass.range_constexpr(ldtm_width // 8): - dstate_smem_pack = tuple( - fp32_to_fp16(dstate_smem_vec[g * 8 + 2 * t], dstate_smem_vec[g * 8 + 2 * t + 1], dtype=cfg.io_dtype) for t in range(4) - ) - dstate_dk = sub * ldtm_width + g * 8 - dstate_addr = (dstate_dk // 64) * (cfg.d_v * 64) + cg1_tidx * 64 + swizzle_xor_128b(cg1_tidx, dstate_dk % 64) - cutlass.inttoptr(sDstate_base_int + dstate_addr * 2, cutlass.AddressSpace.smem, cfg.io_dtype).store( - cutlass.Vector.from_elements(dstate_smem_pack, cutlass.Int32).bitcast(cfg.io_dtype), alignment=16 - ) - nvvm.fence_proxy("async.shared", space="cta") - bars.mb_dstate_smem_ready[0].arrive() - - if chunk_idx < wstart + 1: - cg1_state_dot_dstate_index = advance(cg1_state_dot_dstate_index, 1) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_dk_tmastg_ready[dk_stg_idx].arrive() + bars.mb_gate_done[gate_idx].arrive() - # ---- dstate drain: with an initial state this is d_initial_state ---------------------- + # ---- dstate store: with an initial state this is its gradient ---------------- if num_item_chunks > 0: dstate_idx = dstate_acc_index.idx bars.mb_dstate_acc_ready[dstate_idx].wait(dstate_acc_index.phase) dstate_acc_index = advance(dstate_acc_index, cfg.tmem_dstate_acc_stages) if cutlass.const_expr(cfg.use_dstate0): - if wstart == 0: + if write_start == 0: gDstate0 = mDstate0_out[None, None, head_idx, batch_idx] - for sub in cutlass.range_constexpr(num_state_subs): - dstate0_vec = nvvm.tcgen05_ld( - "32x32b", nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_dstate_acc_col + sub * ldtm_width, cutlass.Float32), num=32 - ) + for i in cutlass.range_constexpr(num_ldtms): + dstate0_vec = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_lo_addr + tmem_dstate_acc_col + i * ldtm_width, cutlass.Float32), num=32) for kk in cutlass.range_constexpr(32): - gDstate0[cg1_tidx, sub * ldtm_width + kk] = dstate0_vec[kk] + gDstate0[cg2_tidx, i * ldtm_width + kk] = dstate0_vec[kk] if cutlass.const_expr(not cfg.use_dstate_in): bars.mb_dstate_scale_acc_done[dstate_idx].arrive() else: if cutlass.const_expr(cfg.use_dstate0): - write_passthrough = wstart == 0 + write_passthrough = write_start == 0 if write_passthrough: gDstate0 = mDstate0_out[None, None, head_idx, batch_idx] if cutlass.const_expr(cfg.use_dstate_in): gDstate_in = mDstate_in[None, None, head_idx, batch_idx] - for sub in cutlass.range_constexpr(num_state_subs): + for i in cutlass.range_constexpr(num_ldtms): for kk in cutlass.range_constexpr(32): - gDstate0[cg1_tidx, sub * ldtm_width + kk] = gDstate_in[cg1_tidx, sub * ldtm_width + kk] + gDstate0[cg2_tidx, i * ldtm_width + kk] = gDstate_in[cg2_tidx, i * ldtm_width + kk] else: - for sub in cutlass.range_constexpr(num_state_subs): + for i in cutlass.range_constexpr(num_ldtms): for kk in cutlass.range_constexpr(32): - gDstate0[cg1_tidx, sub * ldtm_width + kk] = cutlass.Float32(0.0) + gDstate0[cg2_tidx, i * ldtm_width + kk] = cutlass.Float32(0.0) - tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one) bars.mb_tmem_done[0].arrive() - for _ in range(cfg.tmem_dstate_inp_stages): - bars.mb_dstate_inp_done[dstate_inp_index.idx].wait(dstate_inp_index.phase) - dstate_inp_index = advance(dstate_inp_index, cfg.tmem_dstate_inp_stages) + for _ in range(cfg.tmem_dstate_input_stages): + bars.mb_dstate_input_done[dstate_input_index.idx].wait(dstate_input_index.phase) + dstate_input_index = advance(dstate_input_index, cfg.tmem_dstate_input_stages) for _ in range(cfg.smem_dk_stages): bars.mb_dk_tmastg_done[dk_index.idx].wait(dk_index.phase) dk_index = advance(dk_index, cfg.smem_dk_stages) - for _ in range(cfg.smem_dv_stages): - bars.mb_dv_tmastg_done[dv_index.idx].wait(dv_index.phase) - dv_index = advance(dv_index, cfg.smem_dv_stages) @cute.jit @@ -3184,7 +3131,7 @@ def build_descs_body( base_dq, base_dk, base_dv, - desc_ws: cute.Tensor, + desc_workspace: cute.Tensor, cu_seqlens: cute.Tensor, q: cute.Tensor, k: cute.Tensor, @@ -3209,14 +3156,14 @@ def build_descs_body( prologue kernel after its order pass; warps past the array count fall through the widx guards.""" arr_words = n_batch * cutlass.Int32(TENSOR_MAP_QWORDS) - sub0 = cute.make_tensor(desc_ws.iterator, cute.make_layout((arr_words,), stride=(1,))) - sub1 = cute.make_tensor(desc_ws.iterator + arr_words, cute.make_layout((arr_words,), stride=(1,))) - sub2 = cute.make_tensor(desc_ws.iterator + 2 * arr_words, cute.make_layout((arr_words,), stride=(1,))) - sub3 = cute.make_tensor(desc_ws.iterator + 3 * arr_words, cute.make_layout((arr_words,), stride=(1,))) - sub4 = cute.make_tensor(desc_ws.iterator + 4 * arr_words, cute.make_layout((arr_words,), stride=(1,))) - sub5 = cute.make_tensor(desc_ws.iterator + 5 * arr_words, cute.make_layout((arr_words,), stride=(1,))) - sub6 = cute.make_tensor(desc_ws.iterator + 6 * arr_words, cute.make_layout((arr_words,), stride=(1,))) - sub7 = cute.make_tensor(desc_ws.iterator + 7 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub0 = cute.make_tensor(desc_workspace.iterator, cute.make_layout((arr_words,), stride=(1,))) + sub1 = cute.make_tensor(desc_workspace.iterator + arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub2 = cute.make_tensor(desc_workspace.iterator + 2 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub3 = cute.make_tensor(desc_workspace.iterator + 3 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub4 = cute.make_tensor(desc_workspace.iterator + 4 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub5 = cute.make_tensor(desc_workspace.iterator + 5 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub6 = cute.make_tensor(desc_workspace.iterator + 6 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub7 = cute.make_tensor(desc_workspace.iterator + 7 * arr_words, cute.make_layout((arr_words,), stride=(1,))) if widx == 0: if nvvm.elect_sync(): @@ -3265,7 +3212,7 @@ def prologue_kernel( base_dq: cutlass.GridConstant[tma.TensorMap], base_dk: cutlass.GridConstant[tma.TensorMap], base_dv: cutlass.GridConstant[tma.TensorMap], - desc_ws: cute.Tensor, + desc_workspace: cute.Tensor, cu_seqlens: cute.Tensor, q: cute.Tensor, k: cute.Tensor, @@ -3278,7 +3225,7 @@ def prologue_kernel( mStaging: Optional[cute.Tensor], mCount: cute.Tensor, mWorkItems: cute.Tensor, - mSched: Optional[cute.Tensor], + mScheduler: Optional[cute.Tensor], n_batch: cutlass.Int32, q_rs: cutlass.Int32, k_rs: cutlass.Int32, @@ -3308,7 +3255,7 @@ def prologue_kernel( True, b_t, ORDER_THREADS, - ORDER_ELEMS, + ORDER_ELEMENTS, tidx, n_heads_out, n_heads_out * n_batch, @@ -3316,7 +3263,7 @@ def prologue_kernel( mStaging, mCount, mWorkItems, - mSched, + mScheduler, sKey, sIdx, sSpread, @@ -3331,7 +3278,7 @@ def prologue_kernel( base_dq, base_dk, base_dv, - desc_ws, + desc_workspace, cu_seqlens, q, k, @@ -3372,7 +3319,7 @@ def prologue( work_item_staging: Optional[cute.Tensor], work_count: cute.Tensor, work_items: cute.Tensor, - sched_all: Optional[cute.Tensor], + scheduler_all: Optional[cute.Tensor], tensormap_workspace: cute.Tensor, stream: cuda.CUstream, ): @@ -3389,7 +3336,7 @@ def prologue( d_k_state = state_checkpoints.shape[2] d_v_state = state_checkpoints.shape[3] bpe = io_dtype.width // 8 - granule_elems = 128 // bpe + granule_elements = 128 // bpe bt = b_t q_row_stride, q_head_stride = q.stride[0], q.stride[1] @@ -3411,13 +3358,13 @@ def prologue( dk_headed = cute.make_tensor(dk.iterator, cute.make_layout((dk.shape[2], heads_out, seqlen), stride=(1, dk.stride[1], dk.stride[0]))) dv_headed = cute.make_tensor(dv.iterator, cute.make_layout((d_v, heads_out, seqlen), stride=(1, dv.stride[1], dv.stride[0]))) swz128 = tma.TensorMapSwizzle.s128b - base_desc_q = tma.create_tensor_map_tiled_from_view(q_headed, box_dims=(bt, 1, granule_elems), stride_order=(2, 1, 0), swizzle=swz128) - base_desc_k = tma.create_tensor_map_tiled_from_view(k_headed, box_dims=(bt, 1, granule_elems), stride_order=(2, 1, 0), swizzle=swz128) - base_desc_v = tma.create_tensor_map_tiled_from_view(v_headed, box_dims=(granule_elems, 1, bt), stride_order=(0, 1, 2), swizzle=swz128) - base_desc_do = tma.create_tensor_map_tiled_from_view(do_headed, box_dims=(granule_elems, 1, bt), stride_order=(0, 1, 2), swizzle=swz128) - base_desc_dq = tma.create_tensor_map_tiled_from_view(dq_headed, box_dims=(granule_elems, 1, bt), stride_order=(0, 1, 2), swizzle=swz128) - base_desc_dk = tma.create_tensor_map_tiled_from_view(dk_headed, box_dims=(granule_elems, 1, bt), stride_order=(0, 1, 2), swizzle=swz128) - base_desc_dv = tma.create_tensor_map_tiled_from_view(dv_headed, box_dims=(granule_elems, 1, bt), stride_order=(0, 1, 2), swizzle=swz128) + base_desc_q = tma.create_tensor_map_tiled_from_view(q_headed, box_dims=(bt, 1, granule_elements), stride_order=(2, 1, 0), swizzle=swz128) + base_desc_k = tma.create_tensor_map_tiled_from_view(k_headed, box_dims=(bt, 1, granule_elements), stride_order=(2, 1, 0), swizzle=swz128) + base_desc_v = tma.create_tensor_map_tiled_from_view(v_headed, box_dims=(granule_elements, 1, bt), stride_order=(0, 1, 2), swizzle=swz128) + base_desc_do = tma.create_tensor_map_tiled_from_view(do_headed, box_dims=(granule_elements, 1, bt), stride_order=(0, 1, 2), swizzle=swz128) + base_desc_dq = tma.create_tensor_map_tiled_from_view(dq_headed, box_dims=(granule_elements, 1, bt), stride_order=(0, 1, 2), swizzle=swz128) + base_desc_dk = tma.create_tensor_map_tiled_from_view(dk_headed, box_dims=(granule_elements, 1, bt), stride_order=(0, 1, 2), swizzle=swz128) + base_desc_dv = tma.create_tensor_map_tiled_from_view(dv_headed, box_dims=(granule_elements, 1, bt), stride_order=(0, 1, 2), swizzle=swz128) checkpoint_view = cute.make_tensor( state_checkpoints.iterator, cute.make_layout( @@ -3452,7 +3399,7 @@ def prologue( work_item_staging, work_count, work_items, - sched_all, + scheduler_all, cutlass.Int32(batch_size), cutlass.Int32(q_row_stride), cutlass.Int32(k_row_stride), @@ -3487,7 +3434,7 @@ def host( dstate_in: Optional[cute.Tensor], work_items: Optional[cute.Tensor], work_count: Optional[cute.Tensor], - sched_ctr: Optional[cute.Tensor], + scheduler_counter: Optional[cute.Tensor], scale: cutlass.Float32, tensormap_workspace: cute.Tensor, stream: cuda.CUstream, @@ -3498,37 +3445,37 @@ def host( batch_size = cu_seqlens.shape[0] - 1 heads_out = h_q if h_q >= h_v else h_v - # ---- SMEM sizing: per-buffer element cosizes --------------------------------- + # ---- SMEM sizing: per-buffer element cosizes ------------------------------------- bpe = cfg.io_dtype.width // 8 - q_tile_elems = cfg.b_t * cfg.d_k - k_tile_elems = cfg.b_t * cfg.d_k - v_tile_elems = cfg.d_v * cfg.b_t - do_tile_elems = cfg.d_v * cfg.b_t - state_tile_elems = cfg.d_k * cfg.d_v - tinv_tile_elems = cfg.b_t * cfg.b_t - a_tile_elems = cfg.b_t * cfg.b_t - dq_tile_elems = cfg.b_t * cfg.d_k - dk_tile_elems = cfg.b_t * cfg.d_k - dv_tile_elems = cfg.d_v * cfg.b_t - cfg.q_cosize = q_tile_elems * cfg.smem_q_stages - cfg.k_cosize = k_tile_elems * cfg.smem_k_stages - cfg.v_cosize = v_tile_elems * cfg.smem_v_stages - cfg.do_cosize = do_tile_elems * cfg.smem_do_stages - cfg.state_cosize = state_tile_elems * cfg.smem_state_stages - cfg.t_inv_cosize = tinv_tile_elems * cfg.smem_t_inv_stages - cfg.a_cosize = a_tile_elems * cfg.smem_a_stages - cfg.dq_cosize = dq_tile_elems * cfg.smem_dq_stages - cfg.dk_cosize = dk_tile_elems * cfg.smem_dk_stages - cfg.dv_cosize = dv_tile_elems * cfg.smem_dv_stages + q_tile_elements = cfg.b_t * cfg.d_k + k_tile_elements = cfg.b_t * cfg.d_k + v_tile_elements = cfg.d_v * cfg.b_t + do_tile_elements = cfg.d_v * cfg.b_t + state_tile_elements = cfg.d_k * cfg.d_v + tinv_tile_elements = cfg.b_t * cfg.b_t + a_tile_elements = cfg.b_t * cfg.b_t + dq_tile_elements = cfg.b_t * cfg.d_k + dk_tile_elements = cfg.b_t * cfg.d_k + dv_tile_elements = cfg.d_v * cfg.b_t + cfg.q_cosize = q_tile_elements * cfg.smem_q_stages + cfg.k_cosize = k_tile_elements * cfg.smem_k_stages + cfg.v_cosize = v_tile_elements * cfg.smem_v_stages + cfg.do_cosize = do_tile_elements * cfg.smem_do_stages + cfg.state_cosize = state_tile_elements * cfg.smem_state_stages + cfg.t_inv_cosize = tinv_tile_elements * cfg.smem_t_inv_stages + cfg.a_cosize = a_tile_elements * cfg.smem_a_stages + cfg.dq_cosize = dq_tile_elements * cfg.smem_dq_stages + cfg.dk_cosize = dk_tile_elements * cfg.smem_dk_stages + cfg.dv_cosize = dv_tile_elements * cfg.smem_dv_stages cumsumlog_smem_layout_staged = cute.make_layout((cfg.b_t, 1, cfg.smem_gate_stages)) beta_smem_layout_staged = cute.make_layout((cfg.b_t, 1, cfg.smem_beta_stages)) - cfg.tma_q_bytes = q_tile_elems * bpe - cfg.tma_k_bytes = k_tile_elems * bpe - cfg.tma_v_bytes = v_tile_elems * bpe - cfg.tma_do_bytes = do_tile_elems * bpe - cfg.tma_state_bytes = state_tile_elems * bpe + cfg.tma_q_bytes = q_tile_elements * bpe + cfg.tma_k_bytes = k_tile_elements * bpe + cfg.tma_v_bytes = v_tile_elements * bpe + cfg.tma_do_bytes = do_tile_elements * bpe + cfg.tma_state_bytes = state_tile_elements * bpe cfg.n_heads_out = heads_out cfg.q_ratio = heads_out // h_q @@ -3536,7 +3483,7 @@ def host( cfg.v_ratio = heads_out // h_v num_descs = batch_size - # ---- launch ------------------------------------------------------------------ + # ---- launch ---------------------------------------------------------------------- grid_shape = (cfg.max_active_clusters, 1, 1) kernel( @@ -3562,7 +3509,7 @@ def host( dstate_in, work_items, work_count, - sched_ctr, + scheduler_counter, tensormap_workspace, cutlass.Int32(num_descs), ).launch( @@ -3598,7 +3545,7 @@ def kernel( mDstate_in, mWorkItems: cute.Tensor, mCount: cute.Tensor, - mSched: Optional[cute.Tensor], + mScheduler: Optional[cute.Tensor], tensormap_workspace: cute.Tensor, n_desc: cutlass.Int32, ): @@ -3609,8 +3556,8 @@ def kernel( num_ctas = cute.arch.grid_dim()[0] total_tiles = mCount[0] - if cutlass.const_expr(cfg.dyn_sched): - assert mSched is not None, "mSched must be provided if dyn_sched is True" + if cutlass.const_expr(cfg.dynamic_scheduling): + assert mScheduler is not None, "mScheduler must be provided if dynamic_scheduling is True" if cutlass.const_expr(cfg.is_GQA): h_r = cfg.h_q // cfg.h_v @@ -3729,13 +3676,12 @@ def kernel( SMEM = cutlass.AddressSpace.smem bars = make_gdn_bars(cfg) - sSched = cutlass.Array(cutlass.Int32, cfg.sched_stages, space=cutlass.AddressSpace.smem, alignment=16) + sScheduler = cutlass.Array(cutlass.Int32, cfg.scheduler_stages, space=cutlass.AddressSpace.smem, alignment=16) tmem_base_slot = cutlass.Array(cutlass.Int32, 1, space=SMEM, alignment=16) cumsumlog_raw = cutlass.Array(cutlass.Float32, cute.cosize(cumsumlog_smem_layout_staged), space=SMEM, alignment=128) cumprod_raw = cutlass.Array(cutlass.Float32, cute.cosize(cumsumlog_smem_layout_staged), space=SMEM, alignment=128) beta_raw = cutlass.Array(cutlass.Float32, cute.cosize(beta_smem_layout_staged), space=SMEM, alignment=128) - bpe = cfg.io_dtype.width // 8 SWZ = 2 LEAD = 16 STRIDE = 8 * 128 @@ -3749,8 +3695,8 @@ def kernel( alignment=cfg.buffer_align_bytes, ) sQ = SmemTile( - base=sQ_raw.data_ptr().toint(), - elems_per_stage=(cfg.q_cosize // cfg.smem_q_stages) * bpe, + base=sQ_raw.data_ptr(), + elems_per_stage=(cfg.q_cosize // cfg.smem_q_stages), stages=cfg.smem_q_stages, leading_byte_offset=LEAD, stride_byte_offset=STRIDE, @@ -3763,24 +3709,24 @@ def kernel( alignment=cfg.buffer_align_bytes, ) sK = SmemTile( - base=sK_raw.data_ptr().toint(), - elems_per_stage=(cfg.k_cosize // cfg.smem_k_stages) * bpe, + base=sK_raw.data_ptr(), + elems_per_stage=(cfg.k_cosize // cfg.smem_k_stages), stages=cfg.smem_k_stages, leading_byte_offset=LEAD, stride_byte_offset=STRIDE, layout=SWZ, ) sK_trans = SmemTile( - base=sK_raw.data_ptr().toint(), - elems_per_stage=(cfg.k_cosize // cfg.smem_k_stages) * bpe, + base=sK_raw.data_ptr(), + elems_per_stage=(cfg.k_cosize // cfg.smem_k_stages), stages=cfg.smem_k_stages, leading_byte_offset=KT_LEAD, stride_byte_offset=STRIDE, layout=SWZ, ) sQ_trans = SmemTile( - base=sQ_raw.data_ptr().toint(), - elems_per_stage=(cfg.q_cosize // cfg.smem_q_stages) * bpe, + base=sQ_raw.data_ptr(), + elems_per_stage=(cfg.q_cosize // cfg.smem_q_stages), stages=cfg.smem_q_stages, leading_byte_offset=KT_LEAD, stride_byte_offset=STRIDE, @@ -3792,17 +3738,17 @@ def kernel( space=cutlass.AddressSpace.smem, alignment=cfg.buffer_align_bytes, ) - sdO = SmemTile( - base=sdO_raw.data_ptr().toint(), - elems_per_stage=(cfg.do_cosize // cfg.smem_do_stages) * bpe, + sdO_trans = SmemTile( + base=sdO_raw.data_ptr(), + elems_per_stage=(cfg.do_cosize // cfg.smem_do_stages), stages=cfg.smem_do_stages, leading_byte_offset=V_LEAD, stride_byte_offset=STRIDE, layout=SWZ, ) - sdO_kmaj = SmemTile( - base=sdO_raw.data_ptr().toint(), - elems_per_stage=(cfg.do_cosize // cfg.smem_do_stages) * bpe, + sdO = SmemTile( + base=sdO_raw.data_ptr(), + elems_per_stage=(cfg.do_cosize // cfg.smem_do_stages), stages=cfg.smem_do_stages, leading_byte_offset=LEAD, stride_byte_offset=STRIDE, @@ -3814,17 +3760,17 @@ def kernel( space=cutlass.AddressSpace.smem, alignment=cfg.buffer_align_bytes, ) - sState = SmemTile( - base=sState_raw.data_ptr().toint(), - elems_per_stage=(cfg.state_cosize // cfg.smem_state_stages) * bpe, + sState_trans = SmemTile( + base=sState_raw.data_ptr(), + elems_per_stage=(cfg.state_cosize // cfg.smem_state_stages), stages=cfg.smem_state_stages, leading_byte_offset=STATE_LEAD, stride_byte_offset=STRIDE, layout=SWZ, ) - sState_kmaj = SmemTile( - base=sState_raw.data_ptr().toint(), - elems_per_stage=(cfg.state_cosize // cfg.smem_state_stages) * bpe, + sState = SmemTile( + base=sState_raw.data_ptr(), + elems_per_stage=(cfg.state_cosize // cfg.smem_state_stages), stages=cfg.smem_state_stages, leading_byte_offset=LEAD, stride_byte_offset=STRIDE, @@ -3837,16 +3783,16 @@ def kernel( alignment=cfg.buffer_align_bytes, ) sTinv = SmemTile( - base=sTinv_raw.data_ptr().toint(), - elems_per_stage=(cfg.t_inv_cosize // cfg.smem_t_inv_stages) * bpe, + base=sTinv_raw.data_ptr(), + elems_per_stage=(cfg.t_inv_cosize // cfg.smem_t_inv_stages), stages=cfg.smem_t_inv_stages, leading_byte_offset=LEAD, stride_byte_offset=STRIDE, layout=SWZ, ) sTinv_trans = SmemTile( - base=sTinv_raw.data_ptr().toint(), - elems_per_stage=(cfg.t_inv_cosize // cfg.smem_t_inv_stages) * bpe, + base=sTinv_raw.data_ptr(), + elems_per_stage=(cfg.t_inv_cosize // cfg.smem_t_inv_stages), stages=cfg.smem_t_inv_stages, leading_byte_offset=(cfg.b_t // 2) * 128, stride_byte_offset=STRIDE, @@ -3859,8 +3805,8 @@ def kernel( alignment=cfg.buffer_align_bytes, ) sKK = SmemTile( - base=sKK_raw.data_ptr().toint(), - elems_per_stage=(cfg.t_inv_cosize // cfg.smem_t_inv_stages) * bpe, + base=sKK_raw.data_ptr(), + elems_per_stage=(cfg.t_inv_cosize // cfg.smem_t_inv_stages), stages=cfg.smem_t_inv_stages, leading_byte_offset=LEAD, stride_byte_offset=STRIDE, @@ -3873,32 +3819,32 @@ def kernel( alignment=cfg.buffer_align_bytes, ) sA = SmemTile( - base=sA_raw.data_ptr().toint(), - elems_per_stage=(cfg.a_cosize // cfg.smem_a_stages) * bpe, + base=sA_raw.data_ptr(), + elems_per_stage=(cfg.a_cosize // cfg.smem_a_stages), stages=cfg.smem_a_stages, leading_byte_offset=LEAD, stride_byte_offset=STRIDE, layout=SWZ, ) sA_trans = SmemTile( - base=sA_raw.data_ptr().toint(), - elems_per_stage=(cfg.a_cosize // cfg.smem_a_stages) * bpe, + base=sA_raw.data_ptr(), + elems_per_stage=(cfg.a_cosize // cfg.smem_a_stages), stages=cfg.smem_a_stages, leading_byte_offset=(cfg.b_t // 2) * 128, stride_byte_offset=STRIDE, layout=SWZ, ) sDa = SmemTile( - base=sA_raw.data_ptr().toint(), - elems_per_stage=(cfg.a_cosize // cfg.smem_a_stages) * bpe, + base=sA_raw.data_ptr(), + elems_per_stage=(cfg.a_cosize // cfg.smem_a_stages), stages=1, leading_byte_offset=LEAD, stride_byte_offset=STRIDE, layout=SWZ, ) sDa_trans = SmemTile( - base=sA_raw.data_ptr().toint(), - elems_per_stage=(cfg.a_cosize // cfg.smem_a_stages) * bpe, + base=sA_raw.data_ptr(), + elems_per_stage=(cfg.a_cosize // cfg.smem_a_stages), stages=1, leading_byte_offset=(cfg.b_t // 2) * 128, stride_byte_offset=STRIDE, @@ -3911,39 +3857,39 @@ def kernel( alignment=cfg.buffer_align_bytes, ) sDm = SmemTile( - base=sDm_raw.data_ptr().toint(), - elems_per_stage=(cfg.a_cosize // cfg.smem_a_stages) * bpe, + base=sDm_raw.data_ptr(), + elems_per_stage=(cfg.a_cosize // cfg.smem_a_stages), stages=1, leading_byte_offset=LEAD, stride_byte_offset=STRIDE, layout=SWZ, ) sDm_trans = SmemTile( - base=sDm_raw.data_ptr().toint(), - elems_per_stage=(cfg.a_cosize // cfg.smem_a_stages) * bpe, + base=sDm_raw.data_ptr(), + elems_per_stage=(cfg.a_cosize // cfg.smem_a_stages), stages=1, leading_byte_offset=(cfg.b_t // 2) * 128, stride_byte_offset=STRIDE, layout=SWZ, ) - # sub-bank split: V + sDstate + dQ + dK + dV allocated last + # sub-bank split: V + dstate_entry + dQ + dK + dV allocated last sV_raw = cutlass.Array( cfg.io_dtype, cfg.v_cosize, space=cutlass.AddressSpace.smem, alignment=cfg.buffer_align_bytes, ) - sV = SmemTile( - base=sV_raw.data_ptr().toint(), - elems_per_stage=(cfg.v_cosize // cfg.smem_v_stages) * bpe, + sV_trans = SmemTile( + base=sV_raw.data_ptr(), + elems_per_stage=(cfg.v_cosize // cfg.smem_v_stages), stages=cfg.smem_v_stages, leading_byte_offset=V_LEAD, stride_byte_offset=STRIDE, layout=SWZ, ) - sV_kmaj = SmemTile( - base=sV_raw.data_ptr().toint(), - elems_per_stage=(cfg.v_cosize // cfg.smem_v_stages) * bpe, + sV = SmemTile( + base=sV_raw.data_ptr(), + elems_per_stage=(cfg.v_cosize // cfg.smem_v_stages), stages=cfg.smem_v_stages, leading_byte_offset=LEAD, stride_byte_offset=STRIDE, @@ -3955,9 +3901,9 @@ def kernel( space=cutlass.AddressSpace.smem, alignment=cfg.buffer_align_bytes, ) - sDstate = SmemTile( - base=sDstate_raw.data_ptr().toint(), - elems_per_stage=cfg.d_k * cfg.d_v * bpe, + sDstate_trans = SmemTile( + base=sDstate_raw.data_ptr(), + elems_per_stage=cfg.d_k * cfg.d_v, stages=1, leading_byte_offset=STATE_LEAD, stride_byte_offset=STRIDE, @@ -3970,8 +3916,8 @@ def kernel( alignment=cfg.buffer_align_bytes, ) sdQ = SmemTile( - base=sdQ_raw.data_ptr().toint(), - elems_per_stage=(cfg.dq_cosize // cfg.smem_dq_stages) * bpe, + base=sdQ_raw.data_ptr(), + elems_per_stage=(cfg.dq_cosize // cfg.smem_dq_stages), stages=cfg.smem_dq_stages, leading_byte_offset=LEAD, stride_byte_offset=STRIDE, @@ -3984,8 +3930,8 @@ def kernel( alignment=cfg.buffer_align_bytes, ) sdK = SmemTile( - base=sdK_raw.data_ptr().toint(), - elems_per_stage=(cfg.dk_cosize // cfg.smem_dk_stages) * bpe, + base=sdK_raw.data_ptr(), + elems_per_stage=(cfg.dk_cosize // cfg.smem_dk_stages), stages=cfg.smem_dk_stages, leading_byte_offset=LEAD, stride_byte_offset=STRIDE, @@ -3998,14 +3944,13 @@ def kernel( alignment=cfg.buffer_align_bytes, ) sdV = SmemTile( - base=sdV_raw.data_ptr().toint(), - elems_per_stage=(cfg.dv_cosize // cfg.smem_dv_stages) * bpe, + base=sdV_raw.data_ptr(), + elems_per_stage=(cfg.dv_cosize // cfg.smem_dv_stages), stages=cfg.smem_dv_stages, leading_byte_offset=LEAD, stride_byte_offset=STRIDE, layout=SWZ, ) - sdV_kmaj = sdV sdstate_flat = cute.make_tensor( cute.make_ptr(cfg.io_dtype, sDstate_raw.data_ptr().toint(), mem_space=cute.AddressSpace.smem, assumed_align=cfg.buffer_align_bytes), cute.make_layout(cfg.d_k * cfg.d_v), @@ -4027,10 +3972,10 @@ def kernel( beta_smem_layout_staged, ) - # ---- mbarrier init (all threads) --------------------------------------------- - for s_ in range(cfg.sched_stages): - bars.mb_sched_ready[s_].init() - bars.mb_sched_done[s_].init() + # ---- mbarrier init (all threads) ------------------------------------------------- + for s in range(cfg.scheduler_stages): + bars.mb_scheduler_ready[s].init() + bars.mb_scheduler_done[s].init() for s in range(cfg.smem_q_stages): bars.mb_q_ready[s].init() bars.mb_q_mma_done[s].init() @@ -4038,7 +3983,7 @@ def kernel( for s in range(cfg.smem_k_stages): bars.mb_k_ready[s].init() bars.mb_k_mma_done[s].init() - bars.mb_k_cg0_done[s].init() + bars.mb_k_cg2_done[s].init() for s in range(cfg.smem_v_stages): bars.mb_v_ready[s].init() bars.mb_v_mma_done[s].init() @@ -4083,13 +4028,13 @@ def kernel( for s in range(cfg.smem_a_stages): bars.mb_a_ready[s].init() bars.mb_a_done[s].init() - for s in range(cfg.tmem_dstate_inp_stages): - bars.mb_dstate_inp_ready[s].init() - bars.mb_dstate_inp_done[s].init() + for s in range(cfg.tmem_dstate_input_stages): + bars.mb_dstate_input_ready[s].init() + bars.mb_dstate_input_done[s].init() for b in ( - bars.mb_do_prime_inp_ready, - bars.mb_du_inp_ready, - bars.mb_dyp_inp_ready, + bars.mb_do_prime_input_ready, + bars.mb_du_input_ready, + bars.mb_dyp_input_ready, ): b[0].init() for s in range(cfg.smem_dq_stages): @@ -4115,6 +4060,8 @@ def kernel( bars.mb_dm_acc_done[0].init() bars.mb_dbeta_cg1_ready[0].init() bars.mb_dgate_cg1_ready[0].init() + bars.mb_dgate_cg0_ready[0].init() + bars.mb_cg2_chunk_ready[0].init() bars.mb_state_dot_dstate_done[0].init() bars.mb_dk_state_path_acc_ready[0].init() bars.mb_tmem_done[0].init() @@ -4122,7 +4069,7 @@ def kernel( nvvm.fence_mbarrier_init() nvvm.barrier_cta_sync() - # ---- warp specialization ----------------------------------------------------- + # ---- warp specialization --------------------------------------------------------- if warp_idx >= cfg.compute_group_0_warp_ids[0] and warp_idx <= cfg.compute_group_0_warp_ids[-1]: compute0_warp_group( @@ -4143,12 +4090,9 @@ def kernel( sA=sA, sDa=sDa, sDm=sDm, - sK=sK, - sdQ=sdQ, - sDstate=sDstate, sstate_flat=sstate_flat, sdstate_flat=sdstate_flat, - sSched=sSched, + sScheduler=sScheduler, bars=bars, ) @@ -4160,31 +4104,47 @@ def kernel( num_ctas, cu_seqlens, mWorkItems, - mDstate0, - mDstate_in, tidx, warp_idx=warp_idx, tmem_base_slot=tmem_base_slot, scale=scale, sQ=sQ, - sK=sK, sV=sV, sdO=sdO, sCumsumlog=sCumsumlog, sCumprod=sCumprod, sBeta=sBeta, sdQ=sdQ, - sdK=sdK, sdV=sdV, - sDstate=sDstate, + sdK=sdK, sDa=sDa, sDm=sDm, - sSched=sSched, + sScheduler=sScheduler, bars=bars, ) - elif warp_idx == cfg.mma_warp_id: - mma_warp( + if warp_idx >= cfg.compute_group_2_warp_ids[0] and warp_idx <= cfg.compute_group_2_warp_ids[-1]: + compute2_warp_group( + cfg, + total_tiles, + bidx, + num_ctas, + mWorkItems, + mDstate0, + mDstate_in, + tidx, + tmem_base_slot=tmem_base_slot, + sCumsumlog=sCumsumlog, + sCumprod=sCumprod, + sK=sK, + sdK=sdK, + sDstate_trans=sDstate_trans, + sScheduler=sScheduler, + bars=bars, + ) + + elif warp_idx == cfg.tcgen05_mma_warp_id: + tcgen05_mma_warp( cfg, total_tiles, bidx, @@ -4196,23 +4156,23 @@ def kernel( sQ_trans=sQ_trans, sK=sK, sK_trans=sK_trans, + sV_trans=sV_trans, sV=sV, - sV_kmaj=sV_kmaj, + sdO_trans=sdO_trans, sdO=sdO, - sdO_kmaj=sdO_kmaj, + sState_trans=sState_trans, sState=sState, - sState_kmaj=sState_kmaj, sTinv=sTinv, sTinv_trans=sTinv_trans, sA=sA, sA_trans=sA_trans, sDa=sDa, sDa_trans=sDa_trans, - sDstate=sDstate, + sDstate_trans=sDstate_trans, sDm=sDm, sDm_trans=sDm_trans, - sdV_kmaj=sdV_kmaj, - sSched=sSched, + sdV=sdV, + sScheduler=sScheduler, bars=bars, ) @@ -4224,7 +4184,7 @@ def kernel( num_ctas, cu_seqlens, mWorkItems, - mSched=mSched, + mScheduler=mScheduler, sQ_raw=sQ_raw, sK_raw=sK_raw, sV_raw=sV_raw, @@ -4235,7 +4195,7 @@ def kernel( desc_v_base=desc_v_base, desc_do_base=desc_do_base, desc_checkpoint_base=desc_checkpoint_base, - sSched=sSched, + sScheduler=sScheduler, bars=bars, ) @@ -4257,7 +4217,7 @@ def kernel( sCumsumlog=sCumsumlog, sCumprod=sCumprod, sBeta=sBeta, - sSched=sSched, + sScheduler=sScheduler, bars=bars, ) if warp_idx == cfg.epilogue_warp_id: @@ -4274,7 +4234,7 @@ def kernel( desc_dq_base=desc_dq_base, desc_dk_base=desc_dk_base, desc_dv_base=desc_dv_base, - sSched=sSched, + sScheduler=sScheduler, bars=bars, ) @@ -4299,28 +4259,30 @@ class GdnBwdCfg: safe_gate: bool = False beta_sigmoid: bool = False - # ---- fixed constants stamped from CFG by build_cfg --------------------------- + # ---- fixed constants stamped from CFG at build time ------------------------------ b_t: int = CFG.B_T d_k: int = CFG.D_K d_v: int = CFG.D_V compute_group_0_warp_ids: Tuple[int, ...] = CFG.COMPUTE_GROUP_0_WARP_IDS compute_group_1_warp_ids: Tuple[int, ...] = CFG.COMPUTE_GROUP_1_WARP_IDS - mma_warp_id: int = CFG.MMA_WARP_ID + compute_group_2_warp_ids: Tuple[int, ...] = CFG.COMPUTE_GROUP_2_WARP_IDS + tcgen05_mma_warp_id: int = CFG.TCGEN05_MMA_WARP_ID tma_qkv_warp_id: int = CFG.TMA_QKV_WARP_ID load_gate_beta_warp_id: int = CFG.LOAD_GATE_BETA_WARP_ID epilogue_warp_id: int = CFG.EPILOGUE_WARP_ID num_regs_compute_group_0: int = CFG.NUM_REGS_COMPUTE_GROUP_0 num_regs_compute_group_1: int = CFG.NUM_REGS_COMPUTE_GROUP_1 + num_regs_compute_group_2: int = CFG.NUM_REGS_COMPUTE_GROUP_2 num_regs_other: int = CFG.NUM_REGS_OTHER threads_per_warp: int = CFG.THREADS_PER_WARP threads_per_cta: int = 0 cluster_shape_mnk: Tuple[int, int, int] = CFG.CLUSTER_SHAPE_MNK - dyn_sched: bool = False - sched_stages: int = CFG.SMEM_SCHED_STAGES + dynamic_scheduling: bool = False + scheduler_stages: int = CFG.SMEM_SCHEDULER_STAGES - # ---- named barrier slots (ids 1-6; 0 is the CTA-wide sync) ------------------- - tmem_alloc_barrier_id: int = 1 - tmem_alloc_barrier_threads: int = 0 + # ---- named barrier slots (ids 1-6; 0 is the CTA-wide sync) ----------------------- + tmem_lifecycle_barrier_id: int = 1 + tmem_user_threads: int = 0 inverse_barrier_id: int = 2 inverse_barrier_threads: int = 0 inverse_inner_barrier_id: int = 3 @@ -4329,8 +4291,10 @@ class GdnBwdCfg: init_state_store_barrier_threads: int = 0 cg1_barrier_id: int = 5 cg1_barrier_threads: int = 0 + cg2_barrier_id: int = 6 + cg2_barrier_threads: int = 0 - # ---- SMEM / TMEM stage counts + TMEM column offsets -------------------------- + # ---- SMEM / TMEM stage counts + TMEM column offsets ------------------------------ smem_q_stages: int = CFG.SMEM_Q_STAGES smem_k_stages: int = CFG.SMEM_K_STAGES smem_v_stages: int = CFG.SMEM_V_STAGES @@ -4345,18 +4309,18 @@ class GdnBwdCfg: smem_beta_stages: int = 2 tmem_dstate_acc_stages: int = CFG.TMEM_DH_ACC_STAGES tmem_dvdk_acc_stages: int = CFG.TMEM_DVDK_ACC_STAGES - tmem_dstate_inp_stages: int = CFG.TMEM_DH_INP_STAGES - tmem_shared_inp_stages: int = CFG.TMEM_SHARED_INP_STAGES + tmem_dstate_input_stages: int = CFG.TMEM_DH_INP_STAGES + tmem_shared_input_stages: int = CFG.TMEM_SHARED_INP_STAGES tmem_shared_acc_stages: int = CFG.TMEM_SHARED_ACC_STAGES tmem_dstate_acc_offset: int = 0 tmem_dvdk_acc_offset: int = 0 - tmem_dstate_inp_offset: int = 0 + tmem_dstate_input_offset: int = 0 tmem_shared_acc_offset: int = 0 - tmem_shared_inp_offset: int = 0 + tmem_shared_input_offset: int = 0 tmem_y_offset: int = 0 buffer_align_bytes: int = CFG.BUFFER_ALIGN_BYTES - # ---- stamped by host at trace time (shape-derived) -------------------------- + # ---- stamped by host at trace time (shape-derived) ------------------------------- q_cosize: int = 0 k_cosize: int = 0 v_cosize: int = 0 @@ -4389,7 +4353,7 @@ def build_cfg( log_gate: bool = False, safe_gate: bool = False, beta_sigmoid: bool = False, - dyn_sched: bool = False, + dynamic_scheduling: bool = False, ) -> GdnBwdCfg: """Build the per-compile ``GdnBwdCfg`` (io_dtype in {Float16, BFloat16}; acc is always Float32).""" @@ -4406,22 +4370,24 @@ def build_cfg( log_gate=log_gate, safe_gate=safe_gate, beta_sigmoid=beta_sigmoid, - dyn_sched=dyn_sched, + dynamic_scheduling=dynamic_scheduling, ) n_cg0 = len(cfg.compute_group_0_warp_ids) n_cg1 = len(cfg.compute_group_1_warp_ids) - cfg.threads_per_cta = cfg.threads_per_warp * (4 + n_cg0 + n_cg1) - cfg.tmem_alloc_barrier_threads = cfg.threads_per_warp * (1 + n_cg0 + n_cg1) + n_cg2 = len(cfg.compute_group_2_warp_ids) + cfg.threads_per_cta = cfg.threads_per_warp * (4 + n_cg0 + n_cg1 + n_cg2) + cfg.tmem_user_threads = cfg.threads_per_warp * (1 + n_cg0 + n_cg1 + n_cg2) cfg.inverse_barrier_threads = cfg.threads_per_warp * n_cg0 cfg.inverse_inner_barrier_threads = cfg.threads_per_warp * 2 cfg.init_state_store_barrier_threads = cfg.threads_per_warp * n_cg1 cfg.cg1_barrier_threads = cfg.threads_per_warp * n_cg1 + cfg.cg2_barrier_threads = cfg.threads_per_warp * n_cg2 cfg.tmem_dstate_acc_offset = 0 cfg.tmem_dvdk_acc_offset = cfg.tmem_dstate_acc_offset + cfg.tmem_dstate_acc_stages * 128 - cfg.tmem_dstate_inp_offset = cfg.tmem_dvdk_acc_offset + cfg.tmem_dvdk_acc_stages * 64 - cfg.tmem_shared_acc_offset = cfg.tmem_dstate_inp_offset + cfg.tmem_dstate_inp_stages * 64 - cfg.tmem_shared_inp_offset = cfg.tmem_shared_acc_offset + cfg.tmem_shared_acc_stages * 64 - cfg.tmem_y_offset = cfg.tmem_shared_inp_offset + cfg.tmem_shared_inp_stages * (cfg.b_t // 2) + cfg.tmem_dstate_input_offset = cfg.tmem_dvdk_acc_offset + cfg.tmem_dvdk_acc_stages * 64 + cfg.tmem_shared_acc_offset = cfg.tmem_dstate_input_offset + cfg.tmem_dstate_input_stages * 64 + cfg.tmem_shared_input_offset = cfg.tmem_shared_acc_offset + cfg.tmem_shared_acc_stages * 64 + cfg.tmem_y_offset = cfg.tmem_shared_input_offset + cfg.tmem_shared_input_stages * (cfg.b_t // 2) return cfg @@ -4443,7 +4409,7 @@ def get_compiled_cache( log_gate: bool = False, safe_gate: bool = False, beta_sigmoid: bool = False, - dyn_sched: bool = False, + dynamic_scheduling: bool = False, run_order: bool = False, order_gen: bool = False, ): @@ -4460,7 +4426,7 @@ def compile( log_gate: bool = False, safe_gate: bool = False, beta_sigmoid: bool = False, - dyn_sched: bool = False, + dynamic_scheduling: bool = False, *, num_sm: int, h_q: int, @@ -4484,7 +4450,7 @@ def compile( dstate_in_cute=None, work_items_cute=None, work_count_cute=None, - sched_ctr_cute=None, + scheduler_counter_cute=None, scale=None, workspace_cute=None, stream=None, @@ -4500,7 +4466,7 @@ def compile( log_gate=log_gate, safe_gate=safe_gate, beta_sigmoid=beta_sigmoid, - dyn_sched=dyn_sched, + dynamic_scheduling=dynamic_scheduling, ) cfg.h_q = h_q cfg.h_k = h_k @@ -4527,7 +4493,7 @@ def compile( dstate_in_cute, work_items_cute, work_count_cute, - sched_ctr_cute, + scheduler_counter_cute, scale, workspace_cute, stream, @@ -4556,8 +4522,8 @@ def chunk_gdn_bwd_sm100( d_final_state=None, work_items=None, work_count=None, - sched_ctr=None, - sched_all=None, + scheduler_counter=None, + scheduler_all=None, work_item_scratch=None, order_in_prologue: bool = False, log_gate: bool = False, @@ -4576,8 +4542,9 @@ def chunk_gdn_bwd_sm100( ``state_checkpoints`` row 0 like every other chunk (the forward writes the initial state there); ``d_initial_state`` (fp32) then also receives the initial-state gradient. The two go together. ``state_checkpoints`` is - always the PLAIN per-chunk checkpoint series. All tensors are contiguous, - DLPack-compatible CUDA tensors on the same device. + always the PLAIN per-chunk checkpoint series. All tensors are DLPack-compatible + CUDA tensors on the same device with a stride-1 innermost dim (outer + strides are runtime arguments). Compile-cache-and-replay. Args: @@ -4595,10 +4562,10 @@ def chunk_gdn_bwd_sm100( forward states from the prefill kernel's checkpoint output (``checkpoint_every_n_tokens=B_T``) dq/dk/dv: pre-allocated output gradients, shaped/typed like q/k/v at HO heads - dgate: pre-allocated ``(total_tokens, HO)`` float32 gate gradient - (``safe_gate`` leaves it in the transformed gate space) - dbeta: pre-allocated ``(total_tokens, HO)`` beta gradient; float32, or - io dtype and wrt the raw logits under ``use_beta_sigmoid`` + dgate/dbeta: pre-allocated ``(total_tokens, HO)`` float32 gate/beta + gradients (``safe_gate`` leaves dgate in the transformed gate + space; ``use_beta_sigmoid`` leaves dbeta in post-sigmoid space, + io dtype) cu_seqlens: ``(num_seqs + 1,)`` int32 use_initial_state: the forward ran with an initial state, so chunk 0 has an entering state to load from ``state_checkpoints`` row 0 @@ -4609,8 +4576,8 @@ def chunk_gdn_bwd_sm100( use_beta_sigmoid: ``beta`` holds logits; sigmoid in-kernel work_items: ``(max_items, 8)`` int32 work-item table from ``common/split_k.py`` (REQUIRED; an uncut table row is the whole - (b, h) sequence). Each item computes chunks ``[wstart, cend)`` - backward and writes gradients only for ``[wstart, wend)``. + (b, h) sequence). Each item computes chunks ``[write_start, compute_end)`` + backward and writes gradients only for ``[write_start, write_end)``. work_count: ``(1,)`` int32 device-side item count (REQUIRED) workspace: ``(>= tensormap_workspace_bytes(module, B) // 8,)`` int64, 128-byte aligned; holds the per-(b,h) TMA descriptors @@ -4624,7 +4591,9 @@ def chunk_gdn_bwd_sm100( B = cu_seqlens.shape[0] - 1 is_GQA = HQ >= HV if work_items is None or work_count is None: - raise ValueError("work_items/work_count are required (the split-table stage builds them for every launch)") + raise ValueError( + "work_items/work_count are required (built by the split-table stage, or by an order-generating prologue when work_item_scratch is None)" + ) io_dtype = get_dtype(q.dtype) for name, hh in (("HQ", HQ), ("HK", HK), ("HV", HV)): if HO % hh != 0: @@ -4632,11 +4601,11 @@ def chunk_gdn_bwd_sm100( cu_stream = cuda.CUstream(int(stream)) - dyn_sched = sched_ctr is not None + dynamic_scheduling = scheduler_counter is not None run_order = bool(order_in_prologue) order_gen = run_order and work_item_scratch is None - if run_order and sched_all is None: - raise ValueError("order_in_prologue requires sched_all (the prologue zeroes both consumers' sched rings)") + if run_order and scheduler_all is None: + raise ValueError("order_in_prologue requires scheduler_all (the prologue zeroes both consumers' sched rings)") if safe_gate and (a_log is None or dt_bias is None): raise ValueError("safe_gate requires a_log and dt_bias") if not safe_gate: @@ -4655,7 +4624,7 @@ def chunk_gdn_bwd_sm100( log_gate, safe_gate, use_beta_sigmoid, - dyn_sched, + dynamic_scheduling, run_order, order_gen, ) @@ -4673,9 +4642,9 @@ def chunk_gdn_bwd_sm100( work_items_cute = from_dlpack(work_items, assumed_align=16) work_items_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1), divisibility=1) work_count_cute = from_dlpack(work_count, assumed_align=4).mark_layout_dynamic() - sched_ctr_cute = None - if dyn_sched: - sched_ctr_cute = from_dlpack(sched_ctr, assumed_align=4).mark_layout_dynamic() + scheduler_counter_cute = None + if dynamic_scheduling: + scheduler_counter_cute = from_dlpack(scheduler_counter, assumed_align=4).mark_layout_dynamic() a_log_cute = from_dlpack(a_log, assumed_align=4) if a_log is not None else None dt_bias_cute = from_dlpack(dt_bias, assumed_align=4) if dt_bias is not None else None cache["compiled"] = compile( @@ -4687,7 +4656,7 @@ def chunk_gdn_bwd_sm100( log_gate=log_gate, safe_gate=safe_gate, beta_sigmoid=use_beta_sigmoid, - dyn_sched=dyn_sched, + dynamic_scheduling=dynamic_scheduling, num_sm=multiprocessor_count(current_device()), h_q=HQ, h_k=HK, @@ -4710,7 +4679,7 @@ def chunk_gdn_bwd_sm100( dstate_in_cute=dstate_in_cute, work_items_cute=work_items_cute, work_count_cute=work_count_cute, - sched_ctr_cute=sched_ctr_cute, + scheduler_counter_cute=scheduler_counter_cute, scale=scale, workspace_cute=workspace_cute, stream=cu_stream, @@ -4728,10 +4697,10 @@ def chunk_gdn_bwd_sm100( work_count_pl = from_dlpack(work_count, assumed_align=4).mark_layout_dynamic() work_items_pl = from_dlpack(work_items, assumed_align=16) work_items_pl.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1), divisibility=1) - sched_all_pl = None + scheduler_all_pl = None if run_order: - sched_all_pl = from_dlpack(sched_all, assumed_align=4).mark_layout_dynamic() - ws_pl = from_dlpack(workspace, assumed_align=128).mark_layout_dynamic() + scheduler_all_pl = from_dlpack(scheduler_all, assumed_align=4).mark_layout_dynamic() + workspace_pl = from_dlpack(workspace, assumed_align=128).mark_layout_dynamic() cache["prologue"] = cute.compile( prologue, io_dtype, @@ -4750,8 +4719,8 @@ def chunk_gdn_bwd_sm100( staging_pl, work_count_pl, work_items_pl, - sched_all_pl, - ws_pl, + scheduler_all_pl, + workspace_pl, cu_stream, options="--enable-tvm-ffi", ) @@ -4768,7 +4737,7 @@ def chunk_gdn_bwd_sm100( work_item_scratch if (run_order and not order_gen) else None, work_count, work_items, - sched_all if run_order else None, + scheduler_all if run_order else None, workspace, cu_stream, ) @@ -4791,7 +4760,7 @@ def chunk_gdn_bwd_sm100( d_final_state, work_items, work_count, - sched_ctr, + scheduler_counter, scale, workspace, cu_stream, @@ -4818,8 +4787,8 @@ def run_bwd( d_final_state, work_items, work_count, - sched_ctr, - sched_all, + scheduler_counter, + scheduler_all, work_item_scratch, tensormap_workspace, scale, @@ -4844,7 +4813,7 @@ def run_bwd( work_item_scratch, work_count, work_items, - sched_all, + scheduler_all, tensormap_workspace, cu_stream, ) @@ -4867,7 +4836,7 @@ def run_bwd( d_final_state, work_items, work_count, - sched_ctr, + scheduler_counter, scale, tensormap_workspace, cu_stream, diff --git a/python/cudnn/linear_attention/frost/kernel/gdn_prefill_config.py b/python/cudnn/linear_attention/frost/kernel/gdn_prefill_config.py index 0860db59f..175c94f56 100644 --- a/python/cudnn/linear_attention/frost/kernel/gdn_prefill_config.py +++ b/python/cudnn/linear_attention/frost/kernel/gdn_prefill_config.py @@ -39,7 +39,7 @@ class Cfg: COMPUTE_GROUP_1_WARP_IDS: Tuple[int, ...] = (4, 5, 6, 7) # kv_decay_v / v-k*state / epi ops LOAD_GATE_BETA_WARP_ID: int = 8 # gate/beta chunk loads + TMEM lifecycle TMA_QKV_WARP_ID: int = 9 - MMA_WARP_ID: int = 10 # sole tcgen05 issuer: fused KK/QK pairs + KS/QS/U/QKV/KV per chunk + TCGEN05_MMA_WARP_ID: int = 10 # sole tcgen05 issuer: fused KK/QK pairs + KS/QS/U/QKV/KV per chunk EPILOGUE_WARP_ID: int = 11 # --- register split --- @@ -52,7 +52,7 @@ class Cfg: CLUSTER_SHAPE_MNK: Tuple[int, int, int] = (1, 1, 1) # --- SMEM stage counts --- - SMEM_SCHED_STAGES: int = 2 + SMEM_SCHEDULER_STAGES: int = 2 SMEM_KQ_STAGES: int = 4 SMEM_V_STAGES: int = 2 SMEM_T_INV_STAGES: int = 3 diff --git a/python/cudnn/linear_attention/frost/kernel/gdn_prefill_f16.py b/python/cudnn/linear_attention/frost/kernel/gdn_prefill_f16.py index 1e6aa5bcc..cc214209e 100644 --- a/python/cudnn/linear_attention/frost/kernel/gdn_prefill_f16.py +++ b/python/cudnn/linear_attention/frost/kernel/gdn_prefill_f16.py @@ -58,20 +58,20 @@ O store 16384 2 checkpoint staging DK*DV*2 1 <-- enable_checkpoints only cumsumlog / cumprod / Beta 256 3 - sched ticket ring 4 2 <-- dyn_sched publish ring + scheduler ticket ring 4 2 <-- dynamic_scheduling publish ring TMEM layout (512 columns): Buffer Cols state 128 <-- DKxDV fp32 = 128x128x4B Q*state / O acc 64 <-- BTxDV fp32 accumulator - state inp 64 <-- fp16 state staging (GEMMs 3/4 A operand) + state input 64 <-- fp16 state staging (GEMMs 3/4 A operand) cg0 shared acc 128 <-- 2-stage ring: KK0/KK1 then QK0/QK1 cg1 shared acc 64 <-- 1-stage ring: KS then U - Y + U input / decayed-U inp 64 <-- slot 0 = Y then U input, slot 1 = decayed U (b16) + Y + U input / decayed-U input 64 <-- slot 0 = Y then U input, slot 1 = decayed U (b16) Warp assignments (12 warps = 384 threads): - warps 0-3 : compute group 0 - T-pairwise x2, KK_epi x2, pair inverse, - A_epi x2 + warps 0-3 : compute group 0 - T-pairwise x2, KK epilogue x2, pair inverse, + A epilogue x2 warps 4-7 : compute group 1 - state restage/rescale, Y = V - K*state, state*Q_epi, U_epi, QKV_epilogue warp 8 : Gate/Beta loads @@ -95,8 +95,7 @@ from cutlass.cutlass_dsl import min from ..common.thd import emit_checkpoint_seq_descs, emit_seq_descs, TENSOR_MAP_QWORDS -from ..common.split_k import ORDER_CAPACITY, ORDER_ELEMS, ORDER_THREADS, decode_work_item, order_body -from ..common.elementwise import softplus +from ..common.split_k import ORDER_CAPACITY, ORDER_ELEMENTS, ORDER_THREADS, decode_work_item, order_body from ..common.host import get_dtype from cudnn.frost.buffers import data_ptr from cudnn.frost.device import current_device, multiprocessor_count @@ -110,8 +109,8 @@ ) from cudnn.frost.tile_dsl.handles import MmaDesc, SmemTile, tma_slice_runtime_desc from cudnn.frost.tile_dsl.mma import mma_ss, mma_step_k8, mma_ts_step, mma_step -from cudnn.frost.tile_dsl.pointwise import fadd2, fp32_to_fp16, f16x2_to_f32, fmul2, opaque_f32_zero, sub_f16x2 -from cudnn.frost.tile_dsl.swizzle import swizzle_lin_128b, swizzle_xor_128b +from cudnn.frost.tile_dsl.pointwise import f16x2_to_f32, fadd2, fmul2, fp32_to_fp16, opaque_f32_zero, sigmoid, softplus, softplus2, sub_f16x2 +from cudnn.frost.tile_dsl.swizzle import swizzle_xor_128b from cudnn.frost.tile_dsl.tma import ( tma_load_tile, tma_store_tile, @@ -123,13 +122,7 @@ class GdnBars(NamedTuple): - """GDN pipeline mbarrier inventory. - - Every pipeline is a ``_ready``/``_done`` MBarrier pair over one ring: a - slot is acquired for filling by waiting ``_done`` and committed by - arriving ``_ready``; the reading side waits ``_ready`` and releases the - slot by arriving ``_done``. - """ + """Every inter-warp handoff as an ``MBarrier`` over its ring.""" mb_kq_ready: MBarrier mb_kq_done: MBarrier @@ -142,17 +135,16 @@ class GdnBars(NamedTuple): mb_beta_done: MBarrier mb_state_acc_ready: MBarrier - mb_state_acc_scale_done: MBarrier mb_o_acc_ready: MBarrier mb_o_final_acc_ready: MBarrier mb_o_state_scale_acc_done: MBarrier mb_cg0_acc_ready: MBarrier mb_cg0_acc_done: MBarrier - mb_state_inp_ready: MBarrier - mb_y_inp_ready: MBarrier - mb_u_inp_ready: MBarrier - mb_decay_u_inp_ready: MBarrier + mb_state_input_ready: MBarrier + mb_y_input_ready: MBarrier + mb_u_input_ready: MBarrier + mb_decay_u_input_ready: MBarrier mb_t_inv_ready: MBarrier mb_t_inv_done: MBarrier @@ -170,14 +162,14 @@ class GdnBars(NamedTuple): mb_tmem_done: MBarrier - mb_sched_ready: MBarrier - mb_sched_done: MBarrier + mb_scheduler_ready: MBarrier + mb_scheduler_done: MBarrier def make_gdn_bars(cfg) -> GdnBars: - """GdnBars factory. MUST be called from inside ``kernel`` (allocates SMEM).""" + """GdnBars factory.""" ONE_LANE = 1 - MMA_ARRIVERS = len([cfg.mma_warp_id]) + MMA_ARRIVERS = len([cfg.tcgen05_mma_warp_id]) KQ_RELEASE_SITES = 1 GATE_WARP = cfg.threads_per_warp * len([cfg.load_gate_beta_warp_id]) EPI_WARP = cfg.threads_per_warp * len([cfg.epilogue_warp_id]) @@ -198,7 +190,6 @@ def alloc(n): mb_beta_ready=MBarrier(alloc(cfg.smem_beta_stages), stages=cfg.smem_beta_stages, init_count=GATE_WARP, producer=Producer.THREAD), mb_beta_done=MBarrier(alloc(cfg.smem_beta_stages), stages=cfg.smem_beta_stages, init_count=CG0_THREADS, producer=Producer.THREAD), mb_state_acc_ready=MBarrier(alloc(cfg.tmem_state_acc_stages), stages=cfg.tmem_state_acc_stages, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), - mb_state_acc_scale_done=MBarrier(alloc(cfg.tmem_state_acc_stages), stages=cfg.tmem_state_acc_stages, init_count=CG1_THREADS, producer=Producer.THREAD), mb_o_acc_ready=MBarrier(alloc(cfg.tmem_q_state_acc_stages), stages=cfg.tmem_q_state_acc_stages, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), mb_o_final_acc_ready=MBarrier( alloc(cfg.tmem_q_state_acc_stages), stages=cfg.tmem_q_state_acc_stages, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT @@ -208,10 +199,10 @@ def alloc(n): ), mb_cg0_acc_ready=MBarrier(alloc(cfg.tmem_cg0_acc_stages), stages=cfg.tmem_cg0_acc_stages, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), mb_cg0_acc_done=MBarrier(alloc(cfg.tmem_cg0_acc_stages), stages=cfg.tmem_cg0_acc_stages, init_count=CG0_THREADS // 2, producer=Producer.THREAD), - mb_state_inp_ready=MBarrier(alloc(cfg.tmem_state_inp_stages), stages=cfg.tmem_state_inp_stages, init_count=CG1_THREADS, producer=Producer.THREAD), - mb_y_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), - mb_u_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), - mb_decay_u_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_state_input_ready=MBarrier(alloc(cfg.tmem_state_input_stages), stages=cfg.tmem_state_input_stages, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_y_input_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_u_input_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_decay_u_input_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), mb_t_inv_ready=MBarrier(alloc(cfg.smem_t_inv_stages), stages=cfg.smem_t_inv_stages, init_count=CG0_THREADS, producer=Producer.THREAD), mb_t_inv_done=MBarrier(alloc(cfg.smem_t_inv_stages), stages=cfg.smem_t_inv_stages, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), mb_a_ready=MBarrier(alloc(cfg.smem_a_stages), stages=cfg.smem_a_stages, init_count=CG0_THREADS // 2, producer=Producer.THREAD), @@ -225,28 +216,19 @@ def alloc(n): ), mb_checkpoint_tmastg_done=MBarrier(alloc(cfg.smem_checkpoint_stages), stages=cfg.smem_checkpoint_stages, init_count=EPI_WARP, producer=Producer.THREAD), mb_tmem_done=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), - mb_sched_ready=MBarrier(alloc(cfg.sched_stages), stages=cfg.sched_stages, init_count=1, producer=Producer.THREAD), - mb_sched_done=MBarrier(alloc(cfg.sched_stages), stages=cfg.sched_stages, init_count=11, producer=Producer.THREAD), + mb_scheduler_ready=MBarrier(alloc(cfg.scheduler_stages), stages=cfg.scheduler_stages, init_count=1, producer=Producer.THREAD), + mb_scheduler_done=MBarrier(alloc(cfg.scheduler_stages), stages=cfg.scheduler_stages, init_count=11, producer=Producer.THREAD), ) @cute.jit -def invert_diagonal_NxN(cfg, base_int, d, tidx, N: int = 8): +def invert_diagonal_NxN(cfg, base, d_idx, tidx, N: int = 8): """Gauss-Jordan inversion of one diagonal NxN block in-place (f16 SMEM).""" tidx_in_group = tidx % N BT = cfg.b_t - row_lin_base = (d * N + tidx_in_group) * BT + d * N - row_phys = swizzle_lin_128b(row_lin_base, row_stride_log2=6) - row_ptr = ( - cute.make_ptr( - cfg.io_dtype, - base_int, - mem_space=cute.AddressSpace.smem, - assumed_align=cfg.buffer_align_bytes, - ) - + row_phys - ) + row_coord = d_idx * N + tidx_in_group + row_ptr = base + row_coord * BT + swizzle_xor_128b(row_coord, d_idx * N) row = [(row_ptr + j).load().to(cutlass.Float32) for j in range(N)] for i in cutlass.range_constexpr(N): @@ -263,229 +245,146 @@ def invert_diagonal_NxN(cfg, base_int, d, tidx, N: int = 8): @cute.jit -def blockwise_diagonal_8x8_to_16x16(cfg, base_int, d0, lane_id): +def blockwise_diagonal_8x8_to_16x16(cfg, base, d_idx, lane_idx): """Off-diagonal correction 8x8 -> 16x16 (C <- -D^{-1} C A^{-1}).""" - bpe = cfg.io_dtype.width // 8 - ldsm_x1_lane_off = (lane_id % 8) * 64 - d = nvvm.ldmatrix( - cutlass.inttoptr( - base_int + swizzle_lin_128b((d0 + 8) * 64 + d0 + 8 + ldsm_x1_lane_off, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16 - ), - 1, - nvvm.MMALayout.ROW, - ) - c = nvvm.ldmatrix( - cutlass.inttoptr( - base_int + swizzle_lin_128b((d0 + 8) * 64 + d0 + ldsm_x1_lane_off, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16 - ), - 1, - nvvm.MMALayout.COL, - ) - - # ---- T = -(D^{-1} @ C) ------------------------------------------------------- + BT = cfg.b_t + row_lo = d_idx + lane_idx % 8 + row_hi = row_lo + 8 + off_d_inv = row_hi * BT + swizzle_xor_128b(row_hi, d_idx + 8) + off_c = row_hi * BT + swizzle_xor_128b(row_hi, d_idx) + off_a_inv = row_lo * BT + swizzle_xor_128b(row_lo, d_idx) + d_inv_frag = nvvm.ldmatrix(base + off_d_inv, 1, nvvm.MMALayout.ROW) + c_frag = nvvm.ldmatrix(base + off_c, 1, nvvm.MMALayout.COL) + + # ---- T = -(D^-1 @ C) ------------------------------------------------------------- c_regs = cutlass.Array(cutlass.Float32, 4, alignment=16, space=cutlass.AddressSpace.rmem) for i in cutlass.range_constexpr(4): c_regs[i] = cutlass.Float32(0.0) - mma_step_k8(c_regs, [d, d], [c], k_step=0, M=16, N=8, ab_dtype=cfg.io_dtype) + mma_step_k8(c_regs, [d_inv_frag, d_inv_frag], [c_frag], k_step=0, M=16, N=8, ab_dtype=cfg.io_dtype) for i in cutlass.range_constexpr(4): c_regs[i] = -c_regs[i] a_pack = [fp32_to_fp16(c_regs[2 * j], c_regs[2 * j + 1], dtype=cfg.io_dtype) for j in range(2)] - # ---- C = T @ A^{-1} ---------------------------------------------------------- - ai_frag = nvvm.ldmatrix( - cutlass.inttoptr(base_int + swizzle_lin_128b(d0 * 64 + d0 + ldsm_x1_lane_off, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16), - 1, - nvvm.MMALayout.COL, - ) + # ---- C = T @ A^-1 ---------------------------------------------------------------- + a_inv_frag = nvvm.ldmatrix(base + off_a_inv, 1, nvvm.MMALayout.COL) o_regs = cutlass.Array(cutlass.Float32, 4, alignment=16, space=cutlass.AddressSpace.rmem) for i in cutlass.range_constexpr(4): o_regs[i] = cutlass.Float32(0.0) - mma_step_k8(o_regs, a_pack, [ai_frag], k_step=0, M=16, N=8, ab_dtype=cfg.io_dtype) + mma_step_k8(o_regs, a_pack, [a_inv_frag], k_step=0, M=16, N=8, ab_dtype=cfg.io_dtype) o_pack = fp32_to_fp16(o_regs[0], o_regs[1], dtype=cfg.io_dtype) - # ---- store corrected C ------------------------------------------------------- - nvvm.stmatrix( - cutlass.inttoptr( - base_int + swizzle_lin_128b((d0 + 8) * 64 + d0 + ldsm_x1_lane_off, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16 - ), - o_pack, - nvvm.MMALayout.ROW, - ) + # ---- store corrected C ----------------------------------------------------------- + nvvm.stmatrix(base + off_c, o_pack, nvvm.MMALayout.ROW) @cute.jit -def blockwise_diagonal_16x16_to_32x32(cfg, base_int, d0, lane_id): +def blockwise_diagonal_16x16_to_32x32(cfg, base, d_idx, lane_idx): """Off-diagonal correction 16x16 -> 32x32.""" - bpe = cfg.io_dtype.width // 8 - ldsm_x4_lane_off = (lane_id % 16) * 64 + (lane_id // 16) * 8 - d = list( - nvvm.ldmatrix( - cutlass.inttoptr( - base_int + swizzle_lin_128b((d0 + 16) * 64 + d0 + 16 + ldsm_x4_lane_off, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16 - ), - 4, - nvvm.MMALayout.ROW, - ) - ) - c = list( - nvvm.ldmatrix( - cutlass.inttoptr( - base_int + swizzle_lin_128b((d0 + 16) * 64 + d0 + ldsm_x4_lane_off, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16 - ), - 4, - nvvm.MMALayout.COL, - ) - ) - - # ---- T = -(D^{-1} @ C) ------------------------------------------------------- + BT = cfg.b_t + lane_row = lane_idx % 16 + lane_col = (lane_idx // 16) * 8 + row_lo = d_idx + lane_row + row_hi = row_lo + 16 + off_d_inv = row_hi * BT + swizzle_xor_128b(row_hi, d_idx + 16 + lane_col) + off_c = row_hi * BT + swizzle_xor_128b(row_hi, d_idx + lane_col) + off_a_inv = row_lo * BT + swizzle_xor_128b(row_lo, d_idx + lane_col) + d_inv_frags = list(nvvm.ldmatrix(base + off_d_inv, 4, nvvm.MMALayout.ROW)) + c_frags = list(nvvm.ldmatrix(base + off_c, 4, nvvm.MMALayout.COL)) + + # ---- T = -(D^-1 @ C) ------------------------------------------------------------- c_regs = cutlass.Array(cutlass.Float32, 8, alignment=16, space=cutlass.AddressSpace.rmem) for i in cutlass.range_constexpr(8): c_regs[i] = cutlass.Float32(0.0) - mma_step(c_regs, d, c, k_step=0, M=16, N=16, ab_dtype=cfg.io_dtype) + mma_step(c_regs, d_inv_frags, c_frags, k_step=0, M=16, N=16, ab_dtype=cfg.io_dtype) for i in cutlass.range_constexpr(8): c_regs[i] = -c_regs[i] a_pack = [fp32_to_fp16(c_regs[2 * j], c_regs[2 * j + 1], dtype=cfg.io_dtype) for j in range(4)] - # ---- C = T @ A^{-1} ---------------------------------------------------------- - ai_frag = list( - nvvm.ldmatrix( - cutlass.inttoptr( - base_int + swizzle_lin_128b(d0 * 64 + d0 + ldsm_x4_lane_off, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16 - ), - 4, - nvvm.MMALayout.COL, - ) - ) + # ---- C = T @ A^-1 ---------------------------------------------------------------- + a_inv_frags = list(nvvm.ldmatrix(base + off_a_inv, 4, nvvm.MMALayout.COL)) o_regs = cutlass.Array(cutlass.Float32, 8, alignment=16, space=cutlass.AddressSpace.rmem) for i in cutlass.range_constexpr(8): o_regs[i] = cutlass.Float32(0.0) - mma_step(o_regs, a_pack, ai_frag, k_step=0, M=16, N=16, ab_dtype=cfg.io_dtype) + mma_step(o_regs, a_pack, a_inv_frags, k_step=0, M=16, N=16, ab_dtype=cfg.io_dtype) o_pack = [fp32_to_fp16(o_regs[2 * j], o_regs[2 * j + 1], dtype=cfg.io_dtype) for j in range(4)] - # ---- store corrected C ------------------------------------------------------- - nvvm.stmatrix( - cutlass.inttoptr( - base_int + swizzle_lin_128b((d0 + 16) * 64 + d0 + ldsm_x4_lane_off, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16 - ), - o_pack, - nvvm.MMALayout.ROW, - ) + # ---- store corrected C ----------------------------------------------------------- + nvvm.stmatrix(base + off_c, o_pack, nvvm.MMALayout.ROW) @cute.jit -def blockwise_diagonal_32x32_to_64x64(cfg, base_int, warp_id, lane_id): +def blockwise_diagonal_32x32_to_64x64(cfg, base, warp_id, lane_idx): """Off-diagonal correction 32x32 -> 64x64 (2 warps, one 16-row M-band each).""" band = warp_id % 2 - bpe = cfg.io_dtype.width // 8 - ldsm_x4_lane_off = (lane_id % 16) * 64 + (lane_id // 16) * 8 - a_frags = [] + BT = cfg.b_t + lane_row = lane_idx % 16 + lane_col = (lane_idx // 16) * 8 + row_d_inv = 32 + band * 16 + lane_row + d_inv_frags = [] for vs in cutlass.range_constexpr(2): - a_frags += list( - nvvm.ldmatrix( - cutlass.inttoptr( - base_int + swizzle_lin_128b((32 + band * 16) * 64 + 32 + vs * 16 + ldsm_x4_lane_off, row_stride_log2=6) * bpe, - cutlass.AddressSpace.smem, - cutlass.BFloat16, - ), - 4, - nvvm.MMALayout.ROW, - ) - ) - b_frags = [] + d_inv_frags += list(nvvm.ldmatrix(base + row_d_inv * BT + swizzle_xor_128b(row_d_inv, 32 + vs * 16 + lane_col), 4, nvvm.MMALayout.ROW)) + c_frags = [] for vs in cutlass.range_constexpr(4): - b_frags += list( - nvvm.ldmatrix( - cutlass.inttoptr( - base_int + swizzle_lin_128b((32 + (vs // 2) * 16) * 64 + (vs % 2) * 16 + ldsm_x4_lane_off, row_stride_log2=6) * bpe, - cutlass.AddressSpace.smem, - cutlass.BFloat16, - ), - 4, - nvvm.MMALayout.COL, - ) - ) + row_c = 32 + (vs // 2) * 16 + lane_row + c_frags += list(nvvm.ldmatrix(base + row_c * BT + swizzle_xor_128b(row_c, (vs % 2) * 16 + lane_col), 4, nvvm.MMALayout.COL)) - # ---- T = -(D^{-1} @ C) ------------------------------------------------------- + # ---- T = -(D^-1 @ C) ------------------------------------------------------------- c_regs = cutlass.Array(cutlass.Float32, 16, alignment=16, space=cutlass.AddressSpace.rmem) for i in cutlass.range_constexpr(16): c_regs[i] = cutlass.Float32(0.0) for ks in cutlass.range_constexpr(2): - mma_step(c_regs, a_frags, b_frags[ks * 8 : ks * 8 + 8], k_step=ks, M=16, N=32, ab_dtype=cfg.io_dtype) + mma_step(c_regs, d_inv_frags, c_frags[ks * 8 : ks * 8 + 8], k_step=ks, M=16, N=32, ab_dtype=cfg.io_dtype) for i in cutlass.range_constexpr(16): c_regs[i] = -c_regs[i] a_pack = [fp32_to_fp16(c_regs[2 * j], c_regs[2 * j + 1], dtype=cfg.io_dtype) for j in range(8)] - # ---- C = T @ A^{-1} ---------------------------------------------------------- - ai_frags = [] + # ---- C = T @ A^-1 ---------------------------------------------------------------- + a_inv_frags = [] for vs in cutlass.range_constexpr(4): - ai_frags += list( - nvvm.ldmatrix( - cutlass.inttoptr( - base_int + swizzle_lin_128b(((vs // 2) * 16) * 64 + (vs % 2) * 16 + ldsm_x4_lane_off, row_stride_log2=6) * bpe, - cutlass.AddressSpace.smem, - cutlass.BFloat16, - ), - 4, - nvvm.MMALayout.COL, - ) - ) + row_a_inv = (vs // 2) * 16 + lane_row + a_inv_frags += list(nvvm.ldmatrix(base + row_a_inv * BT + swizzle_xor_128b(row_a_inv, (vs % 2) * 16 + lane_col), 4, nvvm.MMALayout.COL)) o_regs = cutlass.Array(cutlass.Float32, 16, alignment=16, space=cutlass.AddressSpace.rmem) for i in cutlass.range_constexpr(16): o_regs[i] = cutlass.Float32(0.0) for ks in cutlass.range_constexpr(2): - mma_step(o_regs, a_pack, ai_frags[ks * 8 : ks * 8 + 8], k_step=ks, M=16, N=32, ab_dtype=cfg.io_dtype) + mma_step(o_regs, a_pack, a_inv_frags[ks * 8 : ks * 8 + 8], k_step=ks, M=16, N=32, ab_dtype=cfg.io_dtype) o_pack = [fp32_to_fp16(o_regs[2 * j], o_regs[2 * j + 1], dtype=cfg.io_dtype) for j in range(8)] - # ---- store corrected C ------------------------------------------------------- + # ---- store corrected C ----------------------------------------------------------- nvvm.barrier_cta_sync_aligned( cfg.inverse_barrier_id, thread_count=cfg.inverse_barrier_threads, ) - nvvm.stmatrix( - cutlass.inttoptr( - base_int + swizzle_lin_128b((32 + band * 16) * 64 + ldsm_x4_lane_off, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16 - ), - o_pack[0:4], - nvvm.MMALayout.ROW, - ) - nvvm.stmatrix( - cutlass.inttoptr( - base_int + swizzle_lin_128b((32 + band * 16) * 64 + 16 + ldsm_x4_lane_off, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16 - ), - o_pack[4:8], - nvvm.MMALayout.ROW, - ) - - -# ---- Dynamic tile scheduler ------------------------------------------------------ + nvvm.stmatrix(base + row_d_inv * BT + swizzle_xor_128b(row_d_inv, lane_col), o_pack[0:4], nvvm.MMALayout.ROW) + nvvm.stmatrix(base + row_d_inv * BT + swizzle_xor_128b(row_d_inv, 16 + lane_col), o_pack[4:8], nvvm.MMALayout.ROW) @cute.jit -def sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ctas): +def scheduler_publish_next(cfg, bars, sScheduler, mScheduler, scheduler_state, tile_idx, num_ctas, elect_one): """TMA-LDG-warp side: pull the next tile off the global ticket, publish it.""" - if cutlass.const_expr(cfg.dyn_sched): - bars.mb_sched_done[sched_state.idx].wait(sched_state.phase) - if nvvm.elect_sync(): - fetched = cutlass.Int32(nvvm.atomicrmw("add", mSched.iterator, cutlass.Int32(1), mem_order="relaxed", syncscope="gpu")) - sSched[sched_state.idx] = num_ctas + fetched + if cutlass.const_expr(cfg.dynamic_scheduling): + bars.mb_scheduler_done[scheduler_state.idx].wait(scheduler_state.phase) + if elect_one: + fetched = cutlass.Int32(nvvm.atomicrmw("add", mScheduler.iterator, cutlass.Int32(1), mem_order="relaxed", syncscope="gpu")) + sScheduler[scheduler_state.idx] = num_ctas + fetched nvvm.bar_warp_sync(cute.arch.FULL_MASK) - next_tile = sSched[sched_state.idx] - if nvvm.elect_sync(): - bars.mb_sched_ready[sched_state.idx].arrive() - return next_tile, advance(sched_state, cfg.sched_stages) - return tile_idx + num_ctas, sched_state + next_tile = sScheduler[scheduler_state.idx] + if elect_one: + bars.mb_scheduler_ready[scheduler_state.idx].arrive() + return next_tile, advance(scheduler_state, cfg.scheduler_stages) + return tile_idx + num_ctas, scheduler_state @cute.jit -def sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas): +def scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one): """Consumer side: read the TMA-LDG warp's published next tile.""" - if cutlass.const_expr(cfg.dyn_sched): - bars.mb_sched_ready[sched_state.idx].wait(sched_state.phase) - next_tile = sSched[sched_state.idx] - if nvvm.elect_sync(): - bars.mb_sched_done[sched_state.idx].arrive() - return next_tile, advance(sched_state, cfg.sched_stages) - return tile_idx + num_ctas, sched_state + if cutlass.const_expr(cfg.dynamic_scheduling): + bars.mb_scheduler_ready[scheduler_state.idx].wait(scheduler_state.phase) + next_tile = sScheduler[scheduler_state.idx] + if elect_one: + bars.mb_scheduler_done[scheduler_state.idx].arrive() + return next_tile, advance(scheduler_state, cfg.scheduler_stages) + return tile_idx + num_ctas, scheduler_state @cute.jit @@ -502,20 +401,22 @@ def tmastg_warp( sCheckpoint_raw, desc_o_base, desc_checkpoint_base, - sSched, + sScheduler, bars, ): """Epilogue warp role (warp 11): persistent scheduler loop issuing the per-chunk O and state-checkpoint TMA stores.""" - elect_one = nvvm.elect_sync() nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) + o_index = PipelineState.start(phase=0) - lidx = tidx % cfg.threads_per_warp - sched_state = PipelineState.start(phase=0) - tile_idx = cutlass.Int32(bidx) + scheduler_state = PipelineState.start(phase=0) + + lane_idx = tidx % cfg.threads_per_warp + elect_one = nvvm.elect_sync() + tile_idx = cutlass.Int32(bidx) bpe = cfg.io_dtype.width // 8 - elems_per_128b = 128 // bpe + elements_per_128b = 128 // bpe sO_tma = SmemTile( base=sO_raw, elems_per_stage=(cfg.o_cosize // cfg.smem_o_stages), @@ -524,11 +425,11 @@ def tmastg_warp( stride_byte_offset=0, layout=0, tma_loads_per_tile=2, - tma_granu_elems=elems_per_128b, + tma_granu_elems=elements_per_128b, tma_subtile_stride_elems=4096, ) if cutlass.const_expr(cfg.enable_checkpoints): - checkpoint_elems_per_128b = 64 + checkpoint_elements_per_128b = 64 sCheckpoint_tma = SmemTile( base=sCheckpoint_raw, elems_per_stage=(cfg.checkpoint_cosize // cfg.smem_checkpoint_stages), @@ -536,18 +437,20 @@ def tmastg_warp( leading_byte_offset=0, stride_byte_offset=0, layout=0, - tma_loads_per_tile=cfg.d_v // checkpoint_elems_per_128b, - tma_granu_elems=checkpoint_elems_per_128b, - tma_subtile_stride_elems=cfg.d_k * checkpoint_elems_per_128b, + tma_loads_per_tile=cfg.d_v // checkpoint_elements_per_128b, + tma_granu_elems=checkpoint_elements_per_128b, + tma_subtile_stride_elems=cfg.d_k * checkpoint_elements_per_128b, ) checkpoint_store_cnt = cutlass.Int32(0) - ckpt_chunks = checkpoint_every_n_tokens // cutlass.Int32(cfg.b_t) + checkpoint_chunks = checkpoint_every_n_tokens // cutlass.Int32(cfg.b_t) heads_out = cutlass.Int32(cfg.n_heads_out) desc_qwords = cutlass.Int32(TENSOR_MAP_QWORDS) while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) - n_local = wend - cstart + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) + n_local = write_end - compute_start head_o = head_idx slot = batch_idx * desc_qwords @@ -556,14 +459,14 @@ def tmastg_warp( tma_tensormap_acquire(desc_o_slot) if cutlass.const_expr(cfg.enable_checkpoints): desc_checkpoint_slot = (desc_checkpoint_base + slot).tospace(cutlass.AddressSpace.generic) - checkpoint_coord = (wstart + ckpt_chunks - cutlass.Int32(1)) // ckpt_chunks - checkpoint_mod = (cstart + cutlass.Int32(1)) % ckpt_chunks + checkpoint_coord = (write_start + checkpoint_chunks - cutlass.Int32(1)) // checkpoint_chunks + checkpoint_mod = (compute_start + cutlass.Int32(1)) % checkpoint_chunks if elect_one: tma_tensormap_acquire(desc_checkpoint_slot) if n_local > 0: if cutlass.const_expr(cfg.enable_checkpoints): - if wstart == 0: + if write_start == 0: checkpoint_stage = checkpoint_store_cnt % cfg.smem_checkpoint_stages checkpoint_phase = (checkpoint_store_cnt // cfg.smem_checkpoint_stages) & cutlass.Int32(1) bars.mb_checkpoint_tmastg_ready[checkpoint_stage].wait(checkpoint_phase) @@ -575,14 +478,14 @@ def tmastg_warp( checkpoint_coord += 1 checkpoint_store_cnt = checkpoint_store_cnt + 1 for local_idx in cutlass.range(n_local): - chunk_idx = cstart + local_idx + chunk_idx = compute_start + local_idx did_o = cutlass.Int32(0) o_idx = o_index.idx bars.mb_o_tmastg_ready[o_idx].wait(o_index.phase) o_index = advance(o_index, cfg.smem_o_stages) - if chunk_idx >= wstart and chunk_idx < wend: + if chunk_idx >= write_start and chunk_idx < write_end: tok_coord = chunk_idx * cutlass.Int32(cfg.b_t) o_slice = tma_slice_runtime_desc(desc_o_slot, cutlass.Int32(0), head_o, tok_coord) tma_store_tile(sO_tma[o_idx], o_slice, acquire=False) @@ -593,7 +496,7 @@ def tmastg_warp( if cutlass.const_expr(cfg.enable_checkpoints): checkpoint_stage = checkpoint_store_cnt % cfg.smem_checkpoint_stages checkpoint_phase = (checkpoint_store_cnt // cfg.smem_checkpoint_stages) & cutlass.Int32(1) - if chunk_idx >= wstart - 1 and chunk_idx < wend - 1: + if chunk_idx >= write_start - 1 and chunk_idx < write_end - 1: if checkpoint_mod == 0: bars.mb_checkpoint_tmastg_ready[checkpoint_stage].wait(checkpoint_phase) checkpoint_slice = tma_slice_runtime_desc(desc_checkpoint_slot, cutlass.Int32(0), cutlass.Int32(0), checkpoint_coord, head_o) @@ -602,7 +505,7 @@ def tmastg_warp( checkpoint_coord += 1 did_checkpoint = cutlass.Int32(1) checkpoint_mod = checkpoint_mod + cutlass.Int32(1) - checkpoint_mod = cutlass.Int32(0) if checkpoint_mod == ckpt_chunks else checkpoint_mod + checkpoint_mod = cutlass.Int32(0) if checkpoint_mod == checkpoint_chunks else checkpoint_mod if cutlass.const_expr(cfg.enable_checkpoints): if did_o == 1 and did_checkpoint == 1: @@ -625,7 +528,7 @@ def tmastg_warp( tma_store_wait(0) bars.mb_o_tmastg_done[o_idx].arrive() - tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one) @cute.jit @@ -644,33 +547,37 @@ def gate_beta_warp( sCumsumlog, sCumprod, sBeta, - sSched, + sScheduler, bars, ): """Gate/Beta producer (warp 8): persistent scheduler loop + the cumsum/cumprod/Beta chunk loads.""" nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) + elect_one = nvvm.elect_sync() + gate_index = PipelineState.start(phase=1) beta_index = PipelineState.start(phase=1) - lidx = tidx % cfg.threads_per_warp + scheduler_state = PipelineState.start(phase=0) + + lane_idx = tidx % cfg.threads_per_warp - a_l2 = cutlass.Float32(0.0) + a = cutlass.Float32(0.0) bias = cutlass.Float32(0.0) - sched_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) - n_local = wend - cstart + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) + n_local = write_end - compute_start if cutlass.const_expr(cfg.safe_gate): if n_local > 0: - # per-head transform constants, fixed for the whole tile - a_l2 = -cute.math.exp2(mA_log[head_idx].to(cutlass.Float32) * cutlass.Float32(RCP_LN2), fastmath=True) * cutlass.Float32(RCP_LN2) + a = -cute.math.exp2(mA_log[head_idx].to(cutlass.Float32) * cutlass.Float32(RCP_LN2), fastmath=True) * cutlass.Float32(RCP_LN2) bias = mDt_bias[head_idx].to(cutlass.Float32) if n_local > 0: for local_idx in cutlass.range(n_local): - # ---- Gate load: GMEM -> SMEM (OOB neutral) ----------------------- - chunk_idx = cstart + local_idx + # ---- Gate load: GMEM -> SMEM (OOB neutral: 1.0 -> log2 = 0.0) -------- + chunk_idx = compute_start + local_idx n_cols = cfg.b_t // cfg.threads_per_warp chunk_offset = batch_start + chunk_idx * cfg.b_t gGateSeq = mGate[None, head_idx] @@ -680,30 +587,32 @@ def gate_beta_warp( gate_phase = gate_index.phase gate_index = advance(gate_index, cfg.smem_gate_stages) - pos_valid = [None] * n_cols - gate_vals = [cutlass.Float32(0.0)] * n_cols oob_neutral = cutlass.Float32(0.0) if cutlass.const_expr(cfg.log_gate) else cutlass.Float32(1.0) - for col in cutlass.range_constexpr(n_cols): - tok = chunk_offset + lidx + col * cfg.threads_per_warp - pos_valid[col] = cute.elem_less(tok, batch_end) - tok_clamped = min(tok, batch_end - 1) - gate_vals[col] = gGateSeq[tok_clamped] if pos_valid[col] else oob_neutral + toks = [chunk_offset + lane_idx + col * cfg.threads_per_warp for col in range(n_cols)] + pos_valid = [tok < batch_end for tok in toks] + gate_vals = [gGateSeq[min(tok, batch_end - 1)] if valid else oob_neutral for tok, valid in zip(toks, pos_valid)] if cutlass.const_expr(cfg.safe_gate): - # raw logits -> log2-domain decay: a_l2 * softplus(g + bias) (split-K scan arithmetic) - for col in cutlass.range_constexpr(n_cols): - contrib = a_l2 * softplus(gate_vals[col] + bias) - gate_vals[col] = contrib if pos_valid[col] else cutlass.Float32(0.0) + for col in cutlass.range_constexpr(0, n_cols, 2): + biased_lo, biased_hi = fadd2(gate_vals[col], gate_vals[col + 1], bias, bias) + sp_lo, sp_hi = softplus2(biased_lo, biased_hi) + contrib_lo, contrib_hi = fmul2(sp_lo, sp_hi, a, a) + gate_vals[col] = contrib_lo if pos_valid[col] else cutlass.Float32(0.0) + gate_vals[col + 1] = contrib_hi if pos_valid[col + 1] else cutlass.Float32(0.0) elif cutlass.const_expr(cfg.log_gate): - for col in cutlass.range_constexpr(n_cols): - gate_vals[col] = gate_vals[col] * cutlass.Float32(RCP_LN2) + rcp_ln2 = opaque_f32_zero() + cutlass.Float32(RCP_LN2) + for col in cutlass.range_constexpr(0, n_cols, 2): + gate_vals[col], gate_vals[col + 1] = fmul2(gate_vals[col], gate_vals[col + 1], rcp_ln2, rcp_ln2) else: - for col in cutlass.range_constexpr(n_cols): - gate_vals[col] = cute.math.log2(gate_vals[col] + 1e-10, fastmath=True) + floor = cutlass.Float32(1e-10) + for col in cutlass.range_constexpr(0, n_cols, 2): + shifted_lo, shifted_hi = fadd2(gate_vals[col], gate_vals[col + 1], floor, floor) + gate_vals[col] = cute.math.log2(shifted_lo, fastmath=True) + gate_vals[col + 1] = cute.math.log2(shifted_hi, fastmath=True) for offset in [1, 2, 4, 8, 16]: for col in cutlass.range_constexpr(n_cols): n = nvvm.shfl_sync(0xFFFFFFFF, gate_vals[col], offset, 0, kind=nvvm.Shfl.UP) - if lidx >= offset: + if lane_idx >= offset: gate_vals[col] = gate_vals[col] + n for col in cutlass.range_constexpr(1, n_cols): last_v = nvvm.shfl_sync( @@ -717,35 +626,33 @@ def gate_beta_warp( bars.mb_gate_done[gate_idx].wait(gate_phase) for col in cutlass.range_constexpr(n_cols): - pos = lidx + col * cfg.threads_per_warp + pos = lane_idx + col * cfg.threads_per_warp sCumsumlog[pos, 0, gate_idx] = gate_vals[col] sCumprod[pos, 0, gate_idx] = cute.math.exp2(gate_vals[col], fastmath=True) bars.mb_gate_ready[gate_idx].arrive() - # ---- Beta load: GMEM -> SMEM (per-element cp.async) -------------------------- + # ---- Beta load: GMEM -> SMEM (per-element cp.async) ------------------ beta_idx = beta_index.idx bars.mb_beta_done[beta_idx].wait(beta_index.phase) beta_index = advance(beta_index, cfg.smem_beta_stages) if cutlass.const_expr(cfg.beta_sigmoid): - # io-dtype logits -> sigmoid (tanh identity) -> fp32 SMEM for col in cutlass.range_constexpr(n_cols): - pos = lidx + col * cfg.threads_per_warp + pos = lane_idx + col * cfg.threads_per_warp beta_value = cutlass.Float32(0.0) if pos_valid[col]: beta_value = gBeta[pos].to(cutlass.Float32) - half = cutlass.Float32(0.5) - beta_value = (cute.math.tanh(beta_value * half, approx=True) * half + half).to(mBeta.element_type).to(cutlass.Float32) + beta_value = sigmoid(beta_value).to(mBeta.element_type).to(cutlass.Float32) sBeta[pos, 0, beta_idx] = beta_value bars.mb_beta_ready[beta_idx].arrive() else: for col in cutlass.range_constexpr(n_cols): - pos = lidx + col * cfg.threads_per_warp + pos = lane_idx + col * cfg.threads_per_warp src = gBeta.iterator + gBeta.layout((pos,)) dst = sBeta.iterator + sBeta.layout((pos, 0, beta_idx)) cp_size = cutlass.Int32(4) * cutlass.Int32(pos_valid[col]) nvvm.cp_async_shared_global(dst, src, 4, nvvm.LoadCacheModifier.CA, cp_size=cp_size) nvvm.cp_async_mbarrier_arrive(bars.mb_beta_ready[beta_idx].smem_ptr, noinc=True) - tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one) for _ in range(cfg.smem_gate_stages): bars.mb_gate_done[gate_index.idx].wait(gate_index.phase) @@ -756,7 +663,7 @@ def gate_beta_warp( @cute.jit -def mma_warp( +def tcgen05_mma_warp( cfg, total_tiles, bidx, @@ -768,34 +675,32 @@ def mma_warp( sKQ_trans, sTinv, sA, - sSched, + sScheduler, bars, ): """MMA issuer role (warp 10): persistent scheduler loop issuing every tcgen05 GEMM.""" - elect_one = nvvm.elect_sync() - nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) + o_acc_index = PipelineState.start(phase=1) o_state_scale_index = PipelineState.start(phase=0) kv_acc_index = PipelineState.start(phase=1) kq_index = PipelineState.start(phase=0) cg0_acc_index = PipelineState.start(phase=1) - kq_fused_index = PipelineState.start(phase=0) + kq_cg0_index = PipelineState.start(phase=0) tinv_index = PipelineState.start(phase=0) a_index = PipelineState.start(phase=0) - state_inp_index = PipelineState.start(phase=0) - y_inp_ready = PipelineState.start(phase=0) - u_inp_ready = PipelineState.start(phase=0) - decay_u_inp_ready = PipelineState.start(phase=0) + state_input_index = PipelineState.start(phase=0) + y_input_ready = PipelineState.start(phase=0) + u_input_ready = PipelineState.start(phase=0) + decay_u_input_ready = PipelineState.start(phase=0) + + elect_one = nvvm.elect_sync() nvvm.tcgen05_alloc(tmem_base_slot, cutlass.Int32(512), group=nvvm.CTAGroup.CTA_1) - nvvm.barrier_cta_sync_aligned( - cfg.tmem_alloc_barrier_id, - thread_count=cfg.tmem_alloc_barrier_threads, - ) + nvvm.barrier_cta_sync_aligned(cfg.tmem_lifecycle_barrier_id, thread_count=cfg.tmem_user_threads) - # ---- chunk-invariant GEMM descriptors ---------------------------------------- + # ---- chunk-invariant GEMM descriptors -------------------------------------------- idesc_qk = nvvm.Tcgen05InstrDesc.build( c_dtype=cutlass.Float32, a_dtype=cfg.io_dtype, @@ -803,7 +708,7 @@ def mma_warp( n_dim=cfg.b_t, m_dim=2 * cfg.b_t, ) - bmm_qk_desc = MmaDesc( + bmm_kq_k_desc = MmaDesc( M=2 * cfg.b_t, N=cfg.b_t, K=cfg.d_k, @@ -823,7 +728,7 @@ def mma_warp( n_dim=cfg.b_t, m_dim=cfg.d_v, ) - bmm_q_state_desc = MmaDesc( + bmm_state_k_desc = MmaDesc( M=cfg.d_v, N=cfg.b_t, K=cfg.d_k, @@ -836,6 +741,7 @@ def mma_warp( idesc=idesc_q_state, kind=nvvm.Tcgen05MMAKind.F16, ) + bmm_state_q_desc = bmm_state_k_desc idesc_qkv_ts = nvvm.Tcgen05InstrDesc.build( c_dtype=cutlass.Float32, a_dtype=cfg.io_dtype, @@ -843,7 +749,7 @@ def mma_warp( n_dim=cfg.b_t, m_dim=cfg.d_v, ) - bmm_qkv_ts_desc = MmaDesc( + bmm_u_a_desc = MmaDesc( M=cfg.d_v, N=cfg.b_t, K=cfg.b_t, @@ -856,6 +762,7 @@ def mma_warp( idesc=idesc_qkv_ts, kind=nvvm.Tcgen05MMAKind.F16, ) + bmm_y_t_inv_desc = bmm_u_a_desc idesc_kv = nvvm.Tcgen05InstrDesc.build( c_dtype=cutlass.Float32, a_dtype=cfg.io_dtype, @@ -864,7 +771,7 @@ def mma_warp( m_dim=cfg.d_k, b_major=1, ) - bmm_kv_desc = MmaDesc( + bmm_decay_u_k_desc = MmaDesc( M=cfg.d_k, N=cfg.d_v, K=cfg.b_t, @@ -880,7 +787,7 @@ def mma_warp( KQ_SEG = (2 * cfg.b_t * 64 * bpe) >> 4 KQ_BOX = (cfg.b_t * 64 * bpe) >> 4 KQ_HALF_K = (cfg.d_k // 16) // 2 - KQ_A_HALF = KQ_HALF_K * bmm_q_state_desc.tmem_advance_A + KQ_A_HALF = KQ_HALF_K * bmm_state_k_desc.tmem_advance_A ACC_STAGE_COLS = cfg.b_t KV_ACC_STAGE_COLS = cfg.d_v @@ -888,177 +795,186 @@ def mma_warp( INP_SLOT_COLS = cfg.b_t // 2 tmem_base = tmem_base_slot.load() - tmem_cg0_acc_col_f = tmem_base + cfg.tmem_cg0_acc_offset - tmem_state_col = tmem_base + cfg.tmem_state_acc_offset - tmem_q_state_col = tmem_base + cfg.tmem_q_state_acc_offset - tmem_state_inp_col = tmem_base + cfg.tmem_state_inp_offset - tmem_inp_col = tmem_base + cfg.tmem_y_decay_u_inp_offset - y_inp_ptr = nvvm.make_tmem_ptr(tmem_inp_col, cutlass.Int8) - u_inp_ptr = nvvm.make_tmem_ptr(tmem_inp_col, cutlass.Int8) - decay_u_inp_ptr = nvvm.make_tmem_ptr(tmem_inp_col + INP_SLOT_COLS, cutlass.Int8) - k_state_acc_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_cg1_acc_offset, cutlass.Float32) + tmem_col = tmem_base & 0xFFFF + tmem_row = tmem_base >> 16 + row_lo_addr = tmem_row << 16 + row_hi_addr = (tmem_row + 16) << 16 + tmem_cg0_acc_col = tmem_col + cfg.tmem_cg0_acc_offset + tmem_state_col = tmem_col + cfg.tmem_state_acc_offset + tmem_q_state_col = tmem_col + cfg.tmem_q_state_acc_offset + tmem_state_input_col = tmem_col + cfg.tmem_state_input_offset + tmem_input_col = tmem_col + cfg.tmem_y_decay_u_input_offset + y_input_ptr = nvvm.make_tmem_ptr(tmem_input_col, cutlass.Int8) + u_input_ptr = nvvm.make_tmem_ptr(tmem_input_col, cutlass.Int8) + decay_u_input_ptr = nvvm.make_tmem_ptr(tmem_input_col + INP_SLOT_COLS, cutlass.Int8) + k_state_acc_ptr = nvvm.make_tmem_ptr(tmem_col + cfg.tmem_cg1_acc_offset, cutlass.Float32) u_acc_ptr = k_state_acc_ptr - sched_state = PipelineState.start(phase=0) + scheduler_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) - n_local = wend - cstart + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) + n_local = write_end - compute_start - # ---- fused KK^T/QK^T pair 0: each member issued ahead of the loop ------- + # ---- fused KK^T/QK^T pair 0: each member issued ahead of the loop ------------ if n_local > 0: - f0_acc_idx = cg0_acc_index.idx - bars.mb_cg0_acc_done[f0_acc_idx].wait(cg0_acc_index.phase) + cg0_acc_idx = cg0_acc_index.idx + bars.mb_cg0_acc_done[cg0_acc_idx].wait(cg0_acc_index.phase) cg0_acc_index = advance(cg0_acc_index, cfg.tmem_cg0_acc_stages) - kqf_idx = kq_fused_index.idx - bars.mb_kq_ready[kqf_idx].wait(kq_fused_index.phase) - kq_fused_index = advance(kq_fused_index, cfg.smem_kq_stages) - desc_kqf = sKQ[kqf_idx].desc() - acc_cg0 = nvvm.make_tmem_ptr(tmem_cg0_acc_col_f + f0_acc_idx * cfg.b_t, cutlass.Float32) - mma_ss(bmm_qk_desc, desc_kqf, desc_kqf, acc_cg0, accumulate=False, k_count=KQ_HALF_K) - mma_ss(bmm_qk_desc, desc_kqf + KQ_SEG, desc_kqf + KQ_SEG, acc_cg0, accumulate=True, k_count=KQ_HALF_K) + kq_cg0_idx = kq_cg0_index.idx + bars.mb_kq_ready[kq_cg0_idx].wait(kq_cg0_index.phase) + kq_cg0_index = advance(kq_cg0_index, cfg.smem_kq_stages) + desc_kq_cg0 = sKQ[kq_cg0_idx].desc() + acc_cg0 = nvvm.make_tmem_ptr(tmem_cg0_acc_col + cg0_acc_idx * cfg.b_t, cutlass.Float32) + mma_ss(bmm_kq_k_desc, desc_kq_cg0, desc_kq_cg0, acc_cg0, accumulate=False, k_count=KQ_HALF_K) + mma_ss(bmm_kq_k_desc, desc_kq_cg0 + KQ_SEG, desc_kq_cg0 + KQ_SEG, acc_cg0, accumulate=True, k_count=KQ_HALF_K) if elect_one: - bars.mb_cg0_acc_ready[f0_acc_idx].arrive(cta_group=1) + bars.mb_cg0_acc_ready[cg0_acc_idx].arrive(cta_group=1) if n_local > 1: - pend_acc_idx = cg0_acc_index.idx - bars.mb_cg0_acc_done[pend_acc_idx].wait(cg0_acc_index.phase) + cg0_acc_idx = cg0_acc_index.idx + bars.mb_cg0_acc_done[cg0_acc_idx].wait(cg0_acc_index.phase) cg0_acc_index = advance(cg0_acc_index, cfg.tmem_cg0_acc_stages) - kqf_idx = kq_fused_index.idx - bars.mb_kq_ready[kqf_idx].wait(kq_fused_index.phase) - kq_fused_index = advance(kq_fused_index, cfg.smem_kq_stages) - desc_kqf = sKQ[kqf_idx].desc() - desc_kqf_b = desc_kqf + KQ_BOX - acc_cg0 = nvvm.make_tmem_ptr(tmem_cg0_acc_col_f + pend_acc_idx * cfg.b_t, cutlass.Float32) - mma_ss(bmm_qk_desc, desc_kqf, desc_kqf_b, acc_cg0, accumulate=False, k_count=KQ_HALF_K) - mma_ss(bmm_qk_desc, desc_kqf + KQ_SEG, desc_kqf_b + KQ_SEG, acc_cg0, accumulate=True, k_count=KQ_HALF_K) + kq_cg0_idx = kq_cg0_index.idx + bars.mb_kq_ready[kq_cg0_idx].wait(kq_cg0_index.phase) + kq_cg0_index = advance(kq_cg0_index, cfg.smem_kq_stages) + desc_kq_cg0 = sKQ[kq_cg0_idx].desc() + desc_k_cg0 = desc_kq_cg0 + KQ_BOX + acc_cg0 = nvvm.make_tmem_ptr(tmem_cg0_acc_col + cg0_acc_idx * cfg.b_t, cutlass.Float32) + mma_ss(bmm_kq_k_desc, desc_kq_cg0, desc_k_cg0, acc_cg0, accumulate=False, k_count=KQ_HALF_K) + mma_ss(bmm_kq_k_desc, desc_kq_cg0 + KQ_SEG, desc_k_cg0 + KQ_SEG, acc_cg0, accumulate=True, k_count=KQ_HALF_K) if elect_one: - bars.mb_cg0_acc_ready[pend_acc_idx].arrive(cta_group=1) + bars.mb_cg0_acc_ready[cg0_acc_idx].arrive(cta_group=1) - for local_idx in cutlass.range(n_local): # noqa: B007 - if cutlass.const_expr(cfg.use_initial_state): - if local_idx == 0: - if elect_one: - bars.mb_state_acc_ready[kv_acc_index.idx].arrive(cta_group=1) - kv_acc_index = advance(kv_acc_index, cfg.tmem_state_acc_stages) - have_state = cutlass.Boolean(True) if cutlass.const_expr(cfg.use_initial_state) else local_idx > 0 + for peel in cutlass.range_constexpr(1 if cfg.use_initial_state else 2): + have_state = cutlass.Boolean(True) if cutlass.const_expr(cfg.use_initial_state) else cutlass.const_expr(peel == 1) + peel_stop = n_local if cutlass.const_expr(peel == 1 or cfg.use_initial_state) else min(n_local, 1) + for local_idx in cutlass.range(peel, peel_stop, 1): # noqa: B007 + kq_idx = kq_index.idx + member = local_idx & 1 + state_input_idx = state_input_index.idx + q_state_acc_idx = o_acc_index.idx + tinv_idx = tinv_index.idx + a_idx = a_index.idx + o_scale_idx = o_state_scale_index.idx + kv_acc_idx = kv_acc_index.idx + kq_member_off = member * KQ_BOX + desc_k = sKQ[kq_idx].desc() + kq_member_off + desc_q = sKQ[kq_idx].desc() + (KQ_BOX - kq_member_off) + desc_tinv = sTinv[tinv_idx].desc() + desc_a = sA[a_idx].desc() + desc_k_trans = sKQ_trans[kq_idx].desc() + kq_member_off + state_a_ptr = nvvm.make_tmem_ptr(tmem_state_input_col + state_input_idx * STATE_INP_STAGE_COLS, cutlass.Int8) + q_state_acc_ptr = nvvm.make_tmem_ptr(tmem_q_state_col + q_state_acc_idx * ACC_STAGE_COLS, cutlass.Float32) + qkv_acc_ptr = nvvm.make_tmem_ptr(tmem_q_state_col + o_scale_idx * ACC_STAGE_COLS, cutlass.Float32) + state_acc_ptr = nvvm.make_tmem_ptr(tmem_state_col + kv_acc_idx * KV_ACC_STAGE_COLS, cutlass.Float32) - kq_idx = kq_index.idx - member = local_idx & 1 - state_inp_idx = state_inp_index.idx - q_state_acc_idx = o_acc_index.idx - tinv_idx = tinv_index.idx - a_idx = a_index.idx - o_scale_idx = o_state_scale_index.idx - kv_acc_idx = kv_acc_index.idx - kq_member_off = member * KQ_BOX - desc_k = sKQ[kq_idx].desc() + kq_member_off - desc_q = sKQ[kq_idx].desc() + (KQ_BOX - kq_member_off) - desc_tinv = sTinv[tinv_idx].desc() - desc_a = sA[a_idx].desc() - desc_kt = sKQ_trans[kq_idx].desc() + kq_member_off - state_a_ptr = nvvm.make_tmem_ptr(tmem_state_inp_col + state_inp_idx * STATE_INP_STAGE_COLS, cutlass.Int8) - q_state_acc_ptr = nvvm.make_tmem_ptr(tmem_q_state_col + q_state_acc_idx * ACC_STAGE_COLS, cutlass.Float32) - qkv_acc_ptr = nvvm.make_tmem_ptr(tmem_q_state_col + o_scale_idx * ACC_STAGE_COLS, cutlass.Float32) - state_acc_ptr = nvvm.make_tmem_ptr(tmem_state_col + kv_acc_idx * KV_ACC_STAGE_COLS, cutlass.Float32) + kq_index = advance(kq_index, cfg.smem_kq_stages) - kq_index = advance(kq_index, cfg.smem_kq_stages) + # ---- QK/KK lookahead (member 1) = [Q;K](S) @ K^T --------------------- + if member == 1: + if local_idx + 2 < n_local: + cg0_acc_idx = cg0_acc_index.idx + bars.mb_cg0_acc_done[cg0_acc_idx].wait(cg0_acc_index.phase) + cg0_acc_index = advance(cg0_acc_index, cfg.tmem_cg0_acc_stages) + kq_cg0_idx = kq_cg0_index.idx + bars.mb_kq_ready[kq_cg0_idx].wait(kq_cg0_index.phase) + kq_cg0_index = advance(kq_cg0_index, cfg.smem_kq_stages) + desc_kq_cg0 = sKQ[kq_cg0_idx].desc() + desc_k_cg0 = desc_kq_cg0 + KQ_BOX + acc_cg0 = nvvm.make_tmem_ptr(tmem_cg0_acc_col + cg0_acc_idx * cfg.b_t, cutlass.Float32) + mma_ss(bmm_kq_k_desc, desc_kq_cg0, desc_k_cg0, acc_cg0, accumulate=False, k_count=KQ_HALF_K) + mma_ss(bmm_kq_k_desc, desc_kq_cg0 + KQ_SEG, desc_k_cg0 + KQ_SEG, acc_cg0, accumulate=True, k_count=KQ_HALF_K) + if elect_one: + bars.mb_cg0_acc_ready[cg0_acc_idx].arrive(cta_group=1) - # ---- QK/KK lookahead (member 1) = [Q;K](S) @ K^T --------------------- - if member == 1: - if local_idx + 2 < n_local: - pend_acc_idx = cg0_acc_index.idx - bars.mb_cg0_acc_done[pend_acc_idx].wait(cg0_acc_index.phase) - cg0_acc_index = advance(cg0_acc_index, cfg.tmem_cg0_acc_stages) - kqf_idx = kq_fused_index.idx - bars.mb_kq_ready[kqf_idx].wait(kq_fused_index.phase) - kq_fused_index = advance(kq_fused_index, cfg.smem_kq_stages) - desc_kqf = sKQ[kqf_idx].desc() - desc_kqf_b = desc_kqf + KQ_BOX - acc_cg0 = nvvm.make_tmem_ptr(tmem_cg0_acc_col_f + pend_acc_idx * cfg.b_t, cutlass.Float32) - mma_ss(bmm_qk_desc, desc_kqf, desc_kqf_b, acc_cg0, accumulate=False, k_count=KQ_HALF_K) - mma_ss(bmm_qk_desc, desc_kqf + KQ_SEG, desc_kqf_b + KQ_SEG, acc_cg0, accumulate=True, k_count=KQ_HALF_K) - if elect_one: - bars.mb_cg0_acc_ready[pend_acc_idx].arrive(cta_group=1) + # ---- k state^T (GEMM 3) = state^T(T) @ K^T --------------------------- + if have_state: + bars.mb_state_input_ready[state_input_idx].wait(state_input_index.phase) + state_input_index = advance(state_input_index, cfg.tmem_state_input_stages) - # ---- K*state^T (GEMM 3) = state^T(T) @ K^T ------------------------------------ - if have_state: - bars.mb_state_inp_ready[state_inp_idx].wait(state_inp_index.phase) - state_inp_index = advance(state_inp_index, cfg.tmem_state_inp_stages) + for k in cutlass.range_constexpr(KQ_HALF_K): + mma_ts_step(bmm_state_k_desc, state_a_ptr, desc_k, k_state_acc_ptr, k, cutlass.Boolean(k > 0)) + for k in cutlass.range_constexpr(KQ_HALF_K): + mma_ts_step(bmm_state_k_desc, state_a_ptr.subview(KQ_A_HALF), desc_k + KQ_SEG, k_state_acc_ptr, k, cutlass.Boolean(True)) + if elect_one: + bars.mb_k_state_acc_ready[0].arrive(cta_group=1) + + # ---- q state^T (GEMM 4) = state^T(T) @ Q^T --------------------------- + o_acc_index = advance(o_acc_index, cfg.tmem_q_state_acc_stages) + if have_state: + for k in cutlass.range_constexpr(KQ_HALF_K): + mma_ts_step(bmm_state_q_desc, state_a_ptr, desc_q, q_state_acc_ptr, k, cutlass.Boolean(k > 0)) + for k in cutlass.range_constexpr(KQ_HALF_K): + mma_ts_step(bmm_state_q_desc, state_a_ptr.subview(KQ_A_HALF), desc_q + KQ_SEG, q_state_acc_ptr, k, cutlass.Boolean(True)) + if elect_one: + bars.mb_o_acc_ready[q_state_acc_idx].arrive(cta_group=1) - for k in cutlass.range_constexpr(KQ_HALF_K): - mma_ts_step(bmm_q_state_desc, state_a_ptr, desc_k, k_state_acc_ptr, k, cutlass.Boolean(k > 0)) - for k in cutlass.range_constexpr(KQ_HALF_K): - mma_ts_step(bmm_q_state_desc, state_a_ptr.subview(KQ_A_HALF), desc_k + KQ_SEG, k_state_acc_ptr, k, cutlass.Boolean(True)) - if elect_one: - bars.mb_k_state_acc_ready[0].arrive(cta_group=1) - - # ---- Q*state^T (GEMM 4) = state^T(T) @ Q^T ------------------------------------ - o_acc_index = advance(o_acc_index, cfg.tmem_q_state_acc_stages) - if have_state: - for k in cutlass.range_constexpr(KQ_HALF_K): - mma_ts_step(bmm_q_state_desc, state_a_ptr, desc_q, q_state_acc_ptr, k, cutlass.Boolean(k > 0)) - for k in cutlass.range_constexpr(KQ_HALF_K): - mma_ts_step(bmm_q_state_desc, state_a_ptr.subview(KQ_A_HALF), desc_q + KQ_SEG, q_state_acc_ptr, k, cutlass.Boolean(True)) + # ---- U^T (GEMM 5) = Y^T(T) @ (T^-1)^T -------------------------------- + bars.mb_t_inv_ready[tinv_idx].wait(tinv_index.phase) + tinv_index = advance(tinv_index, cfg.smem_t_inv_stages) + bars.mb_y_input_ready[0].wait(y_input_ready.phase) + y_input_ready = advance(y_input_ready, 1) + for k in cutlass.range_constexpr(cfg.b_t // 16): + mma_ts_step(bmm_y_t_inv_desc, y_input_ptr, desc_tinv, u_acc_ptr, k, cutlass.Boolean(k > 0)) if elect_one: - bars.mb_o_acc_ready[q_state_acc_idx].arrive(cta_group=1) - - # ---- U^T (GEMM 5) = Y^T(T) @ T_inv^T --------------------------------------- - bars.mb_t_inv_ready[tinv_idx].wait(tinv_index.phase) - tinv_index = advance(tinv_index, cfg.smem_t_inv_stages) - bars.mb_y_inp_ready[0].wait(y_inp_ready.phase) - y_inp_ready = advance(y_inp_ready, 1) - for k in cutlass.range_constexpr(cfg.b_t // 16): - mma_ts_step(bmm_qkv_ts_desc, y_inp_ptr, desc_tinv, u_acc_ptr, k, cutlass.Boolean(k > 0)) - if elect_one: - bars.mb_u_acc_ready[0].arrive(cta_group=1) - bars.mb_t_inv_done[tinv_idx].arrive(cta_group=1) - - # ---- KK/QK lookahead (member 0) = [K;Q](S) @ K^T --------------------- - if member == 0: - if local_idx + 2 < n_local: - f0_acc_idx = cg0_acc_index.idx - bars.mb_cg0_acc_done[f0_acc_idx].wait(cg0_acc_index.phase) - cg0_acc_index = advance(cg0_acc_index, cfg.tmem_cg0_acc_stages) - kqf_idx = kq_fused_index.idx - bars.mb_kq_ready[kqf_idx].wait(kq_fused_index.phase) - kq_fused_index = advance(kq_fused_index, cfg.smem_kq_stages) - desc_kqf = sKQ[kqf_idx].desc() - acc_cg0 = nvvm.make_tmem_ptr(tmem_cg0_acc_col_f + f0_acc_idx * cfg.b_t, cutlass.Float32) - mma_ss(bmm_qk_desc, desc_kqf, desc_kqf, acc_cg0, accumulate=False, k_count=KQ_HALF_K) - mma_ss(bmm_qk_desc, desc_kqf + KQ_SEG, desc_kqf + KQ_SEG, acc_cg0, accumulate=True, k_count=KQ_HALF_K) - if elect_one: - bars.mb_cg0_acc_ready[f0_acc_idx].arrive(cta_group=1) + bars.mb_u_acc_ready[0].arrive(cta_group=1) + bars.mb_t_inv_done[tinv_idx].arrive(cta_group=1) - # ---- O^T (GEMM 6) += U input^T(T) @ A^T ------------------------------ - bars.mb_a_ready[a_idx].wait(a_index.phase) - a_index = advance(a_index, cfg.smem_a_stages) - if have_state: - bars.mb_o_state_scale_acc_done[o_scale_idx].wait(o_state_scale_index.phase) - o_state_scale_index = advance(o_state_scale_index, cfg.tmem_q_state_acc_stages) - bars.mb_u_inp_ready[0].wait(u_inp_ready.phase) - u_inp_ready = advance(u_inp_ready, 1) - for k in cutlass.range_constexpr(cfg.b_t // 16): - mma_ts_step(bmm_qkv_ts_desc, u_inp_ptr, desc_a, qkv_acc_ptr, k, cutlass.Boolean(True) if cutlass.const_expr(k > 0) else have_state) - if elect_one: - bars.mb_a_done[a_idx].arrive(cta_group=1) - bars.mb_o_final_acc_ready[o_scale_idx].arrive(cta_group=1) + # ---- KK/QK lookahead (member 0) = [K;Q](S) @ K^T --------------------- + if member == 0: + if local_idx + 2 < n_local: + cg0_acc_idx = cg0_acc_index.idx + bars.mb_cg0_acc_done[cg0_acc_idx].wait(cg0_acc_index.phase) + cg0_acc_index = advance(cg0_acc_index, cfg.tmem_cg0_acc_stages) + kq_cg0_idx = kq_cg0_index.idx + bars.mb_kq_ready[kq_cg0_idx].wait(kq_cg0_index.phase) + kq_cg0_index = advance(kq_cg0_index, cfg.smem_kq_stages) + desc_kq_cg0 = sKQ[kq_cg0_idx].desc() + acc_cg0 = nvvm.make_tmem_ptr(tmem_cg0_acc_col + cg0_acc_idx * cfg.b_t, cutlass.Float32) + mma_ss(bmm_kq_k_desc, desc_kq_cg0, desc_kq_cg0, acc_cg0, accumulate=False, k_count=KQ_HALF_K) + mma_ss(bmm_kq_k_desc, desc_kq_cg0 + KQ_SEG, desc_kq_cg0 + KQ_SEG, acc_cg0, accumulate=True, k_count=KQ_HALF_K) + if elect_one: + bars.mb_cg0_acc_ready[cg0_acc_idx].arrive(cta_group=1) - # ---- state^T (GEMM 7) += decayed U^T(T) @ K ------------------------------ - bars.mb_decay_u_inp_ready[0].wait(decay_u_inp_ready.phase) - decay_u_inp_ready = advance(decay_u_inp_ready, 1) - kv_acc_index = advance(kv_acc_index, cfg.tmem_state_acc_stages) - for k in cutlass.range_constexpr(cfg.b_t // 16): - mma_ts_step(bmm_kv_desc, decay_u_inp_ptr, desc_kt, state_acc_ptr, k, cutlass.Boolean(True) if cutlass.const_expr(k > 0) else have_state) - if elect_one: - bars.mb_state_acc_ready[kv_acc_idx].arrive(cta_group=1) - bars.mb_kq_done[kq_idx].arrive(cta_group=1) + # ---- O^T (GEMM 6) += U input^T(T) @ A^T ------------------------------ + bars.mb_a_ready[a_idx].wait(a_index.phase) + a_index = advance(a_index, cfg.smem_a_stages) + if have_state: + bars.mb_o_state_scale_acc_done[o_scale_idx].wait(o_state_scale_index.phase) + o_state_scale_index = advance(o_state_scale_index, cfg.tmem_q_state_acc_stages) + bars.mb_u_input_ready[0].wait(u_input_ready.phase) + u_input_ready = advance(u_input_ready, 1) + for k in cutlass.range_constexpr(cfg.b_t // 16): + mma_ts_step(bmm_u_a_desc, u_input_ptr, desc_a, qkv_acc_ptr, k, cutlass.Boolean(True) if cutlass.const_expr(k > 0) else have_state) + if elect_one: + bars.mb_a_done[a_idx].arrive(cta_group=1) + bars.mb_o_final_acc_ready[o_scale_idx].arrive(cta_group=1) + + # ---- state^T (GEMM 7) += decayed U^T(T) @ K -------------------------- + bars.mb_decay_u_input_ready[0].wait(decay_u_input_ready.phase) + decay_u_input_ready = advance(decay_u_input_ready, 1) + kv_acc_index = advance(kv_acc_index, cfg.tmem_state_acc_stages) + for k in cutlass.range_constexpr(cfg.b_t // 16): + mma_ts_step( + bmm_decay_u_k_desc, + decay_u_input_ptr, + desc_k_trans, + state_acc_ptr, + k, + cutlass.Boolean(True) if cutlass.const_expr(k > 0) else have_state, + ) + if elect_one: + bars.mb_state_acc_ready[kv_acc_idx].arrive(cta_group=1) + bars.mb_kq_done[kq_idx].arrive(cta_group=1) - tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one) bars.mb_tmem_done[0].wait(0) nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_1) nvvm.tcgen05_dealloc( - nvvm.make_tmem_ptr(tmem_base, cutlass.Int8), + nvvm.make_tmem_ptr(tmem_col, cutlass.Int8), cutlass.Int32(512), group=nvvm.CTAGroup.CTA_1, ) @@ -1077,34 +993,35 @@ def tmaldg_warp( desc_q_base, desc_k_base, desc_v_base, - mSched, - sSched, + mScheduler, + sScheduler, bars, ): """TMA-LDG warp role (warp 9): persistent scheduler loop + per-chunk Q/K/V G->S TMA loads.""" - elect_one = nvvm.elect_sync() nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) + kq_index = PipelineState.start(phase=1) v_index = PipelineState.start(phase=1) - sched_state = PipelineState.start(phase=1) - tile_idx = cutlass.Int32(bidx) + scheduler_state = PipelineState.start(phase=1) + elect_one = nvvm.elect_sync() + tile_idx = cutlass.Int32(bidx) bpe = cfg.io_dtype.width // 8 - elems_per_128b = 128 // bpe + elements_per_128b = 128 // bpe bt = cfg.b_t - kq_stage_elems = cfg.kq_cosize // cfg.smem_kq_stages - kq_box_elems = kq_stage_elems // 4 + kq_stage_elements = cfg.kq_cosize // cfg.smem_kq_stages + kq_box_elements = kq_stage_elements // 4 sKQ_lo_tma = SmemTile( base=sKQ_raw, - elems_per_stage=kq_stage_elems, + elems_per_stage=kq_stage_elements, stages=cfg.smem_kq_stages, leading_byte_offset=0, stride_byte_offset=0, layout=0, tma_loads_per_tile=2, - tma_granu_elems=elems_per_128b, - tma_subtile_stride_elems=2 * bt * elems_per_128b, + tma_granu_elems=elements_per_128b, + tma_subtile_stride_elems=2 * bt * elements_per_128b, ) sV_tma = SmemTile( base=sV_raw, @@ -1114,14 +1031,16 @@ def tmaldg_warp( stride_byte_offset=0, layout=0, tma_loads_per_tile=2, - tma_granu_elems=elems_per_128b, + tma_granu_elems=elements_per_128b, tma_subtile_stride_elems=4096, ) heads_out = cutlass.Int32(cfg.n_heads_out) desc_qwords = cutlass.Int32(TENSOR_MAP_QWORDS) while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) head_q = head_idx if cfg.q_ratio == 1 else head_idx // cutlass.Int32(cfg.q_ratio) head_k = head_idx if cfg.k_ratio == 1 else head_idx // cutlass.Int32(cfg.k_ratio) @@ -1135,39 +1054,39 @@ def tmaldg_warp( tma_tensormap_acquire(desc_k_slot) tma_tensormap_acquire(desc_v_slot) - if wend > cstart: + if write_end > compute_start: kq_idx = kq_index.idx bars.mb_kq_done[kq_idx].wait(kq_index.phase) kq_index = advance(kq_index, cfg.smem_kq_stages) if elect_one: bars.mb_kq_ready[kq_idx].arrive(n_bytes=cfg.tma_kq_bytes) - tok_coord = cstart * cutlass.Int32(cfg.b_t) + tok_coord = compute_start * cutlass.Int32(cfg.b_t) k_slice = tma_slice_runtime_desc(desc_k_slot, cutlass.Int32(0), head_k, tok_coord) q_slice = tma_slice_runtime_desc(desc_q_slot, cutlass.Int32(0), head_q, tok_coord) kq_tile = sKQ_lo_tma[kq_idx] tma_load_tile(kq_tile, k_slice, bars.mb_kq_ready[kq_idx].smem_ptr, acquire=False) - tma_load_tile(kq_tile.shifted(kq_box_elems), q_slice, bars.mb_kq_ready[kq_idx].smem_ptr, acquire=False) - for chunk_idx in cutlass.range(cstart + 1, wend): + tma_load_tile(kq_tile.shifted(kq_box_elements), q_slice, bars.mb_kq_ready[kq_idx].smem_ptr, acquire=False) + for chunk_idx in cutlass.range(compute_start + 1, write_end): tok_coord = chunk_idx * cutlass.Int32(cfg.b_t) - # ---- K + Q interleaved ------------------------------------------- + # ---- K + Q interleaved ----------------------------------------------- kq_idx = kq_index.idx bars.mb_kq_done[kq_idx].wait(kq_index.phase) kq_index = advance(kq_index, cfg.smem_kq_stages) if elect_one: bars.mb_kq_ready[kq_idx].arrive(n_bytes=cfg.tma_kq_bytes) - member = (chunk_idx - cstart) & 1 + member = (chunk_idx - compute_start) & 1 k_slice = tma_slice_runtime_desc(desc_k_slot, cutlass.Int32(0), head_k, tok_coord) q_slice = tma_slice_runtime_desc(desc_q_slot, cutlass.Int32(0), head_q, tok_coord) kq_tile = sKQ_lo_tma[kq_idx] if member == 0: tma_load_tile(kq_tile, k_slice, bars.mb_kq_ready[kq_idx].smem_ptr, acquire=False) - tma_load_tile(kq_tile.shifted(kq_box_elems), q_slice, bars.mb_kq_ready[kq_idx].smem_ptr, acquire=False) + tma_load_tile(kq_tile.shifted(kq_box_elements), q_slice, bars.mb_kq_ready[kq_idx].smem_ptr, acquire=False) else: tma_load_tile(kq_tile, q_slice, bars.mb_kq_ready[kq_idx].smem_ptr, acquire=False) - tma_load_tile(kq_tile.shifted(kq_box_elems), k_slice, bars.mb_kq_ready[kq_idx].smem_ptr, acquire=False) + tma_load_tile(kq_tile.shifted(kq_box_elements), k_slice, bars.mb_kq_ready[kq_idx].smem_ptr, acquire=False) - # ---- V load ------------------------------------------------------ + # ---- V load ---------------------------------------------------------- v_idx = v_index.idx bars.mb_v_done[v_idx].wait(v_index.phase) v_index = advance(v_index, cfg.smem_v_stages) @@ -1182,11 +1101,11 @@ def tmaldg_warp( v_index = advance(v_index, cfg.smem_v_stages) if elect_one: bars.mb_v_ready[v_idx].arrive(n_bytes=cfg.tma_v_bytes) - v_tok = (wend - cutlass.Int32(1)) * cutlass.Int32(cfg.b_t) + v_tok = (write_end - cutlass.Int32(1)) * cutlass.Int32(cfg.b_t) v_slice = tma_slice_runtime_desc(desc_v_slot, cutlass.Int32(0), head_v, v_tok) tma_load_tile(sV_tma[v_idx], v_slice, bars.mb_v_ready[v_idx].smem_ptr, acquire=False) - tile_idx, sched_state = sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_publish_next(cfg, bars, sScheduler, mScheduler, scheduler_state, tile_idx, num_ctas, elect_one) for _ in range(cfg.smem_kq_stages): bars.mb_kq_done[kq_index.idx].wait(kq_index.phase) @@ -1213,61 +1132,62 @@ def compute0_warp_group( sA, sCheckpoint_raw, checkpoint_every_n_tokens, - sSched, + sScheduler, bars, ): """Compute warp-group 0 role (warps 0-3): persistent scheduler loop computing each chunk pair's T_inv and A epilogues.""" nvvm.setmaxregister(cfg.num_regs_compute_group_0, nvvm.SetMaxRegisterAction.INCREASE) + elect_one = nvvm.elect_sync() + gate_index = PipelineState.start(phase=0) beta_index = PipelineState.start(phase=0) cg0_acc_ready = PipelineState.start(phase=0) tinv_index = PipelineState.start(phase=1) a_index = PipelineState.start(phase=1) - - nvvm.barrier_cta_sync_aligned( - cfg.tmem_alloc_barrier_id, - thread_count=cfg.tmem_alloc_barrier_threads, - ) - tmem_base = tmem_base_slot.load() + scheduler_state = PipelineState.start(phase=0) num_threads_cg0 = cfg.threads_per_warp * len(cfg.compute_group_0_warp_ids) cg0_tidx = tidx % num_threads_cg0 warp_id = cg0_tidx // cfg.threads_per_warp - lane_id = cg0_tidx % cfg.threads_per_warp + lane_idx = cg0_tidx % cfg.threads_per_warp inverse_local_warp = warp_id % 2 + store_row = warp_id * 16 + lane_idx % 16 + store_col = (lane_idx // 16) * 8 + pair_half = warp_id // 2 half_row_base = inverse_local_warp * 32 - bpe = cfg.io_dtype.width // 8 num_vals = 32 FRAG_COLS = 16 ACC_N_FRAGS = cfg.b_t // FRAG_COLS - store_row = warp_id * 16 + lane_id % 16 - store_row_frag = lane_id % 16 - store_col = (lane_id // 16) * 8 - tmem_warp_row = warp_id * cfg.threads_per_warp - tmem_cg0_acc_col = tmem_base + cfg.tmem_cg0_acc_offset + store_row_frag = lane_idx % 16 ACC_STAGE_COLS = cfg.b_t mask_zero = opaque_f32_zero() - crow_lo = warp_id * 16 + lane_id // 4 - crow_hi = crow_lo + 8 - - sched_state = PipelineState.start(phase=0) + chunk_row_lo = warp_id * 16 + lane_idx // 4 + chunk_row_hi = chunk_row_lo + 8 tile_idx = cutlass.Int32(bidx) + + nvvm.barrier_cta_sync_aligned(cfg.tmem_lifecycle_barrier_id, thread_count=cfg.tmem_user_threads) + tmem_base = tmem_base_slot.load() + tmem_col = tmem_base & 0xFFFF + tmem_row = tmem_base >> 16 + row_lo_addr = tmem_row << 16 + row_hi_addr = (tmem_row + 16) << 16 + tmem_cg0_acc_col = tmem_col + cfg.tmem_cg0_acc_offset while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) - n_local = wend - cstart + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) + n_local = write_end - compute_start n_pairs = (n_local + 1) // 2 for pair_i in cutlass.range(n_pairs): - # An odd chunk count leaves the last pair with member 0 only; have_m1 - # is uniform across CG0, so the shared barriers below stay aligned. have_m1 = pair_i * 2 + 1 < n_local do_kk = have_m1 or pair_half == 0 do_a = have_m1 or pair_half == 1 - # ---- Gate rows for this warp's KK / A member roles ------------------- + # ---- Gate rows for this warp's KK / A member roles ----------------------- gate0_idx = gate_index.idx bars.mb_gate_ready[gate0_idx].wait(gate_index.phase) gate_index = advance(gate_index, cfg.smem_gate_stages) @@ -1279,7 +1199,7 @@ def compute0_warp_group( kk_gate_idx = gate1_idx if pair_half == 1 else gate0_idx a_gate_idx = gate0_idx if pair_half == 1 else gate1_idx - row_u0_lo = half_row_base + lane_id // 4 + row_u0_lo = half_row_base + lane_idx // 4 row_u0_hi = row_u0_lo + 8 row_u1_lo = row_u0_lo + 16 row_u1_hi = row_u0_lo + 24 @@ -1293,20 +1213,20 @@ def compute0_warp_group( a_cumsumlog_cols = [] for g in cutlass.range_constexpr(8): for b in cutlass.range_constexpr(2): - ccol = (lane_id % 4) * 2 + g * 8 + b - kk_cumsumlog_cols.append(sCumsumlog[ccol, 0, kk_gate_idx]) - a_cumsumlog_cols.append(sCumsumlog[ccol, 0, a_gate_idx]) + chunk_col = (lane_idx % 4) * 2 + g * 8 + b + kk_cumsumlog_cols.append(sCumsumlog[chunk_col, 0, kk_gate_idx]) + a_cumsumlog_cols.append(sCumsumlog[chunk_col, 0, a_gate_idx]) decay_t_kk = [] decay_t_a = [] for u in cutlass.range_constexpr(2): for k in cutlass.range_constexpr(num_vals): hi_row = ((k // 2) % 2) == 1 - crow_u0 = row_u0_hi if cutlass.const_expr(hi_row) else row_u0_lo - crow_u1 = row_u1_hi if cutlass.const_expr(hi_row) else row_u1_lo - crow = crow_u1 if cutlass.const_expr(u == 1) else crow_u0 - ccol = (lane_id % 4) * 2 + ((k // 4) * 8 + k % 2) - is_lower = crow >= ccol + chunk_row_u0 = row_u0_hi if cutlass.const_expr(hi_row) else row_u0_lo + chunk_row_u1 = row_u1_hi if cutlass.const_expr(hi_row) else row_u1_lo + chunk_row = chunk_row_u1 if cutlass.const_expr(u == 1) else chunk_row_u0 + chunk_col = (lane_idx % 4) * 2 + ((k // 4) * 8 + k % 2) + is_lower = chunk_row >= chunk_col kk_row_cumsumlog = kk_cumsumlog_rows[u * 2 + (1 if hi_row else 0)] a_row_cumsumlog = a_cumsumlog_rows[u * 2 + (1 if hi_row else 0)] col = (k // 4) * 2 + (k % 2) @@ -1329,7 +1249,7 @@ def compute0_warp_group( for r in (row_u0_lo, row_u0_hi, row_u1_lo, row_u1_hi): kk_beta.append(sBeta[r, 0, kk_beta_idx]) - # ---- KK_epi (each warp pair stages its own member) ------------------- + # ---- KK epilogue (each warp pair stages its own member) ------------------ acc0_idx = cg0_acc_ready.idx acc0_phase = cg0_acc_ready.phase cg0_acc_ready = advance(cg0_acc_ready, cfg.tmem_cg0_acc_stages) @@ -1360,15 +1280,11 @@ def compute0_warp_group( kk_base = tinv1_base if pair_half == 1 else tinv0_base if do_kk: bars.mb_cg0_acc_ready[kk_acc_idx].wait(kk_acc_phase) - kk_vec0 = nvvm.tcgen05_ld( - "16x256b", nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_cg0_acc_col + kk_acc_idx * ACC_STAGE_COLS, cutlass.Float32), num=8 - ) - kk_vec1 = nvvm.tcgen05_ld( - "16x256b", nvvm.make_tmem_ptr(((tmem_warp_row + 16) << 16) + tmem_cg0_acc_col + kk_acc_idx * ACC_STAGE_COLS, cutlass.Float32), num=8 - ) + kk_vec_lo = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_lo_addr + tmem_cg0_acc_col + kk_acc_idx * ACC_STAGE_COLS, cutlass.Float32), num=8) + kk_vec_hi = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_hi_addr + tmem_cg0_acc_col + kk_acc_idx * ACC_STAGE_COLS, cutlass.Float32), num=8) bars.mb_t_inv_done[kk_tinv_idx].wait(kk_tinv_phase) for u in cutlass.range_constexpr(2): - kk_vec = kk_vec1 if cutlass.const_expr(u == 1) else kk_vec0 + kk_vec = kk_vec_hi if cutlass.const_expr(u == 1) else kk_vec_lo kk_pack = [] for k in cutlass.range_constexpr(num_vals // 2): b0 = kk_beta[u * 2 + 1] if cutlass.const_expr((k % 2) == 1) else kk_beta[u * 2] @@ -1378,18 +1294,12 @@ def compute0_warp_group( st_row = half_row_base + u * 16 + store_row_frag for c in cutlass.range_constexpr(ACC_N_FRAGS): nvvm.stmatrix( - cutlass.inttoptr( - kk_base + (st_row * cfg.b_t + swizzle_xor_128b(st_row, store_col + c * FRAG_COLS)) * bpe, - cutlass.AddressSpace.smem, - cutlass.BFloat16, - ), + kk_base + st_row * cfg.b_t + swizzle_xor_128b(st_row, store_col + c * FRAG_COLS), [kk_pack[c * 4 + 0], kk_pack[c * 4 + 1], kk_pack[c * 4 + 2], kk_pack[c * 4 + 3]], nvvm.MMALayout.ROW, ) - # ---- pair inverse: warps 0-1 own matrix 0, warps 2-3 matrix 1 -------- - # With no member 1 there is one matrix, so warps 2-3 idle through the - # per-warp steps; the barriers below stay unconditional. + # ---- pair inverse: warps 0-1 own matrix 0, warps 2-3 matrix 1 ------------ inv_base = tinv0_base if have_m1: inv_base = tinv1_base if warp_id >= 2 else tinv0_base @@ -1401,16 +1311,16 @@ def compute0_warp_group( thread_count=cfg.inverse_barrier_threads, ) if do_inv: - invert_diagonal_NxN(cfg, inv_base, (inverse_local_warp * cfg.threads_per_warp + lane_id) // 8, cg0_tidx, 8) + invert_diagonal_NxN(cfg, inv_base, (inverse_local_warp * cfg.threads_per_warp + lane_idx) // 8, cg0_tidx, 8) nvvm.barrier_cta_sync_aligned( cfg.inverse_barrier_id, thread_count=cfg.inverse_barrier_threads, ) # 8x8 -> 16x16 (both matrices per warp) - blockwise_diagonal_8x8_to_16x16(cfg, tinv0_base, warp_id * 16, lane_id) + blockwise_diagonal_8x8_to_16x16(cfg, tinv0_base, warp_id * 16, lane_idx) if have_m1: - blockwise_diagonal_8x8_to_16x16(cfg, tinv1_base, warp_id * 16, lane_id) + blockwise_diagonal_8x8_to_16x16(cfg, tinv1_base, warp_id * 16, lane_idx) nvvm.barrier_cta_sync_aligned( cfg.inverse_barrier_id, thread_count=cfg.inverse_barrier_threads, @@ -1418,34 +1328,28 @@ def compute0_warp_group( # 16x16 -> 32x32, one tile per warp within the group if do_inv: - blockwise_diagonal_16x16_to_32x32(cfg, inv_base, inverse_local_warp * 32, lane_id) + blockwise_diagonal_16x16_to_32x32(cfg, inv_base, inverse_local_warp * 32, lane_idx) nvvm.barrier_cta_sync_aligned( cfg.inverse_barrier_id, thread_count=cfg.inverse_barrier_threads, ) # 32x32 -> 64x64, two warps per matrix - # all four warps: this step carries its own CG0 barrier, and on the tail - # inv_base is matrix 0 for every warp so the duplicate band is identical - blockwise_diagonal_32x32_to_64x64(cfg, inv_base, inverse_local_warp, lane_id) + blockwise_diagonal_32x32_to_64x64(cfg, inv_base, inverse_local_warp, lane_idx) nvvm.barrier_cta_sync_aligned( cfg.inverse_barrier_id, thread_count=cfg.inverse_barrier_threads, ) - # ---- Beta column-scaling + publish, stage 0 -------------------------- + # ---- Beta column-scaling + publish, stage 0 ------------------------------ beta_col = [] for k in cutlass.range_constexpr(num_vals): - beta_col.append(sBeta[(lane_id % 4) * 2 + ((k // 4) * 8 + k % 2), 0, beta0_idx]) + beta_col.append(sBeta[(lane_idx % 4) * 2 + ((k // 4) * 8 + k % 2), 0, beta0_idx]) tinv_frags = [] for c in cutlass.range_constexpr(ACC_N_FRAGS): tinv_frags += list( nvvm.ldmatrix( - cutlass.inttoptr( - tinv0_base + (store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS)) * bpe, - cutlass.AddressSpace.smem, - cutlass.BFloat16, - ), + tinv0_base + store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS), 4, nvvm.MMALayout.ROW, ) @@ -1457,11 +1361,7 @@ def compute0_warp_group( tinv_pack.append(fp32_to_fp16(s0, s1, dtype=cfg.io_dtype)) for c in cutlass.range_constexpr(ACC_N_FRAGS): nvvm.stmatrix( - cutlass.inttoptr( - tinv0_base + (store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS)) * bpe, - cutlass.AddressSpace.smem, - cutlass.BFloat16, - ), + tinv0_base + store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS), [tinv_pack[c * 4 + 0], tinv_pack[c * 4 + 1], tinv_pack[c * 4 + 2], tinv_pack[c * 4 + 3]], nvvm.MMALayout.ROW, ) @@ -1473,16 +1373,12 @@ def compute0_warp_group( # ---- Beta column-scaling + publish, stage 1 -------------------------- beta_col = [] for k in cutlass.range_constexpr(num_vals): - beta_col.append(sBeta[(lane_id % 4) * 2 + ((k // 4) * 8 + k % 2), 0, beta1_idx]) + beta_col.append(sBeta[(lane_idx % 4) * 2 + ((k // 4) * 8 + k % 2), 0, beta1_idx]) tinv_frags = [] for c in cutlass.range_constexpr(ACC_N_FRAGS): tinv_frags += list( nvvm.ldmatrix( - cutlass.inttoptr( - tinv1_base + (store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS)) * bpe, - cutlass.AddressSpace.smem, - cutlass.BFloat16, - ), + tinv1_base + store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS), 4, nvvm.MMALayout.ROW, ) @@ -1494,11 +1390,7 @@ def compute0_warp_group( tinv_pack.append(fp32_to_fp16(s0, s1, dtype=cfg.io_dtype)) for c in cutlass.range_constexpr(ACC_N_FRAGS): nvvm.stmatrix( - cutlass.inttoptr( - tinv1_base + (store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS)) * bpe, - cutlass.AddressSpace.smem, - cutlass.BFloat16, - ), + tinv1_base + store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS), [tinv_pack[c * 4 + 0], tinv_pack[c * 4 + 1], tinv_pack[c * 4 + 2], tinv_pack[c * 4 + 3]], nvvm.MMALayout.ROW, ) @@ -1506,7 +1398,7 @@ def compute0_warp_group( bars.mb_t_inv_ready[tinv1_idx].arrive() bars.mb_beta_done[beta1_idx].arrive() - # ---- A_epi (opposite member, both halves in parallel) ---------------- + # ---- A epilogue (opposite member, both halves in parallel) --------------- a0_idx = a_index.idx a0_phase = a_index.phase a_index = advance(a_index, cfg.smem_a_stages) @@ -1521,17 +1413,13 @@ def compute0_warp_group( if do_a: a_base = sA[my_a_idx].base - a_vec0 = nvvm.tcgen05_ld( - "16x256b", nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_cg0_acc_col + a_acc_idx * ACC_STAGE_COLS, cutlass.Float32), num=8 - ) - a_vec1 = nvvm.tcgen05_ld( - "16x256b", nvvm.make_tmem_ptr(((tmem_warp_row + 16) << 16) + tmem_cg0_acc_col + a_acc_idx * ACC_STAGE_COLS, cutlass.Float32), num=8 - ) + a_vec_lo = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_lo_addr + tmem_cg0_acc_col + a_acc_idx * ACC_STAGE_COLS, cutlass.Float32), num=8) + a_vec_hi = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_hi_addr + tmem_cg0_acc_col + a_acc_idx * ACC_STAGE_COLS, cutlass.Float32), num=8) nvvm.tcgen05_wait("load") bars.mb_cg0_acc_done[a_acc_idx].arrive() bars.mb_a_done[my_a_idx].wait(my_a_phase) for u in cutlass.range_constexpr(2): - a_vec = a_vec1 if cutlass.const_expr(u == 1) else a_vec0 + a_vec = a_vec_hi if cutlass.const_expr(u == 1) else a_vec_lo a_pack = [] for k in cutlass.range_constexpr(num_vals // 2): p0, p1 = fmul2(a_vec[2 * k], a_vec[2 * k + 1], decay_t_a[u * num_vals + 2 * k], decay_t_a[u * num_vals + 2 * k + 1]) @@ -1540,18 +1428,14 @@ def compute0_warp_group( st_row = half_row_base + u * 16 + store_row_frag for c in cutlass.range_constexpr(ACC_N_FRAGS): nvvm.stmatrix( - cutlass.inttoptr( - a_base + (st_row * cfg.b_t + swizzle_xor_128b(st_row, store_col + c * FRAG_COLS)) * bpe, - cutlass.AddressSpace.smem, - cutlass.BFloat16, - ), + a_base + st_row * cfg.b_t + swizzle_xor_128b(st_row, store_col + c * FRAG_COLS), [a_pack[c * 4 + 0], a_pack[c * 4 + 1], a_pack[c * 4 + 2], a_pack[c * 4 + 3]], nvvm.MMALayout.ROW, ) nvvm.fence_proxy("async.shared", space="cta") bars.mb_a_ready[my_a_idx].arrive() - tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one) for _ in range(cfg.smem_t_inv_stages): bars.mb_t_inv_done[tinv_index.idx].wait(tinv_index.phase) tinv_index = advance(tinv_index, cfg.smem_t_inv_stages) @@ -1572,7 +1456,7 @@ def compute1_warp_group( warp_idx, tmem_base_slot, scale, - sV, + sV_trans, sCumsumlog, sCumprod, sBeta, @@ -1581,12 +1465,12 @@ def compute1_warp_group( mState_init, mState_out, checkpoint_every_n_tokens, - sSched, + sScheduler, bars, ): """Compute warp-group 1 role (warps 4-7): persistent scheduler loop running the per-chunk state-update and output epilogues.""" - elect_one = nvvm.elect_sync() + nvvm.setmaxregister(cfg.num_regs_compute_group_1, nvvm.SetMaxRegisterAction.INCREASE) v_index = PipelineState.start(phase=0) gate_index = PipelineState.start(phase=0) @@ -1595,388 +1479,368 @@ def compute1_warp_group( o_final_acc_ready_index = PipelineState.start(phase=0) k_state_ready_index = PipelineState.start(phase=0) u_acc_ready_index = PipelineState.start(phase=0) - state_acc_seed_index = PipelineState.start(phase=1) o_index = PipelineState.start(phase=1) - state_inp_cnt = cutlass.Int32(0) - kv_done_idx = cutlass.Int32(0) - - nvvm.setmaxregister(cfg.num_regs_compute_group_1, nvvm.SetMaxRegisterAction.INCREASE) - nvvm.barrier_cta_sync_aligned( - cfg.tmem_alloc_barrier_id, - thread_count=cfg.tmem_alloc_barrier_threads, - ) - tmem_base = tmem_base_slot.load() num_threads_cg1 = cfg.threads_per_warp * len(cfg.compute_group_1_warp_ids) cg1_tidx = tidx % num_threads_cg1 - lane_id = cg1_tidx % cfg.threads_per_warp - tmem_warp_row = (cg1_tidx // cfg.threads_per_warp) * cfg.threads_per_warp + lane_idx = cg1_tidx % cfg.threads_per_warp + + elect_one = nvvm.elect_sync() + state_input_cnt = cutlass.Int32(0) ldtm_width = 32 - sttm_width = ldtm_width // 2 - num_state_subs = cutlass.const_expr(cfg.d_v // ldtm_width) - tmem_state_col = tmem_base + cfg.tmem_state_acc_offset - tmem_state_inp_col = tmem_base + cfg.tmem_state_inp_offset - tmem_q_state_col = tmem_base + cfg.tmem_q_state_acc_offset - tmem_inp_col = tmem_base + cfg.tmem_y_decay_u_inp_offset + sttm_width = 16 + num_ldtms = cutlass.const_expr(cfg.d_v // ldtm_width) ACC_STAGE_COLS = cfg.b_t INP_SLOT_COLS = cfg.b_t // 2 - tmem_k_state_col = tmem_base + cfg.tmem_cg1_acc_offset - tmem_u_acc_col = tmem_k_state_col - tmem_y_inp_col = tmem_inp_col - tmem_u_inp_col = tmem_inp_col - tmem_decay_v_col = tmem_inp_col + INP_SLOT_COLS - v_o_smem_tok = cg1_tidx % 8 + (cg1_tidx // 16 % 2) * 8 - v_o_smem_col = (cg1_tidx // 8 % 2) * 8 + (cg1_tidx // 32 % 2) * 32 - v_o_smem_subtile_off = (cg1_tidx // 64) * 4096 - v_stage_elems = cfg.v_cosize // cfg.smem_v_stages - o_stage_elems = cfg.o_cosize // cfg.smem_o_stages - sV_base = cute.make_ptr(cfg.io_dtype, sV[0].base, mem_space=cute.AddressSpace.smem, assumed_align=cfg.buffer_align_bytes) - sO_base = cute.make_ptr(cfg.io_dtype, sO[0].base, mem_space=cute.AddressSpace.smem, assumed_align=cfg.buffer_align_bytes) + v_o_row = cg1_tidx % 8 + (cg1_tidx // 16 % 2) * 8 + v_o_col = (cg1_tidx // 8 % 2) * 8 + (cg1_tidx // 32 % 2) * 32 + v_o_segment = (cg1_tidx // 64) * 4096 + v_stage_elements = cfg.v_cosize // cfg.smem_v_stages + o_stage_elements = cfg.o_cosize // cfg.smem_o_stages + sV_base = sV_trans[0].base + sO_base = sO[0].base num_vals = 32 + + nvvm.barrier_cta_sync_aligned(cfg.tmem_lifecycle_barrier_id, thread_count=cfg.tmem_user_threads) + tmem_base = tmem_base_slot.load() + tmem_col = tmem_base & 0xFFFF + tmem_row = tmem_base >> 16 + row_lo_addr = tmem_row << 16 + row_hi_addr = (tmem_row + 16) << 16 + tmem_state_col = tmem_col + cfg.tmem_state_acc_offset + tmem_state_input_col = tmem_col + cfg.tmem_state_input_offset + tmem_q_state_col = tmem_col + cfg.tmem_q_state_acc_offset + tmem_input_col = tmem_col + cfg.tmem_y_decay_u_input_offset + tmem_k_state_col = tmem_col + cfg.tmem_cg1_acc_offset + tmem_u_acc_col = tmem_k_state_col + tmem_y_input_col = tmem_input_col + tmem_u_input_col = tmem_input_col + tmem_decay_v_col = tmem_input_col + INP_SLOT_COLS if cutlass.const_expr(cfg.enable_checkpoints): - sCheckpoint_base_int = sCheckpoint_raw.data_ptr().toint() + sCheckpoint_base = sCheckpoint_raw.data_ptr() checkpoint_cnt = cutlass.Int32(0) checkpoint_smem_row = cg1_tidx % 8 + (cg1_tidx // 16 % 2) * 8 checkpoint_smem_col = (cg1_tidx // 8 % 2) * 8 + (cg1_tidx // 32 % 2) * 32 - sched_state = PipelineState.start(phase=0) + scheduler_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) - n_local = wend - cstart + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) + n_local = write_end - compute_start if cutlass.const_expr(cfg.enable_checkpoints): - ckpt_chunks = checkpoint_every_n_tokens // cutlass.Int32(cfg.b_t) - checkpoint_mod = cstart % ckpt_chunks + checkpoint_chunks = checkpoint_every_n_tokens // cutlass.Int32(cfg.b_t) + checkpoint_mod = compute_start % checkpoint_chunks if n_local > 0: if cutlass.const_expr(cfg.use_initial_state): - # ---- initial-state seed: initial_state GMEM -> state TMEM --------------- gState_init = mState_init[None, None, head_idx, batch_idx] - kv_init_idx = state_acc_seed_index.idx - bars.mb_state_acc_scale_done[kv_init_idx].wait(state_acc_seed_index.phase) - state_acc_seed_index = advance(state_acc_seed_index, cfg.tmem_state_acc_stages) - seed_from_initial_state = cstart == 0 - if seed_from_initial_state: - for sub in cutlass.range_constexpr(num_state_subs): + seed_state = compute_start == 0 + if seed_state: + for i in cutlass.range_constexpr(num_ldtms): words = [] for k in cutlass.range_constexpr(32): - v = gState_init[cg1_tidx, sub * ldtm_width + k] + v = gState_init[cg1_tidx, i * ldtm_width + k] if cutlass.const_expr(cfg.state_dtype != cfg.acc_dtype): v = v.to(cfg.acc_dtype) words.append(v) nvvm.tcgen05_st( "32x32b", - nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_state_col + sub * ldtm_width, cutlass.Float32), + nvvm.make_tmem_ptr(row_lo_addr + tmem_state_col + i * ldtm_width, cutlass.Float32), cutlass.Vector.from_elements(tuple(words), cutlass.Float32), ) + nvvm.tcgen05_wait("store") else: - for sub in cutlass.range_constexpr(num_state_subs): + for i in cutlass.range_constexpr(num_ldtms): nvvm.tcgen05_st( "32x32b", - nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_state_col + sub * ldtm_width, cutlass.Float32), + nvvm.make_tmem_ptr(row_lo_addr + tmem_state_col + i * ldtm_width, cutlass.Float32), cutlass.Vector.from_elements(tuple(cutlass.Float32(0.0) for _ in range(32)), cutlass.Float32), ) - nvvm.tcgen05_wait("store") - - nvvm.barrier_cta_sync_aligned( - cfg.init_state_store_barrier_id, - thread_count=cfg.init_state_store_barrier_threads, - ) - - for local_idx in cutlass.range(n_local): # noqa: B007 - chunk_idx = cstart + local_idx - if cutlass.const_expr(cfg.enable_checkpoints): - do_checkpoint_now = checkpoint_mod == 0 - checkpoint_mod = checkpoint_mod + cutlass.Int32(1) - checkpoint_mod = cutlass.Int32(0) if checkpoint_mod == ckpt_chunks else checkpoint_mod - if cutlass.const_expr(cfg.enable_checkpoints and not cfg.use_initial_state): - if chunk_idx == 0 and wstart == 0: - checkpoint_stage = checkpoint_cnt % cfg.smem_checkpoint_stages - checkpoint_phase_done = cutlass.Int32(1) ^ ((checkpoint_cnt // cfg.smem_checkpoint_stages) & cutlass.Int32(1)) - bars.mb_checkpoint_tmastg_done[checkpoint_stage].wait(checkpoint_phase_done) - checkpoint_zero_ptr = cutlass.inttoptr( - sCheckpoint_base_int + checkpoint_stage * cfg.d_k * cfg.d_v * 2, cutlass.AddressSpace.smem, cutlass.Int32 - ) - for z in cutlass.range_constexpr(cfg.d_k * cfg.d_v // 2 // num_threads_cg1): - (checkpoint_zero_ptr + cg1_tidx + z * num_threads_cg1).store(cutlass.Int32(0)) - nvvm.fence_proxy("async.shared", space="cta") - if elect_one: - bars.mb_checkpoint_tmastg_ready[checkpoint_stage].arrive() - checkpoint_cnt = checkpoint_cnt + 1 - valid_state = local_idx > 0 - if cutlass.const_expr(cfg.use_initial_state): - valid_state = cutlass.Boolean(True) - state_acc_seed_index = advance(state_acc_seed_index, cfg.tmem_state_acc_stages) - - gate_idx = gate_index.idx - bars.mb_gate_ready[gate_idx].wait(gate_index.phase) - gate_index = advance(gate_index, cfg.smem_gate_stages) - cumprod_total = sCumprod[sCumprod.shape[0] - 1, 0, gate_idx] - - # ---- state restage + rescale ------------------------------------- - if valid_state: - kv_idx = kv_acc_index.idx - bars.mb_state_acc_ready[kv_idx].wait(kv_acc_index.phase) - kv_acc_index = advance(kv_acc_index, cfg.tmem_state_acc_stages) - kv_done_idx = kv_idx - - state_regs = [[cutlass.Float32(0.0) for _ in range(num_state_subs)] for _ in range(32)] - state_inp_stage_idx = state_inp_cnt % cfg.tmem_state_inp_stages - state_vecs = [] - for sub in cutlass.range_constexpr(num_state_subs): - state_vecs.append( - nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_state_col + sub * ldtm_width, cutlass.Float32), num=32) - ) - for sub in cutlass.range_constexpr(num_state_subs): - for k in cutlass.range_constexpr(32): - state_regs[k][sub] = state_vecs[sub][k] - state_pack = [fp32_to_fp16(state_regs[2 * j][sub], state_regs[2 * j + 1][sub], dtype=cfg.io_dtype) for j in range(16)] - nvvm.tcgen05_st( - "32x32b", - nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_state_inp_col + sub * sttm_width, cutlass.Int32), - cutlass.Vector.from_elements(tuple(state_pack), cutlass.Int32), - ) nvvm.tcgen05_wait("store") - bars.mb_state_inp_ready[state_inp_stage_idx].arrive() - state_inp_cnt = state_inp_cnt + 1 + for peel in cutlass.range_constexpr(1 if cfg.use_initial_state else 2): + peel_stop = n_local if cutlass.const_expr(peel == 1 or cfg.use_initial_state) else min(n_local, 1) + for local_idx in cutlass.range(peel, peel_stop, 1): # noqa: B007 + chunk_idx = compute_start + local_idx if cutlass.const_expr(cfg.enable_checkpoints): - # ---- state checkpoint ---------------------------------------- - do_checkpoint = do_checkpoint_now and chunk_idx < wend - do_checkpoint = do_checkpoint and chunk_idx >= wstart - if do_checkpoint: + do_checkpoint_now = checkpoint_mod == 0 + checkpoint_mod = checkpoint_mod + cutlass.Int32(1) + checkpoint_mod = cutlass.Int32(0) if checkpoint_mod == checkpoint_chunks else checkpoint_mod + if cutlass.const_expr(cfg.enable_checkpoints and not cfg.use_initial_state): + if chunk_idx == 0 and write_start == 0: checkpoint_stage = checkpoint_cnt % cfg.smem_checkpoint_stages checkpoint_phase_done = cutlass.Int32(1) ^ ((checkpoint_cnt // cfg.smem_checkpoint_stages) & cutlass.Int32(1)) bars.mb_checkpoint_tmastg_done[checkpoint_stage].wait(checkpoint_phase_done) - checkpoint_stage_base = checkpoint_stage * cfg.d_k * cfg.d_v - for sub in cutlass.range_constexpr(num_state_subs): - checkpoint_vec = nvvm.tcgen05_ld( - "32x32b", - nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_state_col + sub * ldtm_width, cutlass.Float32), - num=32, - ) - for g in cutlass.range_constexpr(ldtm_width // 8): - packs = tuple( - fp32_to_fp16(checkpoint_vec[g * 8 + 2 * t], checkpoint_vec[g * 8 + 2 * t + 1], dtype=cfg.io_dtype) for t in range(4) - ) - dk = sub * ldtm_width + g * 8 - checkpoint_addr = checkpoint_stage_base + (dk // 64) * (cfg.d_v * 64) + cg1_tidx * 64 + swizzle_xor_128b(cg1_tidx, dk % 64) - (sCheckpoint_raw.data_ptr() + checkpoint_addr).store( - cutlass.Vector.from_elements(packs, cutlass.Int32).bitcast(cfg.io_dtype), alignment=16 - ) + checkpoint_zero_ptr = cutlass.inttoptr( + (sCheckpoint_base + checkpoint_stage * cfg.d_k * cfg.d_v).toint(), cutlass.AddressSpace.smem, cutlass.Int32 + ) + for z in cutlass.range_constexpr(cfg.d_k * cfg.d_v // 2 // num_threads_cg1): + (checkpoint_zero_ptr + cg1_tidx + z * num_threads_cg1).store(cutlass.Int32(0)) nvvm.fence_proxy("async.shared", space="cta") if elect_one: bars.mb_checkpoint_tmastg_ready[checkpoint_stage].arrive() checkpoint_cnt = checkpoint_cnt + 1 - - for sub in cutlass.range_constexpr(num_state_subs): - state_scaled = [] - for j in cutlass.range_constexpr(16): - s0, s1 = fmul2(state_regs[2 * j][sub], state_regs[2 * j + 1][sub], cumprod_total, cumprod_total) - state_scaled += [s0, s1] + valid_state = cutlass.Boolean(True) if cutlass.const_expr(cfg.use_initial_state) else cutlass.const_expr(peel == 1) + + gate_idx = gate_index.idx + bars.mb_gate_ready[gate_idx].wait(gate_index.phase) + gate_index = advance(gate_index, cfg.smem_gate_stages) + cumprod_total = sCumprod[sCumprod.shape[0] - 1, 0, gate_idx] + + # ---- state stage + rescale --------------------------------------- + if valid_state: + if cutlass.const_expr(cfg.use_initial_state): + if local_idx > 0: + bars.mb_state_acc_ready[kv_acc_index.idx].wait(kv_acc_index.phase) + kv_acc_index = advance(kv_acc_index, cfg.tmem_state_acc_stages) + else: + bars.mb_state_acc_ready[kv_acc_index.idx].wait(kv_acc_index.phase) + kv_acc_index = advance(kv_acc_index, cfg.tmem_state_acc_stages) + + state_input_stage_idx = state_input_cnt % cfg.tmem_state_input_stages + state_regs = [[cutlass.Float32(0.0) for _ in range(num_ldtms)] for _ in range(32)] + for i in cutlass.range_constexpr(num_ldtms): + state_vec = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_lo_addr + tmem_state_col + i * ldtm_width, cutlass.Float32), num=32) + for k in cutlass.range_constexpr(32): + state_regs[k][i] = state_vec[k] + for i in cutlass.range_constexpr(num_ldtms): + state_pack = [fp32_to_fp16(state_regs[2 * j][i], state_regs[2 * j + 1][i], dtype=cfg.io_dtype) for j in range(16)] + nvvm.tcgen05_st( + "32x32b", + nvvm.make_tmem_ptr(row_lo_addr + tmem_state_input_col + i * sttm_width, cutlass.Int32), + cutlass.Vector.from_elements(tuple(state_pack), cutlass.Int32), + ) + nvvm.tcgen05_wait("store") + bars.mb_state_input_ready[state_input_stage_idx].arrive() + state_input_cnt = state_input_cnt + 1 + + if cutlass.const_expr(cfg.enable_checkpoints): + # ---- state checkpoint ------------------------------------ + do_checkpoint = do_checkpoint_now and chunk_idx < write_end + do_checkpoint = do_checkpoint and chunk_idx >= write_start + if do_checkpoint: + checkpoint_stage = checkpoint_cnt % cfg.smem_checkpoint_stages + checkpoint_phase_done = cutlass.Int32(1) ^ ((checkpoint_cnt // cfg.smem_checkpoint_stages) & cutlass.Int32(1)) + bars.mb_checkpoint_tmastg_done[checkpoint_stage].wait(checkpoint_phase_done) + checkpoint_stage_base = checkpoint_stage * cfg.d_k * cfg.d_v + for i in cutlass.range_constexpr(num_ldtms): + for g in cutlass.range_constexpr(ldtm_width // 8): + packs = tuple( + fp32_to_fp16(state_regs[g * 8 + 2 * t][i], state_regs[g * 8 + 2 * t + 1][i], dtype=cfg.io_dtype) for t in range(4) + ) + dk = i * ldtm_width + g * 8 + checkpoint_addr = ( + checkpoint_stage_base + (dk // 64) * (cfg.d_v * 64) + cg1_tidx * 64 + swizzle_xor_128b(cg1_tidx, dk % 64) + ) + (sCheckpoint_raw.data_ptr() + checkpoint_addr).store( + cutlass.Vector.from_elements(packs, cutlass.Int32).bitcast(cfg.io_dtype), alignment=16 + ) + nvvm.fence_proxy("async.shared", space="cta") + if elect_one: + bars.mb_checkpoint_tmastg_ready[checkpoint_stage].arrive() + checkpoint_cnt = checkpoint_cnt + 1 + + for i in cutlass.range_constexpr(num_ldtms): + state_scaled = [] + for j in cutlass.range_constexpr(16): + s0, s1 = fmul2(state_regs[2 * j][i], state_regs[2 * j + 1][i], cumprod_total, cumprod_total) + state_scaled += [s0, s1] + nvvm.tcgen05_st( + "32x32b", + nvvm.make_tmem_ptr(row_lo_addr + tmem_state_col + i * ldtm_width, cutlass.Float32), + cutlass.Vector.from_elements(tuple(state_scaled), cutlass.Float32), + ) + nvvm.tcgen05_wait("store") + + # ---- per-row Gate register builds -------------------------------- + cumprod_vals = [] + for k in cutlass.range_constexpr(num_vals): + cumprod_vals.append(sCumprod[(lane_idx % 4) * 2 + ((k // 4) * 8 + k % 2), 0, gate_idx]) + last_cumsumlog = sCumsumlog[cfg.b_t - 1, 0, gate_idx] + cumsumlog_vals = [] + for k in cutlass.range_constexpr(num_vals): + cumsumlog_vals.append(sCumsumlog[(lane_idx % 4) * 2 + ((k // 4) * 8 + k % 2), 0, gate_idx]) + decay_scale_vals = [] + for k in cutlass.range_constexpr(0, num_vals, 2): + d0, d1 = fadd2(last_cumsumlog, last_cumsumlog, -cumsumlog_vals[k], -cumsumlog_vals[k + 1]) + decay_scale_vals.append(cute.math.exp2(d0, fastmath=True)) + decay_scale_vals.append(cute.math.exp2(d1, fastmath=True)) + bars.mb_gate_done[gate_idx].arrive() + + # ---- Y = V - k state (packed 16-bit) ----------------------------- + v_idx = v_index.idx + bars.mb_v_ready[v_idx].wait(v_index.phase) + v_index = advance(v_index, cfg.smem_v_stages) + + v_frag = [] + for half in cutlass.range_constexpr(2): + v_words = [] + for block in cutlass.range_constexpr(4): + v_raw = nvvm.ldmatrix( + ( + sV_base + + v_idx * v_stage_elements + + v_o_segment + + (v_o_row + block * 16) * 64 + + swizzle_xor_128b(v_o_row + block * 16, v_o_col + half * 16) + ), + 4, + nvvm.MMALayout.COL, + ) + for i in cutlass.range_constexpr(4): + v_words.append(v_raw[i]) + v_frag.append(v_words) + if valid_state: + bars.mb_k_state_acc_ready[0].wait(k_state_ready_index.phase) + k_state_ready_index = advance(k_state_ready_index, 1) + + for half in cutlass.range_constexpr(2): + k_state_vec = nvvm.tcgen05_ld( + "16x256b", + nvvm.make_tmem_ptr(((tmem_row + half * 16) << 16) + tmem_k_state_col, cutlass.Float32), + num=8, + ) + for j in cutlass.range_constexpr(16): + s0, s1 = fmul2(k_state_vec[2 * j], k_state_vec[2 * j + 1], cumprod_vals[2 * j], cumprod_vals[2 * j + 1]) + k_state_word = fp32_to_fp16(s0, s1, dtype=cfg.io_dtype) + v_frag[half][j] = sub_f16x2(v_frag[half][j], k_state_word, cfg.io_dtype) + for half in cutlass.range_constexpr(2): nvvm.tcgen05_st( - "32x32b", - nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_state_col + sub * ldtm_width, cutlass.Float32), - cutlass.Vector.from_elements(tuple(state_scaled), cutlass.Float32), + "16x128b", + nvvm.make_tmem_ptr(((tmem_row + half * 16) << 16) + tmem_y_input_col, cutlass.Int32), + cutlass.Vector.from_elements(tuple(v_frag[half]), cutlass.Int32), ) nvvm.tcgen05_wait("store") - bars.mb_state_acc_scale_done[kv_idx].arrive() - - # ---- per-row Gate register builds -------------------------------- - cumprod_vals = [] - for k in cutlass.range_constexpr(num_vals): - cumprod_vals.append(sCumprod[(lane_id % 4) * 2 + ((k // 4) * 8 + k % 2), 0, gate_idx]) - last_cumsumlog = sCumsumlog[cfg.b_t - 1, 0, gate_idx] - cumsumlog_vals = [] - for k in cutlass.range_constexpr(num_vals): - cumsumlog_vals.append(sCumsumlog[(lane_id % 4) * 2 + ((k // 4) * 8 + k % 2), 0, gate_idx]) - decay_scale_vals = [] - for k in cutlass.range_constexpr(0, num_vals, 2): - d0, d1 = fadd2(last_cumsumlog, last_cumsumlog, -cumsumlog_vals[k], -cumsumlog_vals[k + 1]) - decay_scale_vals.append(cute.math.exp2(d0, fastmath=True)) - decay_scale_vals.append(cute.math.exp2(d1, fastmath=True)) - bars.mb_gate_done[gate_idx].arrive() - - # ---- Y = V - K*state (packed 16-bit) ----------------------------- - v_idx = v_index.idx - bars.mb_v_ready[v_idx].wait(v_index.phase) - v_index = advance(v_index, cfg.smem_v_stages) - - v_frags = [[cutlass.Int32(0), cutlass.Int32(0)] for _ in range(16)] - for c in cutlass.range_constexpr(8): - m0 = cutlass.const_expr(c % 4) - sub = cutlass.const_expr(c // 4) - v_frag = nvvm.ldmatrix( - ( - sV_base - + v_idx * v_stage_elems - + v_o_smem_subtile_off - + (v_o_smem_tok + m0 * 16) * 64 - + swizzle_xor_128b(v_o_smem_tok + m0 * 16, v_o_smem_col + sub * 16) - ).raw_ptr(), - 4, - nvvm.MMALayout.COL, - ) - for i in cutlass.range_constexpr(4): - v_frags[4 * m0 + i][sub] = v_frag[i] - if valid_state: - bars.mb_k_state_acc_ready[0].wait(k_state_ready_index.phase) - k_state_ready_index = advance(k_state_ready_index, 1) - - for sub in cutlass.range_constexpr(2): - k_state_vec = nvvm.tcgen05_ld( - "16x256b", - nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_k_state_col, cutlass.Float32), - num=8, + bars.mb_y_input_ready[0].arrive() + + # ---- q state epilogue: q state *= cumprod * scale ---------------- + if valid_state: + q_state_idx = o_acc_ready_index.idx + bars.mb_o_acc_ready[q_state_idx].wait(o_acc_ready_index.phase) + o_acc_ready_index = advance(o_acc_ready_index, cfg.tmem_q_state_acc_stages) + + q_state_ptrs = [] + q_state_vecs = [] + for half in cutlass.range_constexpr(2): + q_state_ptrs.append( + nvvm.make_tmem_ptr(((tmem_row + half * 16) << 16) + tmem_q_state_col + q_state_idx * ACC_STAGE_COLS, cutlass.Float32) + ) + q_state_vecs.append(nvvm.tcgen05_ld("16x256b", q_state_ptrs[half], num=8)) + for half in cutlass.range_constexpr(2): + q_state_scaled = [] + for j in cutlass.range_constexpr(16): + p0, p1 = fmul2(q_state_vecs[half][2 * j], q_state_vecs[half][2 * j + 1], cumprod_vals[2 * j], cumprod_vals[2 * j + 1]) + s0, s1 = fmul2(p0, p1, scale, scale) + q_state_scaled += [s0, s1] + nvvm.tcgen05_st("16x256b", q_state_ptrs[half], cutlass.Vector.from_elements(tuple(q_state_scaled), cutlass.Float32)) + nvvm.tcgen05_wait("store") + bars.mb_o_state_scale_acc_done[q_state_idx].arrive() + + # ---- U epilogue + decayed-U publish ------------------------------ + bars.mb_u_acc_ready[0].wait(u_acc_ready_index.phase) + u_acc_ready_index = advance(u_acc_ready_index, 1) + bars.mb_v_done[v_idx].arrive() + + u_vecs = [] + for half in cutlass.range_constexpr(2): + u_vecs.append( + nvvm.tcgen05_ld( + "16x256b", + nvvm.make_tmem_ptr(((tmem_row + half * 16) << 16) + tmem_u_acc_col, cutlass.Float32), + num=8, + ) ) - for j in cutlass.range_constexpr(16): - s0, s1 = fmul2(k_state_vec[2 * j], k_state_vec[2 * j + 1], cumprod_vals[2 * j], cumprod_vals[2 * j + 1]) - k_state_word = fp32_to_fp16(s0, s1, dtype=cfg.io_dtype) - v_frags[j][sub] = sub_f16x2(v_frags[j][sub], k_state_word, cfg.io_dtype) - for sub in cutlass.range_constexpr(2): - nvvm.tcgen05_st( - "16x128b", - nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_y_inp_col, cutlass.Int32), - cutlass.Vector.from_elements(tuple(v_frags[j][sub] for j in range(16)), cutlass.Int32), - ) - nvvm.tcgen05_wait("store") - bars.mb_y_inp_ready[0].arrive() - - # ---- state*Q_epi: Q*state *= cumprod * scale ---------------- - if valid_state: - q_state_idx = o_acc_ready_index.idx - bars.mb_o_acc_ready[q_state_idx].wait(o_acc_ready_index.phase) - o_acc_ready_index = advance(o_acc_ready_index, cfg.tmem_q_state_acc_stages) - - q_state_ptrs = [] - q_state_vecs = [] - for sub in cutlass.range_constexpr(2): - q_state_ptrs.append( - nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_q_state_col + q_state_idx * ACC_STAGE_COLS, cutlass.Float32) + u_regs = [[u_vecs[0][k], u_vecs[1][k]] for k in range(32)] + for half in cutlass.range_constexpr(2): + u_pack = [fp32_to_fp16(u_regs[2 * j][half], u_regs[2 * j + 1][half], dtype=cfg.io_dtype) for j in range(16)] + nvvm.tcgen05_st( + "16x128b", + nvvm.make_tmem_ptr(((tmem_row + half * 16) << 16) + tmem_u_input_col, cutlass.Int32), + cutlass.Vector.from_elements(tuple(u_pack), cutlass.Int32), ) - q_state_vecs.append(nvvm.tcgen05_ld("16x256b", q_state_ptrs[sub], num=8)) - for sub in cutlass.range_constexpr(2): - q_state_scaled = [] - for j in cutlass.range_constexpr(16): - p0, p1 = fmul2(q_state_vecs[sub][2 * j], q_state_vecs[sub][2 * j + 1], cumprod_vals[2 * j], cumprod_vals[2 * j + 1]) - s0, s1 = fmul2(p0, p1, scale, scale) - q_state_scaled += [s0, s1] - nvvm.tcgen05_st("16x256b", q_state_ptrs[sub], cutlass.Vector.from_elements(tuple(q_state_scaled), cutlass.Float32)) nvvm.tcgen05_wait("store") + bars.mb_u_input_ready[0].arrive() - bars.mb_o_state_scale_acc_done[q_state_idx].arrive() + for half in cutlass.range_constexpr(2): + for j in cutlass.range_constexpr(16): + u_regs[2 * j][half], u_regs[2 * j + 1][half] = fmul2( + u_regs[2 * j][half], u_regs[2 * j + 1][half], decay_scale_vals[2 * j], decay_scale_vals[2 * j + 1] + ) + decay_pack = [fp32_to_fp16(u_regs[2 * j][half], u_regs[2 * j + 1][half], dtype=cfg.io_dtype) for j in range(16)] + nvvm.tcgen05_st( + "16x128b", + nvvm.make_tmem_ptr(((tmem_row + half * 16) << 16) + tmem_decay_v_col, cutlass.Int32), + cutlass.Vector.from_elements(tuple(decay_pack), cutlass.Int32), + ) + nvvm.tcgen05_wait("store") + bars.mb_decay_u_input_ready[0].arrive() - # ---- U_epi + decayed-U publish ----------------------------------- - bars.mb_u_acc_ready[0].wait(u_acc_ready_index.phase) - u_acc_ready_index = advance(u_acc_ready_index, 1) - bars.mb_v_done[v_idx].arrive() + # ---- output store: O acc TMEM -> SMEM ---------------------------- + o_scale_idx = o_final_acc_ready_index.idx + bars.mb_o_final_acc_ready[o_scale_idx].wait(o_final_acc_ready_index.phase) + o_final_acc_ready_index = advance(o_final_acc_ready_index, cfg.tmem_q_state_acc_stages) - u_regs = [[cutlass.Float32(0.0), cutlass.Float32(0.0)] for _ in range(32)] - u_vecs = [] - for sub in cutlass.range_constexpr(2): - u_vecs.append( - nvvm.tcgen05_ld( + o_regs = [] + for half in cutlass.range_constexpr(2): + o_vec = nvvm.tcgen05_ld( "16x256b", - nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_u_acc_col, cutlass.Float32), + nvvm.make_tmem_ptr(((tmem_row + half * 16) << 16) + tmem_q_state_col + o_scale_idx * ACC_STAGE_COLS, cutlass.Float32), num=8, ) - ) - for sub in cutlass.range_constexpr(2): - for k in cutlass.range_constexpr(32): - u_regs[k][sub] = u_vecs[sub][k] - - u_pack = [fp32_to_fp16(u_regs[2 * j][sub], u_regs[2 * j + 1][sub], dtype=cfg.io_dtype) for j in range(16)] - nvvm.tcgen05_st( - "16x128b", - nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_u_inp_col, cutlass.Int32), - cutlass.Vector.from_elements(tuple(u_pack), cutlass.Int32), - ) - nvvm.tcgen05_wait("store") - bars.mb_u_inp_ready[0].arrive() - - for sub in cutlass.range_constexpr(2): - for j in cutlass.range_constexpr(16): - u_regs[2 * j][sub], u_regs[2 * j + 1][sub] = fmul2( - u_regs[2 * j][sub], u_regs[2 * j + 1][sub], decay_scale_vals[2 * j], decay_scale_vals[2 * j + 1] - ) - decay_pack = [fp32_to_fp16(u_regs[2 * j][sub], u_regs[2 * j + 1][sub], dtype=cfg.io_dtype) for j in range(16)] - nvvm.tcgen05_st( - "16x128b", - nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_decay_v_col, cutlass.Int32), - cutlass.Vector.from_elements(tuple(decay_pack), cutlass.Int32), - ) - nvvm.tcgen05_wait("store") - bars.mb_decay_u_inp_ready[0].arrive() - - # ---- QKV_epilogue: O acc TMEM -> sO SMEM ------------------------- - o_scale_idx = o_final_acc_ready_index.idx - bars.mb_o_final_acc_ready[o_scale_idx].wait(o_final_acc_ready_index.phase) - o_final_acc_ready_index = advance(o_final_acc_ready_index, cfg.tmem_q_state_acc_stages) - - o_regs = [] - for sub in cutlass.range_constexpr(2): - o_vec = nvvm.tcgen05_ld( - "16x256b", - nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_q_state_col + o_scale_idx * ACC_STAGE_COLS, cutlass.Float32), - num=8, - ) - o_regs.append([o_vec[k] for k in range(32)]) - nvvm.tcgen05_wait("load") - o_idx = o_index.idx - bars.mb_o_tmastg_done[o_idx].wait(o_index.phase) - o_index = advance(o_index, cfg.smem_o_stages) - for sub in cutlass.range_constexpr(2): - for m0 in cutlass.range_constexpr(4): - o_pack = [fp32_to_fp16(o_regs[sub][8 * m0 + 2 * j], o_regs[sub][8 * m0 + 2 * j + 1], dtype=cfg.io_dtype) for j in range(4)] - nvvm.stmatrix( - ( - sO_base - + o_idx * o_stage_elems - + v_o_smem_subtile_off - + (v_o_smem_tok + m0 * 16) * 64 - + swizzle_xor_128b(v_o_smem_tok + m0 * 16, v_o_smem_col + sub * 16) - ).raw_ptr(), - o_pack, - nvvm.MMALayout.COL, - ) - nvvm.fence_proxy("async.shared", space="cta") - bars.mb_o_tmastg_ready[o_idx].arrive() - - # ---- final state: state TMEM -> GMEM ----------------------------------- + o_regs.append([o_vec[k] for k in range(32)]) + nvvm.tcgen05_wait("load") + o_idx = o_index.idx + bars.mb_o_tmastg_done[o_idx].wait(o_index.phase) + o_index = advance(o_index, cfg.smem_o_stages) + for half in cutlass.range_constexpr(2): + for block in cutlass.range_constexpr(4): + o_pack = [fp32_to_fp16(o_regs[half][8 * block + 2 * j], o_regs[half][8 * block + 2 * j + 1], dtype=cfg.io_dtype) for j in range(4)] + nvvm.stmatrix( + ( + sO_base + + o_idx * o_stage_elements + + v_o_segment + + (v_o_row + block * 16) * 64 + + swizzle_xor_128b(v_o_row + block * 16, v_o_col + half * 16) + ), + o_pack, + nvvm.MMALayout.COL, + ) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_o_tmastg_ready[o_idx].arrive() + + # ---- final state store: TMEM -> GMEM ----------------------------------------- if n_local > 0: kv_last_idx = kv_acc_index.idx bars.mb_state_acc_ready[kv_last_idx].wait(kv_acc_index.phase) kv_acc_index = advance(kv_acc_index, cfg.tmem_state_acc_stages) if cutlass.const_expr(cfg.store_final_state): - if wend == num_chunks_b: + if write_end == batch_num_chunks: gState_out = mState_out[None, None, head_idx, batch_idx] - for sub in cutlass.range_constexpr(num_state_subs): - state_vec = nvvm.tcgen05_ld( - "32x32b", nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_state_col + sub * ldtm_width, cutlass.Float32), num=32 - ) + for i in cutlass.range_constexpr(num_ldtms): + state_vec = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_lo_addr + tmem_state_col + i * ldtm_width, cutlass.Float32), num=32) for k in cutlass.range_constexpr(32): val = state_vec[k] if cutlass.const_expr(cfg.state_dtype != cfg.acc_dtype): val = val.to(cfg.state_dtype) - gState_out[cg1_tidx, sub * ldtm_width + k] = val - bars.mb_state_acc_scale_done[kv_last_idx].arrive() - else: - bars.mb_state_acc_scale_done[kv_last_idx].arrive() + gState_out[cg1_tidx, i * ldtm_width + k] = val else: if cutlass.const_expr(cfg.store_final_state): - write_passthrough = wend == num_chunks_b + write_passthrough = write_end == batch_num_chunks if write_passthrough: gState_out = mState_out[None, None, head_idx, batch_idx] if cutlass.const_expr(cfg.use_initial_state): gState_in = mState_init[None, None, head_idx, batch_idx] - for sub in cutlass.range_constexpr(num_state_subs): + for i in cutlass.range_constexpr(num_ldtms): for k in cutlass.range_constexpr(32): - gState_out[cg1_tidx, sub * ldtm_width + k] = gState_in[cg1_tidx, sub * ldtm_width + k] + gState_out[cg1_tidx, i * ldtm_width + k] = gState_in[cg1_tidx, i * ldtm_width + k] else: - for sub in cutlass.range_constexpr(num_state_subs): + for i in cutlass.range_constexpr(num_ldtms): for k in cutlass.range_constexpr(32): - gState_out[cg1_tidx, sub * ldtm_width + k] = cutlass.Float32(0.0).to(cfg.state_dtype) + gState_out[cg1_tidx, i * ldtm_width + k] = cutlass.Float32(0.0).to(cfg.state_dtype) - tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one) bars.mb_tmem_done[0].arrive() @@ -1999,7 +1863,7 @@ def build_descs_body( base_v, base_o, base_checkpoint, - desc_ws: cute.Tensor, + desc_workspace: cute.Tensor, cu_seqlens: cute.Tensor, q: cute.Tensor, k: cute.Tensor, @@ -2018,11 +1882,11 @@ def build_descs_body( prologue kernel after its order pass; warps past the array count fall through the widx guards.""" arr_words = n_batch * cutlass.Int32(TENSOR_MAP_QWORDS) - sub0 = cute.make_tensor(desc_ws.iterator, cute.make_layout((arr_words,), stride=(1,))) - sub1 = cute.make_tensor(desc_ws.iterator + arr_words, cute.make_layout((arr_words,), stride=(1,))) - sub2 = cute.make_tensor(desc_ws.iterator + 2 * arr_words, cute.make_layout((arr_words,), stride=(1,))) - sub3 = cute.make_tensor(desc_ws.iterator + 3 * arr_words, cute.make_layout((arr_words,), stride=(1,))) - sub4 = cute.make_tensor(desc_ws.iterator + 4 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub0 = cute.make_tensor(desc_workspace.iterator, cute.make_layout((arr_words,), stride=(1,))) + sub1 = cute.make_tensor(desc_workspace.iterator + arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub2 = cute.make_tensor(desc_workspace.iterator + 2 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub3 = cute.make_tensor(desc_workspace.iterator + 3 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub4 = cute.make_tensor(desc_workspace.iterator + 4 * arr_words, cute.make_layout((arr_words,), stride=(1,))) if widx == 0: if nvvm.elect_sync(): @@ -2051,14 +1915,14 @@ def build_descs_body( @cute.kernel def prologue_kernel( order_gen: cutlass.Constexpr[bool], - has_sched: cutlass.Constexpr[bool], + has_scheduler: cutlass.Constexpr[bool], b_t: cutlass.Constexpr[int], base_q: cutlass.GridConstant[tma.TensorMap], base_k: cutlass.GridConstant[tma.TensorMap], base_v: cutlass.GridConstant[tma.TensorMap], base_o: cutlass.GridConstant[tma.TensorMap], base_checkpoint: cutlass.GridConstant[tma.TensorMap], - desc_ws: cute.Tensor, + desc_workspace: cute.Tensor, cu_seqlens: cute.Tensor, q: cute.Tensor, k: cute.Tensor, @@ -2068,7 +1932,7 @@ def prologue_kernel( mStaging: Optional[cute.Tensor], mCount: cute.Tensor, mWorkItems: cute.Tensor, - mSched: Optional[cute.Tensor], + mScheduler: Optional[cute.Tensor], n_batch: cutlass.Int32, q_row_stride: cutlass.Int32, k_row_stride: cutlass.Int32, @@ -2077,7 +1941,7 @@ def prologue_kernel( checkpoint_row_stride: cutlass.Int32, checkpoint_every_n: cutlass.Int32, ) -> None: - """Single-CTA prologue: LPT-order the work-item table and zero the sched + """Single-CTA prologue: LPT-order the work-item table and zero the scheduler rings via :func:`order_body`, then build the per-batch TMA-descriptor arrays via :func:`build_descs_body`, one warp per array (the extra warps only take part in the order phase).""" @@ -2090,10 +1954,10 @@ def prologue_kernel( n_heads_out = cutlass.Int32(q.shape[1] if q.shape[1] >= v.shape[1] else v.shape[1]) order_body( order_gen, - has_sched, + has_scheduler, b_t, ORDER_THREADS, - ORDER_ELEMS, + ORDER_ELEMENTS, tidx, n_heads_out, n_heads_out * n_batch, @@ -2101,7 +1965,7 @@ def prologue_kernel( mStaging, mCount, mWorkItems, - mSched, + mScheduler, sKey, sIdx, sSpread, @@ -2113,7 +1977,7 @@ def prologue_kernel( base_v, base_o, base_checkpoint, - desc_ws, + desc_workspace, cu_seqlens, q, k, @@ -2135,7 +1999,7 @@ def prologue( io_dtype: cutlass.Constexpr, b_t: cutlass.Constexpr[int], order_gen: cutlass.Constexpr[bool], - has_sched: cutlass.Constexpr[bool], + has_scheduler: cutlass.Constexpr[bool], q: cute.Tensor, k: cute.Tensor, v: cute.Tensor, @@ -2145,7 +2009,7 @@ def prologue( work_item_staging: Optional[cute.Tensor], work_count: cute.Tensor, work_items: cute.Tensor, - sched_ctr: Optional[cute.Tensor], + scheduler_counter: Optional[cute.Tensor], checkpoint_every_n: cutlass.Int32, tensormap_workspace: cute.Tensor, stream: cuda.CUstream, @@ -2160,7 +2024,7 @@ def prologue( heads_out = h_q if h_q >= h_v else h_v d_v = v.shape[2] bpe = io_dtype.width // 8 - elems_per_128b = 128 // bpe + elements_per_128b = 128 // bpe bt = b_t q_row_stride, q_head_stride = q.stride[0], q.stride[1] @@ -2173,20 +2037,20 @@ def prologue( k_headed = cute.make_tensor(k.iterator, cute.make_layout((seqlen, h_k, d_k), stride=(k_row_stride, k_head_stride, 1))) v_headed = cute.make_tensor(v.iterator, cute.make_layout((d_v, h_v, seqlen), stride=(1, v_head_stride, v_row_stride))) swz128 = tma.TensorMapSwizzle.s128b - base_desc_q = tma.create_tensor_map_tiled_from_view(q_headed, box_dims=(bt, 1, elems_per_128b), stride_order=(2, 1, 0), swizzle=swz128) - base_desc_k = tma.create_tensor_map_tiled_from_view(k_headed, box_dims=(bt, 1, elems_per_128b), stride_order=(2, 1, 0), swizzle=swz128) - base_desc_v = tma.create_tensor_map_tiled_from_view(v_headed, box_dims=(elems_per_128b, 1, bt), stride_order=(0, 1, 2), swizzle=swz128) + base_desc_q = tma.create_tensor_map_tiled_from_view(q_headed, box_dims=(bt, 1, elements_per_128b), stride_order=(2, 1, 0), swizzle=swz128) + base_desc_k = tma.create_tensor_map_tiled_from_view(k_headed, box_dims=(bt, 1, elements_per_128b), stride_order=(2, 1, 0), swizzle=swz128) + base_desc_v = tma.create_tensor_map_tiled_from_view(v_headed, box_dims=(elements_per_128b, 1, bt), stride_order=(0, 1, 2), swizzle=swz128) base_desc_o = base_desc_v if cutlass.const_expr(o is not None): o_headed = cute.make_tensor(o.iterator, cute.make_layout((d_v, heads_out, seqlen), stride=(1, o.stride[1], o.stride[0]))) - base_desc_o = tma.create_tensor_map_tiled_from_view(o_headed, box_dims=(elems_per_128b, 1, bt), stride_order=(0, 1, 2), swizzle=swz128) + base_desc_o = tma.create_tensor_map_tiled_from_view(o_headed, box_dims=(elements_per_128b, 1, bt), stride_order=(0, 1, 2), swizzle=swz128) base_desc_checkpoint = base_desc_v if cutlass.const_expr(state_checkpoints_out is not None): d_k_state = state_checkpoints_out.shape[2] d_v_state = state_checkpoints_out.shape[3] - checkpoint_elems_per_128b = 128 // (state_checkpoints_out.element_type.width // 8) + checkpoint_elements_per_128b = 128 // (state_checkpoints_out.element_type.width // 8) checkpoint_view = cute.make_tensor( state_checkpoints_out.iterator, cute.make_layout( @@ -2195,12 +2059,12 @@ def prologue( ), ) base_desc_checkpoint = tma.create_tensor_map_tiled_from_view( - checkpoint_view, box_dims=(checkpoint_elems_per_128b, d_k_state, 1, 1), stride_order=(0, 1, 2, 3), swizzle=swz128 + checkpoint_view, box_dims=(checkpoint_elements_per_128b, d_k_state, 1, 1), stride_order=(0, 1, 2, 3), swizzle=swz128 ) prologue_kernel( order_gen, - has_sched, + has_scheduler, b_t, base_desc_q, base_desc_k, @@ -2217,7 +2081,7 @@ def prologue( work_item_staging, work_count, work_items, - sched_ctr, + scheduler_counter, cutlass.Int32(batch_size), cutlass.Int32(q_row_stride), cutlass.Int32(k_row_stride), @@ -2244,19 +2108,19 @@ def host( state_out: Optional[cute.Tensor], work_items: Optional[cute.Tensor], work_count: Optional[cute.Tensor], - sched_ctr: Optional[cute.Tensor], + scheduler_counter: Optional[cute.Tensor], checkpoint_every_n_tokens: cutlass.Int32, scale: cutlass.Float32, tensormap_workspace: cute.Tensor, stream: cuda.CUstream, ): - h_q = q.shape[1] - h_k = k.shape[1] - h_v = v.shape[1] + h_q = cfg.h_q + h_k = cfg.h_k + h_v = cfg.h_v batch_size = cu_seqlens.shape[0] - 1 heads_out = h_q if h_q >= h_v else h_v - # ---- GQA reshapes: fold the head group into a -------------------------------- + # ---- GQA reshapes: fold the head group into the Q head axis ---------------------- if cutlass.const_expr(cfg.is_GQA): h_r = h_q // h_v h_qv = h_v @@ -2355,26 +2219,26 @@ def host( ), ) - # ---- SMEM sizing: per-buffer element cosizes --------------------------------- + # ---- SMEM sizing: per-buffer element cosizes ------------------------------------- bpe = cfg.io_dtype.width // 8 - kq_tile_elems = 2 * cfg.b_t * cfg.d_k - v_tile_elems = cfg.d_v * cfg.b_t - tinv_tile_elems = cfg.b_t * cfg.b_t - a_tile_elems = cfg.b_t * cfg.b_t - o_tile_elems = cfg.d_v * cfg.b_t - cfg.kq_cosize = kq_tile_elems * cfg.smem_kq_stages - cfg.v_cosize = v_tile_elems * cfg.smem_v_stages - cfg.t_inv_cosize = tinv_tile_elems * cfg.smem_t_inv_stages - cfg.a_cosize = a_tile_elems * cfg.smem_a_stages - cfg.o_cosize = o_tile_elems * cfg.smem_o_stages + kq_tile_elements = 2 * cfg.b_t * cfg.d_k + v_tile_elements = cfg.d_v * cfg.b_t + tinv_tile_elements = cfg.b_t * cfg.b_t + a_tile_elements = cfg.b_t * cfg.b_t + o_tile_elements = cfg.d_v * cfg.b_t + cfg.kq_cosize = kq_tile_elements * cfg.smem_kq_stages + cfg.v_cosize = v_tile_elements * cfg.smem_v_stages + cfg.t_inv_cosize = tinv_tile_elements * cfg.smem_t_inv_stages + cfg.a_cosize = a_tile_elements * cfg.smem_a_stages + cfg.o_cosize = o_tile_elements * cfg.smem_o_stages cfg.checkpoint_cosize = cfg.d_k * cfg.d_v * cfg.smem_checkpoint_stages cumsumlog_smem_layout_staged = cute.make_layout((cfg.b_t, 1, cfg.smem_gate_stages)) beta_smem_layout_staged = cute.make_layout((cfg.b_t, 1, cfg.smem_beta_stages)) - cfg.tma_kq_bytes = kq_tile_elems * bpe - cfg.tma_v_bytes = v_tile_elems * bpe - cfg.tma_o_bytes = o_tile_elems * bpe + cfg.tma_kq_bytes = kq_tile_elements * bpe + cfg.tma_v_bytes = v_tile_elements * bpe + cfg.tma_o_bytes = o_tile_elements * bpe cfg.n_heads_out = heads_out cfg.q_ratio = heads_out // h_q @@ -2382,8 +2246,7 @@ def host( cfg.v_ratio = heads_out // h_v num_descs = batch_size - # ---- launch ------------------------------------------------------------------ - # CUDA-graph-stable launch: fixed SM-count grid; shapes ride on buffer contents + # ---- launch ---------------------------------------------------------------------- grid_shape = (cfg.max_active_clusters, 1, 1) kernel( @@ -2397,7 +2260,7 @@ def host( state_out, work_items, work_count, - sched_ctr, + scheduler_counter, checkpoint_every_n_tokens, scale, cumsumlog_smem_layout_staged, @@ -2429,7 +2292,7 @@ def kernel( mState_out: Optional[cute.Tensor], mWorkItems: cute.Tensor, mCount: cute.Tensor, - mSched: Optional[cute.Tensor], + mScheduler: Optional[cute.Tensor], checkpoint_every_n_tokens: cutlass.Int32, scale: cutlass.Float32, cumsumlog_smem_layout_staged: cute.Layout, @@ -2448,8 +2311,8 @@ def kernel( num_ctas = cute.arch.grid_dim()[0] total_tiles = mCount[0] - if cutlass.const_expr(cfg.dyn_sched): - assert mSched is not None, "mSched must be provided if dyn_sched is True" + if cutlass.const_expr(cfg.dynamic_scheduling): + assert mScheduler is not None, "mScheduler must be provided if dynamic_scheduling is True" if cutlass.const_expr(cfg.use_initial_state): assert mState_init is not None, "mState_init must be provided if use_initial_state is True" @@ -2471,7 +2334,6 @@ def kernel( SMEM = cutlass.AddressSpace.smem - bpe = cfg.io_dtype.width // 8 SWZ = 2 LEAD = 16 STRIDE = 8 * 128 @@ -2484,8 +2346,8 @@ def kernel( alignment=cfg.buffer_align_bytes, ) sO = SmemTile( - base=sO_raw.data_ptr().toint(), - elems_per_stage=(cfg.o_cosize // cfg.smem_o_stages) * bpe, + base=sO_raw.data_ptr(), + elems_per_stage=(cfg.o_cosize // cfg.smem_o_stages), stages=cfg.smem_o_stages, leading_byte_offset=LEAD, stride_byte_offset=STRIDE, @@ -2507,16 +2369,16 @@ def kernel( alignment=cfg.buffer_align_bytes, ) sKQ = SmemTile( - base=sKQ_raw.data_ptr().toint(), - elems_per_stage=(cfg.kq_cosize // cfg.smem_kq_stages) * bpe, + base=sKQ_raw.data_ptr(), + elems_per_stage=(cfg.kq_cosize // cfg.smem_kq_stages), stages=cfg.smem_kq_stages, leading_byte_offset=LEAD, stride_byte_offset=STRIDE, layout=SWZ, ) sKQ_trans = SmemTile( - base=sKQ_raw.data_ptr().toint(), - elems_per_stage=(cfg.kq_cosize // cfg.smem_kq_stages) * bpe, + base=sKQ_raw.data_ptr(), + elems_per_stage=(cfg.kq_cosize // cfg.smem_kq_stages), stages=cfg.smem_kq_stages, leading_byte_offset=2 * KT_LEAD, stride_byte_offset=STRIDE, @@ -2524,7 +2386,7 @@ def kernel( ) bars = make_gdn_bars(cfg) tmem_base_slot = cutlass.Array(cutlass.Int32, 1, space=SMEM, alignment=16) - sSched = cutlass.Array(cutlass.Int32, cfg.sched_stages, space=SMEM, alignment=16) + sScheduler = cutlass.Array(cutlass.Int32, cfg.scheduler_stages, space=SMEM, alignment=16) cumsumlog_raw = cutlass.Array(cutlass.Float32, cute.cosize(cumsumlog_smem_layout_staged), space=SMEM, alignment=128) cumprod_raw = cutlass.Array(cutlass.Float32, cute.cosize(cumsumlog_smem_layout_staged), space=SMEM, alignment=128) beta_raw = cutlass.Array(cutlass.Float32, cute.cosize(beta_smem_layout_staged), space=SMEM, alignment=128) @@ -2535,8 +2397,8 @@ def kernel( alignment=cfg.buffer_align_bytes, ) sTinv = SmemTile( - base=sTinv_raw.data_ptr().toint(), - elems_per_stage=(cfg.t_inv_cosize // cfg.smem_t_inv_stages) * bpe, + base=sTinv_raw.data_ptr(), + elems_per_stage=(cfg.t_inv_cosize // cfg.smem_t_inv_stages), stages=cfg.smem_t_inv_stages, leading_byte_offset=LEAD, stride_byte_offset=STRIDE, @@ -2549,8 +2411,8 @@ def kernel( alignment=cfg.buffer_align_bytes, ) sA = SmemTile( - base=sA_raw.data_ptr().toint(), - elems_per_stage=(cfg.a_cosize // cfg.smem_a_stages) * bpe, + base=sA_raw.data_ptr(), + elems_per_stage=(cfg.a_cosize // cfg.smem_a_stages), stages=cfg.smem_a_stages, leading_byte_offset=LEAD, stride_byte_offset=STRIDE, @@ -2562,9 +2424,9 @@ def kernel( space=cutlass.AddressSpace.smem, alignment=cfg.buffer_align_bytes, ) - sV = SmemTile( - base=sV_raw.data_ptr().toint(), - elems_per_stage=(cfg.v_cosize // cfg.smem_v_stages) * bpe, + sV_trans = SmemTile( + base=sV_raw.data_ptr(), + elems_per_stage=(cfg.v_cosize // cfg.smem_v_stages), stages=cfg.smem_v_stages, leading_byte_offset=V_LEAD, stride_byte_offset=STRIDE, @@ -2583,7 +2445,7 @@ def kernel( beta_smem_layout_staged, ) - # ---- mbarrier init (all threads) --------------------------------------------- + # ---- mbarrier init (all threads) ------------------------------------------------- for s in range(cfg.smem_kq_stages): bars.mb_kq_ready[s].init() bars.mb_kq_done[s].init() @@ -2598,7 +2460,6 @@ def kernel( bars.mb_beta_done[s].init() for s in range(cfg.tmem_state_acc_stages): bars.mb_state_acc_ready[s].init() - bars.mb_state_acc_scale_done[s].init() for s in range(cfg.tmem_q_state_acc_stages): bars.mb_o_acc_ready[s].init() bars.mb_o_final_acc_ready[s].init() @@ -2614,9 +2475,9 @@ def kernel( for s in range(cfg.smem_a_stages): bars.mb_a_ready[s].init() bars.mb_a_done[s].init() - for s in range(cfg.tmem_state_inp_stages): - bars.mb_state_inp_ready[s].init() - for b in (bars.mb_y_inp_ready, bars.mb_u_inp_ready, bars.mb_decay_u_inp_ready): + for s in range(cfg.tmem_state_input_stages): + bars.mb_state_input_ready[s].init() + for b in (bars.mb_y_input_ready, bars.mb_u_input_ready, bars.mb_decay_u_input_ready): b[0].init() for s in range(cfg.smem_o_stages): bars.mb_o_tmastg_ready[s].init() @@ -2624,15 +2485,15 @@ def kernel( for s in range(cfg.smem_checkpoint_stages): bars.mb_checkpoint_tmastg_ready[s].init() bars.mb_checkpoint_tmastg_done[s].init() - for s_ in range(cfg.sched_stages): - bars.mb_sched_ready[s_].init() - bars.mb_sched_done[s_].init() + for s in range(cfg.scheduler_stages): + bars.mb_scheduler_ready[s].init() + bars.mb_scheduler_done[s].init() bars.mb_tmem_done[0].init() nvvm.fence_mbarrier_init() nvvm.barrier_cta_sync() - # ---- warp specialization ----------------------------------------------------- + # ---- warp specialization --------------------------------------------------------- if warp_idx >= cfg.compute_group_0_warp_ids[0] and warp_idx <= cfg.compute_group_0_warp_ids[-1]: compute0_warp_group( @@ -2651,7 +2512,7 @@ def kernel( sA=sA, sCheckpoint_raw=sCheckpoint_raw, checkpoint_every_n_tokens=checkpoint_every_n_tokens, - sSched=sSched, + sScheduler=sScheduler, bars=bars, ) @@ -2667,7 +2528,7 @@ def kernel( warp_idx=warp_idx, tmem_base_slot=tmem_base_slot, scale=scale, - sV=sV, + sV_trans=sV_trans, sCumsumlog=sCumsumlog, sCumprod=sCumprod, sBeta=sBeta, @@ -2676,7 +2537,7 @@ def kernel( mState_init=mState_init, mState_out=mState_out, checkpoint_every_n_tokens=checkpoint_every_n_tokens, - sSched=sSched, + sScheduler=sScheduler, bars=bars, ) @@ -2696,12 +2557,12 @@ def kernel( sCumsumlog=sCumsumlog, sCumprod=sCumprod, sBeta=sBeta, - sSched=sSched, + sScheduler=sScheduler, bars=bars, ) - elif warp_idx == cfg.mma_warp_id: - mma_warp( + elif warp_idx == cfg.tcgen05_mma_warp_id: + tcgen05_mma_warp( cfg, total_tiles, bidx, @@ -2713,7 +2574,7 @@ def kernel( sKQ_trans=sKQ_trans, sTinv=sTinv, sA=sA, - sSched=sSched, + sScheduler=sScheduler, bars=bars, ) @@ -2730,8 +2591,8 @@ def kernel( desc_q_base=desc_q_base, desc_k_base=desc_k_base, desc_v_base=desc_v_base, - mSched=mSched, - sSched=sSched, + mScheduler=mScheduler, + sScheduler=sScheduler, bars=bars, ) @@ -2749,7 +2610,7 @@ def kernel( sCheckpoint_raw=sCheckpoint_raw, desc_o_base=desc_o_base, desc_checkpoint_base=desc_checkpoint_base, - sSched=sSched, + sScheduler=sScheduler, bars=bars, ) @@ -2776,10 +2637,10 @@ class GdnCfg: log_gate: bool = False safe_gate: bool = False beta_sigmoid: bool = False - dyn_sched: bool = False - sched_stages: int = CFG.SMEM_SCHED_STAGES + dynamic_scheduling: bool = False + scheduler_stages: int = CFG.SMEM_SCHEDULER_STAGES - # ---- fixed constants stamped from CFG by build_cfg --------------------------- + # ---- fixed constants stamped from CFG at build time ------------------------------ b_t: int = CFG.B_T d_k: int = CFG.D_K d_v: int = CFG.D_V @@ -2787,7 +2648,7 @@ class GdnCfg: compute_group_1_warp_ids: Tuple[int, ...] = CFG.COMPUTE_GROUP_1_WARP_IDS load_gate_beta_warp_id: int = CFG.LOAD_GATE_BETA_WARP_ID tma_qkv_warp_id: int = CFG.TMA_QKV_WARP_ID - mma_warp_id: int = CFG.MMA_WARP_ID + tcgen05_mma_warp_id: int = CFG.TCGEN05_MMA_WARP_ID epilogue_warp_id: int = CFG.EPILOGUE_WARP_ID num_regs_compute_group_0: int = CFG.NUM_REGS_COMPUTE_GROUP_0 num_regs_compute_group_1: int = CFG.NUM_REGS_COMPUTE_GROUP_1 @@ -2796,15 +2657,13 @@ class GdnCfg: threads_per_cta: int = 0 cluster_shape_mnk: Tuple[int, int, int] = CFG.CLUSTER_SHAPE_MNK - # ---- named barrier slots (ids 1-4; 0 is the CTA-wide sync) ------------------- - tmem_alloc_barrier_id: int = 1 - tmem_alloc_barrier_threads: int = 0 + # ---- named barrier slots (ids 1-4; 0 is the CTA-wide sync) ----------------------- + tmem_lifecycle_barrier_id: int = 1 + tmem_user_threads: int = 0 inverse_barrier_id: int = 2 inverse_barrier_threads: int = 0 - init_state_store_barrier_id: int = 4 - init_state_store_barrier_threads: int = 0 - # ---- SMEM / TMEM stage counts + TMEM column offsets -------------------------- + # ---- SMEM / TMEM stage counts + TMEM column offsets ------------------------------ smem_kq_stages: int = CFG.SMEM_KQ_STAGES smem_v_stages: int = CFG.SMEM_V_STAGES smem_t_inv_stages: int = CFG.SMEM_T_INV_STAGES @@ -2815,18 +2674,18 @@ class GdnCfg: smem_beta_stages: int = CFG.SMEM_BETA_STAGES tmem_state_acc_stages: int = CFG.TMEM_KV_ACC_STAGES tmem_q_state_acc_stages: int = CFG.TMEM_Q_STATE_ACC_STAGES - tmem_state_inp_stages: int = CFG.TMEM_STATE_INP_STAGES + tmem_state_input_stages: int = CFG.TMEM_STATE_INP_STAGES tmem_cg0_acc_stages: int = CFG.TMEM_CG0_ACC_STAGES tmem_cg1_acc_stages: int = CFG.TMEM_CG1_ACC_STAGES tmem_state_acc_offset: int = 0 tmem_q_state_acc_offset: int = 0 - tmem_state_inp_offset: int = 0 + tmem_state_input_offset: int = 0 tmem_cg0_acc_offset: int = 0 tmem_cg1_acc_offset: int = 0 - tmem_y_decay_u_inp_offset: int = 0 + tmem_y_decay_u_input_offset: int = 0 buffer_align_bytes: int = CFG.BUFFER_ALIGN_BYTES - # ---- stamped by host at trace time (shape-derived) -------------------------- + # ---- stamped by host at trace time (shape-derived) ------------------------------- kq_cosize: int = 0 v_cosize: int = 0 t_inv_cosize: int = 0 @@ -2854,7 +2713,7 @@ def build_cfg( log_gate: bool = False, safe_gate: bool = False, beta_sigmoid: bool = False, - dyn_sched: bool = False, + dynamic_scheduling: bool = False, ) -> GdnCfg: """Build the per-compile ``GdnCfg`` (io_dtype ∈ {Float16, BFloat16}; acc is always Float32).""" @@ -2872,7 +2731,7 @@ def build_cfg( log_gate=log_gate, safe_gate=safe_gate, beta_sigmoid=beta_sigmoid, - dyn_sched=dyn_sched, + dynamic_scheduling=dynamic_scheduling, ) cfg.smem_checkpoint_stages = 1 if enable_checkpoints: @@ -2883,15 +2742,14 @@ def build_cfg( n_cg0 = len(cfg.compute_group_0_warp_ids) n_cg1 = len(cfg.compute_group_1_warp_ids) cfg.threads_per_cta = cfg.threads_per_warp * (4 + n_cg0 + n_cg1) - cfg.tmem_alloc_barrier_threads = cfg.threads_per_warp * (1 + n_cg0 + n_cg1) + cfg.tmem_user_threads = cfg.threads_per_warp * (1 + n_cg0 + n_cg1) cfg.inverse_barrier_threads = cfg.threads_per_warp * n_cg0 - cfg.init_state_store_barrier_threads = cfg.threads_per_warp * n_cg1 cfg.tmem_state_acc_offset = 0 cfg.tmem_q_state_acc_offset = cfg.tmem_state_acc_offset + cfg.tmem_state_acc_stages * 128 - cfg.tmem_state_inp_offset = cfg.tmem_q_state_acc_offset + cfg.tmem_q_state_acc_stages * 64 - cfg.tmem_cg0_acc_offset = cfg.tmem_state_inp_offset + cfg.tmem_state_inp_stages * 64 + cfg.tmem_state_input_offset = cfg.tmem_q_state_acc_offset + cfg.tmem_q_state_acc_stages * 64 + cfg.tmem_cg0_acc_offset = cfg.tmem_state_input_offset + cfg.tmem_state_input_stages * 64 cfg.tmem_cg1_acc_offset = cfg.tmem_cg0_acc_offset + cfg.tmem_cg0_acc_stages * 64 - cfg.tmem_y_decay_u_inp_offset = cfg.tmem_cg1_acc_offset + cfg.tmem_cg1_acc_stages * 64 + cfg.tmem_y_decay_u_input_offset = cfg.tmem_cg1_acc_offset + cfg.tmem_cg1_acc_stages * 64 return cfg @@ -2917,7 +2775,7 @@ def get_compiled_cache( log_gate: bool, safe_gate: bool, beta_sigmoid: bool, - dyn_sched: bool, + dynamic_scheduling: bool, order_gen: bool, ): """Return a mutable dict that lazily stores the compiled kernel.""" @@ -2934,9 +2792,12 @@ def compile( log_gate: bool = False, safe_gate: bool = False, beta_sigmoid: bool = False, - dyn_sched: bool = False, + dynamic_scheduling: bool = False, *, num_sm: int, + h_q: int, + h_k: int, + h_v: int, q_cute, k_cute, v_cute, @@ -2950,7 +2811,7 @@ def compile( state_out_cute, work_items_cute=None, work_count_cute=None, - sched_ctr_cute=None, + scheduler_counter_cute=None, checkpoint_every_n_tokens, scale, workspace_cute, @@ -2968,8 +2829,11 @@ def compile( log_gate=log_gate, safe_gate=safe_gate, beta_sigmoid=beta_sigmoid, - dyn_sched=dyn_sched, + dynamic_scheduling=dynamic_scheduling, ) + cfg.h_q = h_q + cfg.h_k = h_k + cfg.h_v = h_v return cute.compile( host, @@ -2987,7 +2851,7 @@ def compile( state_out_cute, work_items_cute, work_count_cute, - sched_ctr_cute, + scheduler_counter_cute, checkpoint_every_n_tokens, scale, workspace_cute, @@ -3011,7 +2875,7 @@ def chunk_gdn_sm100( output_state_checkpoints=None, work_items=None, work_count=None, - sched_ctr=None, + scheduler_counter=None, log_gate: bool = False, safe_gate: bool = False, a_log=None, @@ -3024,8 +2888,8 @@ def chunk_gdn_sm100( ) -> None: """Execute the Blackwell chunked GDN prefill kernel (THD / varlen entry). - All tensors are contiguous, DLPack-compatible CUDA tensors on the same - device. Compile-cache-and-replay: the kernel is compiled once per static + All tensors are DLPack-compatible CUDA tensors on the same device with a + stride-1 innermost dim (outer strides are runtime arguments). Compile-cache-and-replay: the kernel is compiled once per static config (dtypes, head counts, state flags) and replayed afterwards. Args: @@ -3051,8 +2915,8 @@ def chunk_gdn_sm100( series the backward pass consumes. work_items: ``(max_items, 8)`` int32 work-item table from ``common/split_k.py`` (REQUIRED; an uncut table row is the whole - (b, h) sequence). Each item computes chunks ``[cstart, wend)`` - and writes O/checkpoints only for ``[wstart, wend)``. + (b, h) sequence). Each item computes chunks ``[compute_start, write_end)`` + and writes O/checkpoints only for ``[write_start, write_end)``. work_count: ``(1,)`` int32 device-side item count (REQUIRED) log_gate: ``gate`` holds natural-log decay values; the gate warp skips its log2 (rescales by 1/ln2) instead of exponentiating @@ -3075,13 +2939,15 @@ def chunk_gdn_sm100( store_final_state = output_state is not None enable_checkpoints = checkpoint_every_n_tokens > 0 if work_items is None or work_count is None: - raise ValueError("work_items/work_count are required (the split-table stage builds them for every launch)") + raise ValueError( + "work_items/work_count are required (built by the split-table stage, or by an order-generating prologue when work_item_scratch is None)" + ) if safe_gate and (a_log is None or dt_bias is None): raise ValueError("safe_gate requires a_log and dt_bias") if not safe_gate: a_log = None dt_bias = None - dyn_sched = sched_ctr is not None + dynamic_scheduling = scheduler_counter is not None order_gen = work_item_scratch is None io_dtype = get_dtype(q.dtype) @@ -3109,25 +2975,19 @@ def chunk_gdn_sm100( log_gate, safe_gate, use_beta_sigmoid, - dyn_sched, + dynamic_scheduling, order_gen, ) if "compiled" not in cache: - q_cute = from_dlpack(q, assumed_align=16) - q_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) - k_cute = from_dlpack(k, assumed_align=16) - k_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) - v_cute = from_dlpack(v, assumed_align=16) - v_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) - gate_cute = from_dlpack(gate, assumed_align=16) - gate_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1), divisibility=1) + q_cute = from_dlpack(q, assumed_align=16).mark_layout_dynamic(leading_dim=2) + k_cute = from_dlpack(k, assumed_align=16).mark_layout_dynamic(leading_dim=2) + v_cute = from_dlpack(v, assumed_align=16).mark_layout_dynamic(leading_dim=2) + gate_cute = from_dlpack(gate, assumed_align=16).mark_layout_dynamic(leading_dim=1) a_log_cute = from_dlpack(a_log, assumed_align=4) if a_log is not None else None dt_bias_cute = from_dlpack(dt_bias, assumed_align=4) if dt_bias is not None else None - beta_cute = from_dlpack(beta, assumed_align=16) - beta_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1), divisibility=1) - o_cute = from_dlpack(output, assumed_align=16) - o_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) + beta_cute = from_dlpack(beta, assumed_align=16).mark_layout_dynamic(leading_dim=1) + o_cute = from_dlpack(output, assumed_align=16).mark_layout_dynamic(leading_dim=2) cu_seqlens_cute = from_dlpack(cu_seqlens, assumed_align=8 if str(cu_seqlens.dtype).endswith("int64") else 4).mark_layout_dynamic() state_in_cute = None @@ -3146,9 +3006,9 @@ def chunk_gdn_sm100( work_items_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1), divisibility=1) work_count_cute = from_dlpack(work_count, assumed_align=4).mark_layout_dynamic() - sched_ctr_cute = None - if dyn_sched: - sched_ctr_cute = from_dlpack(sched_ctr, assumed_align=4).mark_layout_dynamic() + scheduler_counter_cute = None + if dynamic_scheduling: + scheduler_counter_cute = from_dlpack(scheduler_counter, assumed_align=4).mark_layout_dynamic() cache["compiled"] = compile( io_dtype, @@ -3160,8 +3020,11 @@ def chunk_gdn_sm100( log_gate, safe_gate, use_beta_sigmoid, - dyn_sched, + dynamic_scheduling, num_sm=multiprocessor_count(current_device()), + h_q=HQ, + h_k=k.shape[1], + h_v=HV, q_cute=q_cute, k_cute=k_cute, v_cute=v_cute, @@ -3175,7 +3038,7 @@ def chunk_gdn_sm100( state_out_cute=state_out_cute, work_items_cute=work_items_cute, work_count_cute=work_count_cute, - sched_ctr_cute=sched_ctr_cute, + scheduler_counter_cute=scheduler_counter_cute, checkpoint_every_n_tokens=checkpoint_every_n_tokens, scale=scale, workspace_cute=workspace_cute, @@ -3185,19 +3048,14 @@ def chunk_gdn_sm100( compiled = cache["compiled"] if "prologue" not in cache: - q_pl = from_dlpack(q, assumed_align=16) - q_pl.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) - k_pl = from_dlpack(k, assumed_align=16) - k_pl.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) - v_pl = from_dlpack(v, assumed_align=16) - v_pl.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) - o_pl = from_dlpack(output, assumed_align=16) - o_pl.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) + q_pl = from_dlpack(q, assumed_align=16).mark_layout_dynamic(leading_dim=2) + k_pl = from_dlpack(k, assumed_align=16).mark_layout_dynamic(leading_dim=2) + v_pl = from_dlpack(v, assumed_align=16).mark_layout_dynamic(leading_dim=2) + o_pl = from_dlpack(output, assumed_align=16).mark_layout_dynamic(leading_dim=2) cu_pl = from_dlpack(cu_seqlens, assumed_align=8 if str(cu_seqlens.dtype).endswith("int64") else 4).mark_layout_dynamic() checkpoints_pl = None if enable_checkpoints: - checkpoints_pl = from_dlpack(output_state_checkpoints, assumed_align=16) - checkpoints_pl.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2, 3), divisibility=1) + checkpoints_pl = from_dlpack(output_state_checkpoints, assumed_align=16).mark_layout_dynamic(leading_dim=3) staging_pl = None if not order_gen: staging_pl = from_dlpack(work_item_scratch, assumed_align=16) @@ -3205,16 +3063,16 @@ def chunk_gdn_sm100( work_items_pl = from_dlpack(work_items, assumed_align=16) work_items_pl.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1), divisibility=1) work_count_pl = from_dlpack(work_count, assumed_align=4).mark_layout_dynamic() - sched_pl = None - if dyn_sched: - sched_pl = from_dlpack(sched_ctr, assumed_align=4).mark_layout_dynamic() - ws_pl = from_dlpack(workspace, assumed_align=128).mark_layout_dynamic() + scheduler_pl = None + if dynamic_scheduling: + scheduler_pl = from_dlpack(scheduler_counter, assumed_align=4).mark_layout_dynamic() + workspace_pl = from_dlpack(workspace, assumed_align=128).mark_layout_dynamic() cache["prologue"] = cute.compile( prologue, io_dtype, CFG.B_T, order_gen, - dyn_sched, + dynamic_scheduling, q_pl, k_pl, v_pl, @@ -3224,9 +3082,9 @@ def chunk_gdn_sm100( staging_pl, work_count_pl, work_items_pl, - sched_pl, + scheduler_pl, cutlass.Int32(checkpoint_every_n_tokens), - ws_pl, + workspace_pl, cu_stream, options="--enable-tvm-ffi", ) @@ -3240,7 +3098,7 @@ def chunk_gdn_sm100( work_item_scratch if not order_gen else None, work_count, work_items, - sched_ctr, + scheduler_counter, checkpoint_every_n_tokens, workspace, cu_stream, @@ -3259,7 +3117,7 @@ def chunk_gdn_sm100( output_state, work_items, work_count, - sched_ctr, + scheduler_counter, checkpoint_every_n_tokens, scale, workspace, @@ -3282,7 +3140,7 @@ def run_prefill( output_state_checkpoints, work_items, work_count, - sched_ctr, + scheduler_counter, work_item_scratch, tensormap_workspace, checkpoint_every_n_tokens, @@ -3305,7 +3163,7 @@ def run_prefill( work_item_scratch, work_count, work_items, - sched_ctr, + scheduler_counter, checkpoint_every_n_tokens, tensormap_workspace, cu_stream, @@ -3324,7 +3182,7 @@ def run_prefill( output_state, work_items, work_count, - sched_ctr, + scheduler_counter, checkpoint_every_n_tokens, scale, tensormap_workspace, diff --git a/python/cudnn/linear_attention/frost/kernel/gdn_recompute_config.py b/python/cudnn/linear_attention/frost/kernel/gdn_recompute_config.py index 30675f5a2..8c59ca589 100644 --- a/python/cudnn/linear_attention/frost/kernel/gdn_recompute_config.py +++ b/python/cudnn/linear_attention/frost/kernel/gdn_recompute_config.py @@ -39,7 +39,7 @@ class Cfg: COMPUTE_GROUP_1_WARP_IDS: Tuple[int, ...] = (4, 5, 6, 7) # kv_decay_v / v-k*state / epi ops LOAD_GATE_BETA_WARP_ID: int = 8 # gate/beta chunk loads + TMEM lifecycle TMA_KV_WARP_ID: int = 9 - MMA_WARP_ID: int = 10 # sole tcgen05 issuer: KK pairs + KS/U/KV per chunk + TCGEN05_MMA_WARP_ID: int = 10 # sole tcgen05 issuer: KK pairs + KS/U/KV per chunk EPILOGUE_WARP_ID: int = 11 # --- register split --- @@ -52,7 +52,7 @@ class Cfg: CLUSTER_SHAPE_MNK: Tuple[int, int, int] = (1, 1, 1) # --- SMEM stage counts --- - SMEM_SCHED_STAGES: int = 2 + SMEM_SCHEDULER_STAGES: int = 2 SMEM_KQ_STAGES: int = 4 SMEM_V_STAGES: int = 2 SMEM_T_INV_STAGES: int = 3 diff --git a/python/cudnn/linear_attention/frost/kernel/gdn_recompute_f16.py b/python/cudnn/linear_attention/frost/kernel/gdn_recompute_f16.py index eb71a4b38..5138cefac 100644 --- a/python/cudnn/linear_attention/frost/kernel/gdn_recompute_f16.py +++ b/python/cudnn/linear_attention/frost/kernel/gdn_recompute_f16.py @@ -50,18 +50,18 @@ T_inv 8192 2 checkpoint staging DK*DV*2 1 <-- enable_checkpoints only cumsumlog / cumprod / Beta 256 3 - sched ticket ring 4 2 <-- dyn_sched publish ring + scheduler ticket ring 4 2 <-- dynamic_scheduling publish ring TMEM layout (512 columns): Buffer Cols state 128 <-- DKxDV fp32 = 128x128x4B - state inp 64 <-- fp16 state staging (GEMM 3 A operand) + state input 64 <-- fp16 state staging (GEMM 3 A operand) cg0 shared acc 128 <-- 2-stage ring: KK0/KK1 cg1 shared acc 64 <-- 1-stage ring: KS then U - Y / decayed-U inp 64 <-- slot 0 = Y (V - K*state), slot 1 = decayed U (b16) + Y / decayed-U input 64 <-- slot 0 = Y (V - K*state), slot 1 = decayed U (b16) Warp assignments (12 warps = 384 threads): - warps 0-3 : compute group 0 - T-pairwise x2, KK_epi x2, pair inverse + warps 0-3 : compute group 0 - T-pairwise x2, KK epilogue x2, pair inverse warps 4-7 : compute group 1 - state restage/rescale, Y = V - K*state, U epilogue warp 8 : Gate/Beta loads @@ -85,8 +85,7 @@ from cutlass.cutlass_dsl import min from ..common.thd import emit_checkpoint_seq_descs, emit_seq_descs, TENSOR_MAP_QWORDS -from ..common.split_k import ORDER_CAPACITY, ORDER_ELEMS, ORDER_THREADS, decode_work_item, order_body -from ..common.elementwise import softplus +from ..common.split_k import ORDER_CAPACITY, ORDER_ELEMENTS, ORDER_THREADS, decode_work_item, order_body from ..common.host import get_dtype from cudnn.frost.buffers import data_ptr from cudnn.frost.device import current_device, multiprocessor_count @@ -100,8 +99,8 @@ ) from cudnn.frost.tile_dsl.handles import MmaDesc, SmemTile, tma_slice_runtime_desc from cudnn.frost.tile_dsl.mma import mma_ss, mma_step_k8, mma_ts_step, mma_step -from cudnn.frost.tile_dsl.pointwise import fadd2, fp32_to_fp16, f16x2_to_f32, fmul2, opaque_f32_zero, sub_f16x2 -from cudnn.frost.tile_dsl.swizzle import swizzle_lin_128b, swizzle_xor_128b +from cudnn.frost.tile_dsl.pointwise import f16x2_to_f32, fadd2, fmul2, fp32_to_fp16, opaque_f32_zero, sigmoid, softplus, softplus2, sub_f16x2 +from cudnn.frost.tile_dsl.swizzle import swizzle_xor_128b from cudnn.frost.tile_dsl.tma import ( tma_load_tile, tma_store_tile, @@ -113,13 +112,7 @@ class GdnBars(NamedTuple): - """GDN pipeline mbarrier inventory. - - Every pipeline is a ``_ready``/``_done`` MBarrier pair over one ring: a - slot is acquired for filling by waiting ``_done`` and committed by - arriving ``_ready``; the reading side waits ``_ready`` and releases the - slot by arriving ``_done``. - """ + """Every inter-warp handoff as an ``MBarrier`` over its ring.""" mb_kq_ready: MBarrier mb_kq_done: MBarrier @@ -132,7 +125,6 @@ class GdnBars(NamedTuple): mb_beta_done: MBarrier mb_state_acc_ready: MBarrier - mb_state_acc_scale_done: MBarrier mb_cg0_acc_ready: MBarrier mb_cg0_acc_done: MBarrier mb_k_state_acc_ready: MBarrier @@ -140,22 +132,22 @@ class GdnBars(NamedTuple): mb_t_inv_ready: MBarrier mb_t_inv_done: MBarrier - mb_state_inp_ready: MBarrier - mb_y_inp_ready: MBarrier - mb_decay_u_inp_ready: MBarrier + mb_state_input_ready: MBarrier + mb_y_input_ready: MBarrier + mb_decay_u_input_ready: MBarrier mb_checkpoint_tmastg_ready: MBarrier mb_checkpoint_tmastg_done: MBarrier mb_tmem_done: MBarrier - mb_sched_ready: MBarrier - mb_sched_done: MBarrier + mb_scheduler_ready: MBarrier + mb_scheduler_done: MBarrier def make_gdn_bars(cfg) -> GdnBars: - """GdnBars factory. MUST be called from inside ``kernel`` (allocates SMEM).""" + """GdnBars factory.""" ONE_LANE = 1 - MMA_ARRIVERS = len([cfg.mma_warp_id]) + MMA_ARRIVERS = len([cfg.tcgen05_mma_warp_id]) KQ_RELEASE_SITES = 1 GATE_WARP = cfg.threads_per_warp * len([cfg.load_gate_beta_warp_id]) EPI_WARP = cfg.threads_per_warp * len([cfg.epilogue_warp_id]) @@ -176,43 +168,33 @@ def alloc(n): mb_beta_ready=MBarrier(alloc(cfg.smem_beta_stages), stages=cfg.smem_beta_stages, init_count=GATE_WARP, producer=Producer.THREAD), mb_beta_done=MBarrier(alloc(cfg.smem_beta_stages), stages=cfg.smem_beta_stages, init_count=CG0_THREADS, producer=Producer.THREAD), mb_state_acc_ready=MBarrier(alloc(cfg.tmem_state_acc_stages), stages=cfg.tmem_state_acc_stages, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), - mb_state_acc_scale_done=MBarrier(alloc(cfg.tmem_state_acc_stages), stages=cfg.tmem_state_acc_stages, init_count=CG1_THREADS, producer=Producer.THREAD), mb_cg0_acc_ready=MBarrier(alloc(cfg.tmem_cg0_acc_stages), stages=cfg.tmem_cg0_acc_stages, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), mb_cg0_acc_done=MBarrier(alloc(cfg.tmem_cg0_acc_stages), stages=cfg.tmem_cg0_acc_stages, init_count=CG0_THREADS // 2, producer=Producer.THREAD), mb_k_state_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), mb_u_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), mb_t_inv_ready=MBarrier(alloc(cfg.smem_t_inv_stages), stages=cfg.smem_t_inv_stages, init_count=CG0_THREADS, producer=Producer.THREAD), mb_t_inv_done=MBarrier(alloc(cfg.smem_t_inv_stages), stages=cfg.smem_t_inv_stages, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), - mb_state_inp_ready=MBarrier(alloc(cfg.tmem_state_inp_stages), stages=cfg.tmem_state_inp_stages, init_count=CG1_THREADS, producer=Producer.THREAD), - mb_y_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), - mb_decay_u_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_state_input_ready=MBarrier(alloc(cfg.tmem_state_input_stages), stages=cfg.tmem_state_input_stages, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_y_input_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_decay_u_input_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), mb_checkpoint_tmastg_ready=MBarrier( alloc(cfg.smem_checkpoint_stages), stages=cfg.smem_checkpoint_stages, init_count=len(cfg.compute_group_1_warp_ids), producer=Producer.THREAD ), mb_checkpoint_tmastg_done=MBarrier(alloc(cfg.smem_checkpoint_stages), stages=cfg.smem_checkpoint_stages, init_count=EPI_WARP, producer=Producer.THREAD), mb_tmem_done=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), - mb_sched_ready=MBarrier(alloc(cfg.sched_stages), stages=cfg.sched_stages, init_count=1, producer=Producer.THREAD), - mb_sched_done=MBarrier(alloc(cfg.sched_stages), stages=cfg.sched_stages, init_count=11, producer=Producer.THREAD), + mb_scheduler_ready=MBarrier(alloc(cfg.scheduler_stages), stages=cfg.scheduler_stages, init_count=1, producer=Producer.THREAD), + mb_scheduler_done=MBarrier(alloc(cfg.scheduler_stages), stages=cfg.scheduler_stages, init_count=11, producer=Producer.THREAD), ) @cute.jit -def invert_diagonal_NxN(cfg, base_int, d, tidx, N: int = 8): +def invert_diagonal_NxN(cfg, base, d_idx, tidx, N: int = 8): """Gauss-Jordan inversion of one diagonal NxN block in-place (f16 SMEM).""" tidx_in_group = tidx % N BT = cfg.b_t - row_lin_base = (d * N + tidx_in_group) * BT + d * N - row_phys = swizzle_lin_128b(row_lin_base, row_stride_log2=6) - row_ptr = ( - cute.make_ptr( - cfg.io_dtype, - base_int, - mem_space=cute.AddressSpace.smem, - assumed_align=cfg.buffer_align_bytes, - ) - + row_phys - ) + row_coord = d_idx * N + tidx_in_group + row_ptr = base + row_coord * BT + swizzle_xor_128b(row_coord, d_idx * N) row = [(row_ptr + j).load().to(cutlass.Float32) for j in range(N)] for i in cutlass.range_constexpr(N): @@ -229,213 +211,146 @@ def invert_diagonal_NxN(cfg, base_int, d, tidx, N: int = 8): @cute.jit -def blockwise_diagonal_8x8_to_16x16(cfg, base_int, d0, lane_id): +def blockwise_diagonal_8x8_to_16x16(cfg, base, d_idx, lane_idx): """Off-diagonal correction 8x8 -> 16x16 (C <- -D^{-1} C A^{-1}).""" - bpe = cfg.io_dtype.width // 8 - lds1 = (lane_id % 8) * 64 - d = nvvm.ldmatrix( - cutlass.inttoptr(base_int + swizzle_lin_128b((d0 + 8) * 64 + d0 + 8 + lds1, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16), - 1, - nvvm.MMALayout.ROW, - ) - c = nvvm.ldmatrix( - cutlass.inttoptr(base_int + swizzle_lin_128b((d0 + 8) * 64 + d0 + lds1, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16), - 1, - nvvm.MMALayout.COL, - ) - - # ---- T = -(D^{-1} @ C) ------------------------------------------------------- + BT = cfg.b_t + row_lo = d_idx + lane_idx % 8 + row_hi = row_lo + 8 + off_d_inv = row_hi * BT + swizzle_xor_128b(row_hi, d_idx + 8) + off_c = row_hi * BT + swizzle_xor_128b(row_hi, d_idx) + off_a_inv = row_lo * BT + swizzle_xor_128b(row_lo, d_idx) + d_inv_frag = nvvm.ldmatrix(base + off_d_inv, 1, nvvm.MMALayout.ROW) + c_frag = nvvm.ldmatrix(base + off_c, 1, nvvm.MMALayout.COL) + + # ---- T = -(D^-1 @ C) ------------------------------------------------------------- c_regs = cutlass.Array(cutlass.Float32, 4, alignment=16, space=cutlass.AddressSpace.rmem) for i in cutlass.range_constexpr(4): c_regs[i] = cutlass.Float32(0.0) - mma_step_k8(c_regs, [d, d], [c], k_step=0, M=16, N=8, ab_dtype=cfg.io_dtype) + mma_step_k8(c_regs, [d_inv_frag, d_inv_frag], [c_frag], k_step=0, M=16, N=8, ab_dtype=cfg.io_dtype) for i in cutlass.range_constexpr(4): c_regs[i] = -c_regs[i] a_pack = [fp32_to_fp16(c_regs[2 * j], c_regs[2 * j + 1], dtype=cfg.io_dtype) for j in range(2)] - # ---- C = T @ A^{-1} ---------------------------------------------------------- - ai = nvvm.ldmatrix( - cutlass.inttoptr(base_int + swizzle_lin_128b(d0 * 64 + d0 + lds1, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16), - 1, - nvvm.MMALayout.COL, - ) + # ---- C = T @ A^-1 ---------------------------------------------------------------- + a_inv_frag = nvvm.ldmatrix(base + off_a_inv, 1, nvvm.MMALayout.COL) o_regs = cutlass.Array(cutlass.Float32, 4, alignment=16, space=cutlass.AddressSpace.rmem) for i in cutlass.range_constexpr(4): o_regs[i] = cutlass.Float32(0.0) - mma_step_k8(o_regs, a_pack, [ai], k_step=0, M=16, N=8, ab_dtype=cfg.io_dtype) + mma_step_k8(o_regs, a_pack, [a_inv_frag], k_step=0, M=16, N=8, ab_dtype=cfg.io_dtype) o_pack = fp32_to_fp16(o_regs[0], o_regs[1], dtype=cfg.io_dtype) - # ---- store corrected C ------------------------------------------------------- - nvvm.stmatrix( - cutlass.inttoptr(base_int + swizzle_lin_128b((d0 + 8) * 64 + d0 + lds1, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16), - o_pack, - nvvm.MMALayout.ROW, - ) + # ---- store corrected C ----------------------------------------------------------- + nvvm.stmatrix(base + off_c, o_pack, nvvm.MMALayout.ROW) @cute.jit -def blockwise_diagonal_16x16_to_32x32(cfg, base_int, d0, lane_id): +def blockwise_diagonal_16x16_to_32x32(cfg, base, d_idx, lane_idx): """Off-diagonal correction 16x16 -> 32x32.""" - bpe = cfg.io_dtype.width // 8 - lds4 = (lane_id % 16) * 64 + (lane_id // 16) * 8 - d = list( - nvvm.ldmatrix( - cutlass.inttoptr( - base_int + swizzle_lin_128b((d0 + 16) * 64 + d0 + 16 + lds4, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16 - ), - 4, - nvvm.MMALayout.ROW, - ) - ) - c = list( - nvvm.ldmatrix( - cutlass.inttoptr(base_int + swizzle_lin_128b((d0 + 16) * 64 + d0 + lds4, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16), - 4, - nvvm.MMALayout.COL, - ) - ) - - # ---- T = -(D^{-1} @ C) ------------------------------------------------------- + BT = cfg.b_t + lane_row = lane_idx % 16 + lane_col = (lane_idx // 16) * 8 + row_lo = d_idx + lane_row + row_hi = row_lo + 16 + off_d_inv = row_hi * BT + swizzle_xor_128b(row_hi, d_idx + 16 + lane_col) + off_c = row_hi * BT + swizzle_xor_128b(row_hi, d_idx + lane_col) + off_a_inv = row_lo * BT + swizzle_xor_128b(row_lo, d_idx + lane_col) + d_inv_frags = list(nvvm.ldmatrix(base + off_d_inv, 4, nvvm.MMALayout.ROW)) + c_frags = list(nvvm.ldmatrix(base + off_c, 4, nvvm.MMALayout.COL)) + + # ---- T = -(D^-1 @ C) ------------------------------------------------------------- c_regs = cutlass.Array(cutlass.Float32, 8, alignment=16, space=cutlass.AddressSpace.rmem) for i in cutlass.range_constexpr(8): c_regs[i] = cutlass.Float32(0.0) - mma_step(c_regs, d, c, k_step=0, M=16, N=16, ab_dtype=cfg.io_dtype) + mma_step(c_regs, d_inv_frags, c_frags, k_step=0, M=16, N=16, ab_dtype=cfg.io_dtype) for i in cutlass.range_constexpr(8): c_regs[i] = -c_regs[i] a_pack = [fp32_to_fp16(c_regs[2 * j], c_regs[2 * j + 1], dtype=cfg.io_dtype) for j in range(4)] - # ---- C = T @ A^{-1} ---------------------------------------------------------- - ai = list( - nvvm.ldmatrix( - cutlass.inttoptr(base_int + swizzle_lin_128b(d0 * 64 + d0 + lds4, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16), - 4, - nvvm.MMALayout.COL, - ) - ) + # ---- C = T @ A^-1 ---------------------------------------------------------------- + a_inv_frags = list(nvvm.ldmatrix(base + off_a_inv, 4, nvvm.MMALayout.COL)) o_regs = cutlass.Array(cutlass.Float32, 8, alignment=16, space=cutlass.AddressSpace.rmem) for i in cutlass.range_constexpr(8): o_regs[i] = cutlass.Float32(0.0) - mma_step(o_regs, a_pack, ai, k_step=0, M=16, N=16, ab_dtype=cfg.io_dtype) + mma_step(o_regs, a_pack, a_inv_frags, k_step=0, M=16, N=16, ab_dtype=cfg.io_dtype) o_pack = [fp32_to_fp16(o_regs[2 * j], o_regs[2 * j + 1], dtype=cfg.io_dtype) for j in range(4)] - # ---- store corrected C ------------------------------------------------------- - nvvm.stmatrix( - cutlass.inttoptr(base_int + swizzle_lin_128b((d0 + 16) * 64 + d0 + lds4, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16), - o_pack, - nvvm.MMALayout.ROW, - ) + # ---- store corrected C ----------------------------------------------------------- + nvvm.stmatrix(base + off_c, o_pack, nvvm.MMALayout.ROW) @cute.jit -def blockwise_diagonal_32x32_to_64x64(cfg, base_int, warp_id, lane_id): +def blockwise_diagonal_32x32_to_64x64(cfg, base, warp_id, lane_idx): """Off-diagonal correction 32x32 -> 64x64 (2 warps, one 16-row M-band each).""" band = warp_id % 2 - bpe = cfg.io_dtype.width // 8 - lds4 = (lane_id % 16) * 64 + (lane_id // 16) * 8 - a_frags = [] + BT = cfg.b_t + lane_row = lane_idx % 16 + lane_col = (lane_idx // 16) * 8 + row_d_inv = 32 + band * 16 + lane_row + d_inv_frags = [] for vs in cutlass.range_constexpr(2): - a_frags += list( - nvvm.ldmatrix( - cutlass.inttoptr( - base_int + swizzle_lin_128b((32 + band * 16) * 64 + 32 + vs * 16 + lds4, row_stride_log2=6) * bpe, - cutlass.AddressSpace.smem, - cutlass.BFloat16, - ), - 4, - nvvm.MMALayout.ROW, - ) - ) - b_frags = [] + d_inv_frags += list(nvvm.ldmatrix(base + row_d_inv * BT + swizzle_xor_128b(row_d_inv, 32 + vs * 16 + lane_col), 4, nvvm.MMALayout.ROW)) + c_frags = [] for vs in cutlass.range_constexpr(4): - b_frags += list( - nvvm.ldmatrix( - cutlass.inttoptr( - base_int + swizzle_lin_128b((32 + (vs // 2) * 16) * 64 + (vs % 2) * 16 + lds4, row_stride_log2=6) * bpe, - cutlass.AddressSpace.smem, - cutlass.BFloat16, - ), - 4, - nvvm.MMALayout.COL, - ) - ) + row_c = 32 + (vs // 2) * 16 + lane_row + c_frags += list(nvvm.ldmatrix(base + row_c * BT + swizzle_xor_128b(row_c, (vs % 2) * 16 + lane_col), 4, nvvm.MMALayout.COL)) - # ---- T = -(D^{-1} @ C) ------------------------------------------------------- + # ---- T = -(D^-1 @ C) ------------------------------------------------------------- c_regs = cutlass.Array(cutlass.Float32, 16, alignment=16, space=cutlass.AddressSpace.rmem) for i in cutlass.range_constexpr(16): c_regs[i] = cutlass.Float32(0.0) for ks in cutlass.range_constexpr(2): - mma_step(c_regs, a_frags, b_frags[ks * 8 : ks * 8 + 8], k_step=ks, M=16, N=32, ab_dtype=cfg.io_dtype) + mma_step(c_regs, d_inv_frags, c_frags[ks * 8 : ks * 8 + 8], k_step=ks, M=16, N=32, ab_dtype=cfg.io_dtype) for i in cutlass.range_constexpr(16): c_regs[i] = -c_regs[i] a_pack = [fp32_to_fp16(c_regs[2 * j], c_regs[2 * j + 1], dtype=cfg.io_dtype) for j in range(8)] - # ---- C = T @ A^{-1} ---------------------------------------------------------- - ai_frags = [] + # ---- C = T @ A^-1 ---------------------------------------------------------------- + a_inv_frags = [] for vs in cutlass.range_constexpr(4): - ai_frags += list( - nvvm.ldmatrix( - cutlass.inttoptr( - base_int + swizzle_lin_128b(((vs // 2) * 16) * 64 + (vs % 2) * 16 + lds4, row_stride_log2=6) * bpe, - cutlass.AddressSpace.smem, - cutlass.BFloat16, - ), - 4, - nvvm.MMALayout.COL, - ) - ) + row_a_inv = (vs // 2) * 16 + lane_row + a_inv_frags += list(nvvm.ldmatrix(base + row_a_inv * BT + swizzle_xor_128b(row_a_inv, (vs % 2) * 16 + lane_col), 4, nvvm.MMALayout.COL)) o_regs = cutlass.Array(cutlass.Float32, 16, alignment=16, space=cutlass.AddressSpace.rmem) for i in cutlass.range_constexpr(16): o_regs[i] = cutlass.Float32(0.0) for ks in cutlass.range_constexpr(2): - mma_step(o_regs, a_pack, ai_frags[ks * 8 : ks * 8 + 8], k_step=ks, M=16, N=32, ab_dtype=cfg.io_dtype) + mma_step(o_regs, a_pack, a_inv_frags[ks * 8 : ks * 8 + 8], k_step=ks, M=16, N=32, ab_dtype=cfg.io_dtype) o_pack = [fp32_to_fp16(o_regs[2 * j], o_regs[2 * j + 1], dtype=cfg.io_dtype) for j in range(8)] - # ---- store corrected C ------------------------------------------------------- + # ---- store corrected C ----------------------------------------------------------- nvvm.barrier_cta_sync_aligned( cfg.inverse_barrier_id, thread_count=cfg.inverse_barrier_threads, ) - nvvm.stmatrix( - cutlass.inttoptr(base_int + swizzle_lin_128b((32 + band * 16) * 64 + lds4, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16), - o_pack[0:4], - nvvm.MMALayout.ROW, - ) - nvvm.stmatrix( - cutlass.inttoptr(base_int + swizzle_lin_128b((32 + band * 16) * 64 + 16 + lds4, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16), - o_pack[4:8], - nvvm.MMALayout.ROW, - ) - - -# ---- Dynamic tile scheduler ------------------------------------------------------ + nvvm.stmatrix(base + row_d_inv * BT + swizzle_xor_128b(row_d_inv, lane_col), o_pack[0:4], nvvm.MMALayout.ROW) + nvvm.stmatrix(base + row_d_inv * BT + swizzle_xor_128b(row_d_inv, 16 + lane_col), o_pack[4:8], nvvm.MMALayout.ROW) @cute.jit -def sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ctas): +def scheduler_publish_next(cfg, bars, sScheduler, mScheduler, scheduler_state, tile_idx, num_ctas, elect_one): """TMA-LDG-warp side: pull the next tile off the global ticket, publish it.""" - if cutlass.const_expr(cfg.dyn_sched): - bars.mb_sched_done[sched_state.idx].wait(sched_state.phase) - if nvvm.elect_sync(): - fetched = cutlass.Int32(nvvm.atomicrmw("add", mSched.iterator, cutlass.Int32(1), mem_order="relaxed", syncscope="gpu")) - sSched[sched_state.idx] = num_ctas + fetched + if cutlass.const_expr(cfg.dynamic_scheduling): + bars.mb_scheduler_done[scheduler_state.idx].wait(scheduler_state.phase) + if elect_one: + fetched = cutlass.Int32(nvvm.atomicrmw("add", mScheduler.iterator, cutlass.Int32(1), mem_order="relaxed", syncscope="gpu")) + sScheduler[scheduler_state.idx] = num_ctas + fetched nvvm.bar_warp_sync(cute.arch.FULL_MASK) - next_tile = sSched[sched_state.idx] - if nvvm.elect_sync(): - bars.mb_sched_ready[sched_state.idx].arrive() - return next_tile, advance(sched_state, cfg.sched_stages) - return tile_idx + num_ctas, sched_state + next_tile = sScheduler[scheduler_state.idx] + if elect_one: + bars.mb_scheduler_ready[scheduler_state.idx].arrive() + return next_tile, advance(scheduler_state, cfg.scheduler_stages) + return tile_idx + num_ctas, scheduler_state @cute.jit -def sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas): +def scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one): """Consumer side: read the TMA-LDG warp's published next tile.""" - if cutlass.const_expr(cfg.dyn_sched): - bars.mb_sched_ready[sched_state.idx].wait(sched_state.phase) - next_tile = sSched[sched_state.idx] - if nvvm.elect_sync(): - bars.mb_sched_done[sched_state.idx].arrive() - return next_tile, advance(sched_state, cfg.sched_stages) - return tile_idx + num_ctas, sched_state + if cutlass.const_expr(cfg.dynamic_scheduling): + bars.mb_scheduler_ready[scheduler_state.idx].wait(scheduler_state.phase) + next_tile = sScheduler[scheduler_state.idx] + if elect_one: + bars.mb_scheduler_done[scheduler_state.idx].arrive() + return next_tile, advance(scheduler_state, cfg.scheduler_stages) + return tile_idx + num_ctas, scheduler_state @cute.jit @@ -450,19 +365,22 @@ def tmastg_warp( tidx, sCheckpoint_raw, desc_checkpoint_base, - sSched, + sScheduler, bars, ): """Epilogue warp role (warp 11): persistent scheduler loop issuing the per-chunk checkpoint TMA stores.""" - elect_one = nvvm.elect_sync() nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) - lidx = tidx % cfg.threads_per_warp - sched_state = PipelineState.start(phase=0) + + scheduler_state = PipelineState.start(phase=0) + + lane_idx = tidx % cfg.threads_per_warp + + elect_one = nvvm.elect_sync() tile_idx = cutlass.Int32(bidx) if cutlass.const_expr(cfg.enable_checkpoints): - checkpoint_granu = 64 + checkpoint_granule = 64 sCheckpoint_tma = SmemTile( base=sCheckpoint_raw, elems_per_stage=(cfg.checkpoint_cosize // cfg.smem_checkpoint_stages), @@ -470,37 +388,39 @@ def tmastg_warp( leading_byte_offset=0, stride_byte_offset=0, layout=0, - tma_loads_per_tile=cfg.d_v // checkpoint_granu, - tma_granu_elems=checkpoint_granu, - tma_subtile_stride_elems=cfg.d_k * checkpoint_granu, + tma_loads_per_tile=cfg.d_v // checkpoint_granule, + tma_granu_elems=checkpoint_granule, + tma_subtile_stride_elems=cfg.d_k * checkpoint_granule, ) checkpoint_store_cnt = cutlass.Int32(0) - ckpt_chunks = checkpoint_every_n_tokens // cutlass.Int32(cfg.b_t) + checkpoint_chunks = checkpoint_every_n_tokens // cutlass.Int32(cfg.b_t) heads_out = cutlass.Int32(cfg.n_heads_out) desc_qwords = cutlass.Int32(TENSOR_MAP_QWORDS) while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) - n_local = wend - cstart + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) + n_local = write_end - compute_start head_o = head_idx slot = batch_idx * desc_qwords if cutlass.const_expr(cfg.enable_checkpoints): desc_checkpoint_slot = (desc_checkpoint_base + slot).tospace(cutlass.AddressSpace.generic) - checkpoint_coord = (wstart + ckpt_chunks - cutlass.Int32(1)) // ckpt_chunks - checkpoint_mod = cstart % ckpt_chunks + checkpoint_coord = (write_start + checkpoint_chunks - cutlass.Int32(1)) // checkpoint_chunks + checkpoint_mod = compute_start % checkpoint_chunks if elect_one: tma_tensormap_acquire(desc_checkpoint_slot) if n_local > 0: for local_idx in cutlass.range(n_local): - chunk_idx = cstart + local_idx + chunk_idx = compute_start + local_idx did_checkpoint = cutlass.Int32(0) if cutlass.const_expr(cfg.enable_checkpoints): checkpoint_stage = checkpoint_store_cnt % cfg.smem_checkpoint_stages checkpoint_phase = (checkpoint_store_cnt // cfg.smem_checkpoint_stages) & cutlass.Int32(1) - if chunk_idx >= wstart and chunk_idx < wend: + if chunk_idx >= write_start and chunk_idx < write_end: if checkpoint_mod == 0: bars.mb_checkpoint_tmastg_ready[checkpoint_stage].wait(checkpoint_phase) checkpoint_slice = tma_slice_runtime_desc(desc_checkpoint_slot, cutlass.Int32(0), cutlass.Int32(0), checkpoint_coord, head_o) @@ -509,14 +429,14 @@ def tmastg_warp( checkpoint_coord += 1 did_checkpoint = cutlass.Int32(1) checkpoint_mod = checkpoint_mod + cutlass.Int32(1) - checkpoint_mod = cutlass.Int32(0) if checkpoint_mod == ckpt_chunks else checkpoint_mod + checkpoint_mod = cutlass.Int32(0) if checkpoint_mod == checkpoint_chunks else checkpoint_mod if did_checkpoint == 1: tma_store_wait(0) bars.mb_checkpoint_tmastg_done[checkpoint_stage].arrive() checkpoint_store_cnt = checkpoint_store_cnt + 1 - tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one) @cute.jit @@ -535,33 +455,37 @@ def gate_beta_warp( sCumsumlog, sCumprod, sBeta, - sSched, + sScheduler, bars, ): """Gate/Beta producer (warp 8): persistent scheduler loop + the cumsum/cumprod/Beta chunk loads.""" nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) + elect_one = nvvm.elect_sync() + gate_index = PipelineState.start(phase=1) beta_index = PipelineState.start(phase=1) - lidx = tidx % cfg.threads_per_warp + scheduler_state = PipelineState.start(phase=0) - a_l2 = cutlass.Float32(0.0) + lane_idx = tidx % cfg.threads_per_warp + + a = cutlass.Float32(0.0) bias = cutlass.Float32(0.0) - sched_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) - n_local = wend - cstart + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) + n_local = write_end - compute_start if cutlass.const_expr(cfg.safe_gate): if n_local > 0: - # per-head transform constants, fixed for the whole tile - a_l2 = -cute.math.exp2(mA_log[head_idx].to(cutlass.Float32) * cutlass.Float32(RCP_LN2), fastmath=True) * cutlass.Float32(RCP_LN2) + a = -cute.math.exp2(mA_log[head_idx].to(cutlass.Float32) * cutlass.Float32(RCP_LN2), fastmath=True) * cutlass.Float32(RCP_LN2) bias = mDt_bias[head_idx].to(cutlass.Float32) if n_local > 0: for local_idx in cutlass.range(n_local): - # ---- Gate load: GMEM -> SMEM (OOB neutral) ----------------------- - chunk_idx = cstart + local_idx + # ---- Gate load: GMEM -> SMEM (OOB neutral: 1.0 -> log2 = 0.0) -------- + chunk_idx = compute_start + local_idx n_cols = cfg.b_t // cfg.threads_per_warp chunk_offset = batch_start + chunk_idx * cfg.b_t gGateSeq = mGate[None, head_idx] @@ -571,30 +495,32 @@ def gate_beta_warp( gate_phase = gate_index.phase gate_index = advance(gate_index, cfg.smem_gate_stages) - pos_valid = [None] * n_cols - gate_vals = [cutlass.Float32(0.0)] * n_cols oob_neutral = cutlass.Float32(0.0) if cutlass.const_expr(cfg.log_gate) else cutlass.Float32(1.0) - for col in cutlass.range_constexpr(n_cols): - tok = chunk_offset + lidx + col * cfg.threads_per_warp - pos_valid[col] = cute.elem_less(tok, batch_end) - tok_clamped = min(tok, batch_end - 1) - gate_vals[col] = gGateSeq[tok_clamped] if pos_valid[col] else oob_neutral + toks = [chunk_offset + lane_idx + col * cfg.threads_per_warp for col in range(n_cols)] + pos_valid = [tok < batch_end for tok in toks] + gate_vals = [gGateSeq[min(tok, batch_end - 1)] if valid else oob_neutral for tok, valid in zip(toks, pos_valid)] if cutlass.const_expr(cfg.safe_gate): - # raw logits -> log2-domain decay: a_l2 * softplus(g + bias) (split-K scan arithmetic) - for col in cutlass.range_constexpr(n_cols): - contrib = a_l2 * softplus(gate_vals[col] + bias) - gate_vals[col] = contrib if pos_valid[col] else cutlass.Float32(0.0) + for col in cutlass.range_constexpr(0, n_cols, 2): + biased_lo, biased_hi = fadd2(gate_vals[col], gate_vals[col + 1], bias, bias) + sp_lo, sp_hi = softplus2(biased_lo, biased_hi) + contrib_lo, contrib_hi = fmul2(sp_lo, sp_hi, a, a) + gate_vals[col] = contrib_lo if pos_valid[col] else cutlass.Float32(0.0) + gate_vals[col + 1] = contrib_hi if pos_valid[col + 1] else cutlass.Float32(0.0) elif cutlass.const_expr(cfg.log_gate): - for col in cutlass.range_constexpr(n_cols): - gate_vals[col] = gate_vals[col] * cutlass.Float32(RCP_LN2) + rcp_ln2 = opaque_f32_zero() + cutlass.Float32(RCP_LN2) + for col in cutlass.range_constexpr(0, n_cols, 2): + gate_vals[col], gate_vals[col + 1] = fmul2(gate_vals[col], gate_vals[col + 1], rcp_ln2, rcp_ln2) else: - for col in cutlass.range_constexpr(n_cols): - gate_vals[col] = cute.math.log2(gate_vals[col] + 1e-10, fastmath=True) + floor = cutlass.Float32(1e-10) + for col in cutlass.range_constexpr(0, n_cols, 2): + shifted_lo, shifted_hi = fadd2(gate_vals[col], gate_vals[col + 1], floor, floor) + gate_vals[col] = cute.math.log2(shifted_lo, fastmath=True) + gate_vals[col + 1] = cute.math.log2(shifted_hi, fastmath=True) for offset in [1, 2, 4, 8, 16]: for col in cutlass.range_constexpr(n_cols): n = nvvm.shfl_sync(0xFFFFFFFF, gate_vals[col], offset, 0, kind=nvvm.Shfl.UP) - if lidx >= offset: + if lane_idx >= offset: gate_vals[col] = gate_vals[col] + n for col in cutlass.range_constexpr(1, n_cols): last_v = nvvm.shfl_sync( @@ -608,36 +534,34 @@ def gate_beta_warp( bars.mb_gate_done[gate_idx].wait(gate_phase) for col in cutlass.range_constexpr(n_cols): - pos = lidx + col * cfg.threads_per_warp + pos = lane_idx + col * cfg.threads_per_warp sCumsumlog[pos, 0, gate_idx] = gate_vals[col] sCumprod[pos, 0, gate_idx] = cute.math.exp2(gate_vals[col], fastmath=True) bars.mb_gate_ready[gate_idx].arrive() - # ---- Beta load: GMEM -> SMEM (per-element cp.async) -------------------------- + # ---- Beta load: GMEM -> SMEM (per-element cp.async) ------------------ beta_idx = beta_index.idx bars.mb_beta_done[beta_idx].wait(beta_index.phase) beta_index = advance(beta_index, cfg.smem_beta_stages) if cutlass.const_expr(cfg.beta_sigmoid): - # io-dtype logits -> sigmoid (tanh identity) -> fp32 SMEM for col in cutlass.range_constexpr(n_cols): - pos = lidx + col * cfg.threads_per_warp + pos = lane_idx + col * cfg.threads_per_warp beta_value = cutlass.Float32(0.0) if pos_valid[col]: beta_value = gBeta[pos].to(cutlass.Float32) - half = cutlass.Float32(0.5) - beta_value = (cute.math.tanh(beta_value * half, approx=True) * half + half).to(mBeta.element_type).to(cutlass.Float32) + beta_value = sigmoid(beta_value).to(mBeta.element_type).to(cutlass.Float32) sBeta[pos, 0, beta_idx] = beta_value bars.mb_beta_ready[beta_idx].arrive() else: for col in cutlass.range_constexpr(n_cols): - pos = lidx + col * cfg.threads_per_warp + pos = lane_idx + col * cfg.threads_per_warp src = gBeta.iterator + gBeta.layout((pos,)) dst = sBeta.iterator + sBeta.layout((pos, 0, beta_idx)) cp_size = cutlass.Int32(4) * cutlass.Int32(pos_valid[col]) nvvm.cp_async_shared_global(dst, src, 4, nvvm.LoadCacheModifier.CA, cp_size=cp_size) nvvm.cp_async_mbarrier_arrive(bars.mb_beta_ready[beta_idx].smem_ptr, noinc=True) - tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one) for _ in range(cfg.smem_gate_stages): bars.mb_gate_done[gate_index.idx].wait(gate_index.phase) @@ -648,7 +572,7 @@ def gate_beta_warp( @cute.jit -def mma_warp( +def tcgen05_mma_warp( cfg, total_tiles, bidx, @@ -659,30 +583,28 @@ def mma_warp( sKQ, sKQ_trans, sTinv, - sSched, + sScheduler, bars, ): """MMA issuer role (warp 10): persistent scheduler loop issuing every tcgen05 GEMM.""" - elect_one = nvvm.elect_sync() - nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) + kv_acc_index = PipelineState.start(phase=1) kq_index = PipelineState.start(phase=0) cg0_acc_index = PipelineState.start(phase=1) - kq_fused_index = PipelineState.start(phase=0) + kq_cg0_index = PipelineState.start(phase=0) tinv_index = PipelineState.start(phase=0) - state_inp_index = PipelineState.start(phase=0) - y_inp_ready = PipelineState.start(phase=0) - decay_u_inp_ready = PipelineState.start(phase=0) + state_input_index = PipelineState.start(phase=0) + y_input_ready = PipelineState.start(phase=0) + decay_u_input_ready = PipelineState.start(phase=0) + + elect_one = nvvm.elect_sync() nvvm.tcgen05_alloc(tmem_hold, cutlass.Int32(512), group=nvvm.CTAGroup.CTA_1) - nvvm.barrier_cta_sync_aligned( - cfg.tmem_alloc_barrier_id, - thread_count=cfg.tmem_alloc_barrier_threads, - ) + nvvm.barrier_cta_sync_aligned(cfg.tmem_lifecycle_barrier_id, thread_count=cfg.tmem_user_threads) - # ---- chunk-invariant GEMM descriptors ---------------------------------------- + # ---- chunk-invariant GEMM descriptors -------------------------------------------- idesc_kk = nvvm.Tcgen05InstrDesc.build( c_dtype=cutlass.Float32, a_dtype=cfg.io_dtype, @@ -690,7 +612,7 @@ def mma_warp( n_dim=cfg.b_t, m_dim=2 * cfg.b_t, ) - bmm_kk_desc = MmaDesc( + bmm_k_k_desc = MmaDesc( M=2 * cfg.b_t, N=cfg.b_t, K=cfg.d_k, @@ -710,7 +632,7 @@ def mma_warp( n_dim=cfg.b_t, m_dim=cfg.d_v, ) - bmm_k_state_desc = MmaDesc( + bmm_state_k_desc = MmaDesc( M=cfg.d_v, N=cfg.b_t, K=cfg.d_k, @@ -730,7 +652,7 @@ def mma_warp( n_dim=cfg.b_t, m_dim=cfg.d_v, ) - bmm_u_ts_desc = MmaDesc( + bmm_y_t_inv_desc = MmaDesc( M=cfg.d_v, N=cfg.b_t, K=cfg.b_t, @@ -751,7 +673,7 @@ def mma_warp( m_dim=cfg.d_k, b_major=1, ) - bmm_kv_desc = MmaDesc( + bmm_decay_u_k_desc = MmaDesc( M=cfg.d_k, N=cfg.d_v, K=cfg.b_t, @@ -767,151 +689,159 @@ def mma_warp( KQ_SEG = (2 * cfg.b_t * 64 * bpe) >> 4 KQ_BOX = (cfg.b_t * 64 * bpe) >> 4 KQ_HALF_K = (cfg.d_k // 16) // 2 - KQ_A_HALF = KQ_HALF_K * bmm_k_state_desc.tmem_advance_A + KQ_A_HALF = KQ_HALF_K * bmm_state_k_desc.tmem_advance_A KV_ACC_STAGE_COLS = cfg.d_v STATE_INP_STAGE_COLS = cfg.d_k // 2 INP_SLOT_COLS = cfg.b_t // 2 tmem_base = tmem_hold.load() - tmem_cg0_acc_col_f = tmem_base + cfg.tmem_cg0_acc_offset - tmem_state_col = tmem_base + cfg.tmem_state_acc_offset - tmem_state_inp_col = tmem_base + cfg.tmem_state_inp_offset - tmem_inp_col = tmem_base + cfg.tmem_y_decay_u_inp_offset - y_inp_ptr = nvvm.make_tmem_ptr(tmem_inp_col, cutlass.Int8) - decay_u_inp_ptr = nvvm.make_tmem_ptr(tmem_inp_col + INP_SLOT_COLS, cutlass.Int8) - k_state_acc_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_cg1_acc_offset, cutlass.Float32) + tmem_col = tmem_base & 0xFFFF + tmem_row = tmem_base >> 16 + row_lo_addr = tmem_row << 16 + row_hi_addr = (tmem_row + 16) << 16 + tmem_cg0_acc_col = tmem_col + cfg.tmem_cg0_acc_offset + tmem_state_col = tmem_col + cfg.tmem_state_acc_offset + tmem_state_input_col = tmem_col + cfg.tmem_state_input_offset + tmem_input_col = tmem_col + cfg.tmem_y_decay_u_input_offset + y_input_ptr = nvvm.make_tmem_ptr(tmem_input_col, cutlass.Int8) + decay_u_input_ptr = nvvm.make_tmem_ptr(tmem_input_col + INP_SLOT_COLS, cutlass.Int8) + k_state_acc_ptr = nvvm.make_tmem_ptr(tmem_col + cfg.tmem_cg1_acc_offset, cutlass.Float32) u_acc_ptr = k_state_acc_ptr - sched_state = PipelineState.start(phase=0) + scheduler_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) - n_local = wend - cstart + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) + n_local = write_end - compute_start - # ---- KK pair 0 = K(S) @ K^T, each member issued ahead of the loop ------- + # ---- KK pair 0 = K(S) @ K^T, each member issued ahead of the loop ------------ if n_local > 0: member0_acc_idx = cg0_acc_index.idx bars.mb_cg0_acc_done[member0_acc_idx].wait(cg0_acc_index.phase) cg0_acc_index = advance(cg0_acc_index, cfg.tmem_cg0_acc_stages) - kqf_idx = kq_fused_index.idx - bars.mb_kq_ready[kqf_idx].wait(kq_fused_index.phase) - kq_fused_index = advance(kq_fused_index, cfg.smem_kq_stages) - desc_kqf = sKQ[kqf_idx].desc() - acc_cg0 = nvvm.make_tmem_ptr(tmem_cg0_acc_col_f + member0_acc_idx * cfg.b_t, cutlass.Float32) - mma_ss(bmm_kk_desc, desc_kqf, desc_kqf, acc_cg0, accumulate=False, k_count=KQ_HALF_K) - mma_ss(bmm_kk_desc, desc_kqf + KQ_SEG, desc_kqf + KQ_SEG, acc_cg0, accumulate=True, k_count=KQ_HALF_K) + kq_cg0_idx = kq_cg0_index.idx + bars.mb_kq_ready[kq_cg0_idx].wait(kq_cg0_index.phase) + kq_cg0_index = advance(kq_cg0_index, cfg.smem_kq_stages) + desc_kq_cg0 = sKQ[kq_cg0_idx].desc() + acc_cg0 = nvvm.make_tmem_ptr(tmem_cg0_acc_col + member0_acc_idx * cfg.b_t, cutlass.Float32) + mma_ss(bmm_k_k_desc, desc_kq_cg0, desc_kq_cg0, acc_cg0, accumulate=False, k_count=KQ_HALF_K) + mma_ss(bmm_k_k_desc, desc_kq_cg0 + KQ_SEG, desc_kq_cg0 + KQ_SEG, acc_cg0, accumulate=True, k_count=KQ_HALF_K) if elect_one: bars.mb_cg0_acc_ready[member0_acc_idx].arrive(cta_group=1) if n_local > 1: member1_acc_idx = cg0_acc_index.idx bars.mb_cg0_acc_done[member1_acc_idx].wait(cg0_acc_index.phase) cg0_acc_index = advance(cg0_acc_index, cfg.tmem_cg0_acc_stages) - kqf_idx = kq_fused_index.idx - bars.mb_kq_ready[kqf_idx].wait(kq_fused_index.phase) - kq_fused_index = advance(kq_fused_index, cfg.smem_kq_stages) - desc_kqf = sKQ[kqf_idx].desc() - desc_kqf_member1 = desc_kqf + KQ_BOX - acc_cg0 = nvvm.make_tmem_ptr(tmem_cg0_acc_col_f + member1_acc_idx * cfg.b_t, cutlass.Float32) - mma_ss(bmm_kk_desc, desc_kqf, desc_kqf_member1, acc_cg0, accumulate=False, k_count=KQ_HALF_K) - mma_ss(bmm_kk_desc, desc_kqf + KQ_SEG, desc_kqf_member1 + KQ_SEG, acc_cg0, accumulate=True, k_count=KQ_HALF_K) + kq_cg0_idx = kq_cg0_index.idx + bars.mb_kq_ready[kq_cg0_idx].wait(kq_cg0_index.phase) + kq_cg0_index = advance(kq_cg0_index, cfg.smem_kq_stages) + desc_kq_cg0 = sKQ[kq_cg0_idx].desc() + desc_k_cg0 = desc_kq_cg0 + KQ_BOX + acc_cg0 = nvvm.make_tmem_ptr(tmem_cg0_acc_col + member1_acc_idx * cfg.b_t, cutlass.Float32) + mma_ss(bmm_k_k_desc, desc_kq_cg0, desc_k_cg0, acc_cg0, accumulate=False, k_count=KQ_HALF_K) + mma_ss(bmm_k_k_desc, desc_kq_cg0 + KQ_SEG, desc_k_cg0 + KQ_SEG, acc_cg0, accumulate=True, k_count=KQ_HALF_K) if elect_one: bars.mb_cg0_acc_ready[member1_acc_idx].arrive(cta_group=1) + if cutlass.const_expr(cfg.use_initial_state): + seed_state = compute_start == 0 for local_idx in cutlass.range(n_local): # noqa: B007 if cutlass.const_expr(cfg.use_initial_state): - if local_idx == 0: - if elect_one: - bars.mb_state_acc_ready[kv_acc_index.idx].arrive(cta_group=1) - kv_acc_index = advance(kv_acc_index, cfg.tmem_state_acc_stages) - have_state = cutlass.Boolean(True) if cutlass.const_expr(cfg.use_initial_state) else local_idx > 0 + have_state = local_idx > 0 or seed_state + else: + have_state = local_idx > 0 kq_idx = kq_index.idx member = local_idx & 1 - state_inp_idx = state_inp_index.idx + state_input_idx = state_input_index.idx tinv_idx = tinv_index.idx kv_acc_idx = kv_acc_index.idx kq_member_off = member * KQ_BOX desc_k = sKQ[kq_idx].desc() + kq_member_off desc_tinv = sTinv[tinv_idx].desc() - desc_kt = sKQ_trans[kq_idx].desc() + kq_member_off - state_a_ptr = nvvm.make_tmem_ptr(tmem_state_inp_col + state_inp_idx * STATE_INP_STAGE_COLS, cutlass.Int8) + desc_k_trans = sKQ_trans[kq_idx].desc() + kq_member_off + state_a_ptr = nvvm.make_tmem_ptr(tmem_state_input_col + state_input_idx * STATE_INP_STAGE_COLS, cutlass.Int8) state_acc_ptr = nvvm.make_tmem_ptr(tmem_state_col + kv_acc_idx * KV_ACC_STAGE_COLS, cutlass.Float32) kq_index = advance(kq_index, cfg.smem_kq_stages) - # ---- KK pair lookahead (member 1) = K(S) @ K^T ----------------------- + # ---- KK pair lookahead (member 1) = K(S) @ K^T --------------------------- if member == 1: if local_idx + 2 < n_local: member1_acc_idx = cg0_acc_index.idx bars.mb_cg0_acc_done[member1_acc_idx].wait(cg0_acc_index.phase) cg0_acc_index = advance(cg0_acc_index, cfg.tmem_cg0_acc_stages) - kqf_idx = kq_fused_index.idx - bars.mb_kq_ready[kqf_idx].wait(kq_fused_index.phase) - kq_fused_index = advance(kq_fused_index, cfg.smem_kq_stages) - desc_kqf = sKQ[kqf_idx].desc() - desc_kqf_member1 = desc_kqf + KQ_BOX - acc_cg0 = nvvm.make_tmem_ptr(tmem_cg0_acc_col_f + member1_acc_idx * cfg.b_t, cutlass.Float32) - mma_ss(bmm_kk_desc, desc_kqf, desc_kqf_member1, acc_cg0, accumulate=False, k_count=KQ_HALF_K) - mma_ss(bmm_kk_desc, desc_kqf + KQ_SEG, desc_kqf_member1 + KQ_SEG, acc_cg0, accumulate=True, k_count=KQ_HALF_K) + kq_cg0_idx = kq_cg0_index.idx + bars.mb_kq_ready[kq_cg0_idx].wait(kq_cg0_index.phase) + kq_cg0_index = advance(kq_cg0_index, cfg.smem_kq_stages) + desc_kq_cg0 = sKQ[kq_cg0_idx].desc() + desc_k_cg0 = desc_kq_cg0 + KQ_BOX + acc_cg0 = nvvm.make_tmem_ptr(tmem_cg0_acc_col + member1_acc_idx * cfg.b_t, cutlass.Float32) + mma_ss(bmm_k_k_desc, desc_kq_cg0, desc_k_cg0, acc_cg0, accumulate=False, k_count=KQ_HALF_K) + mma_ss(bmm_k_k_desc, desc_kq_cg0 + KQ_SEG, desc_k_cg0 + KQ_SEG, acc_cg0, accumulate=True, k_count=KQ_HALF_K) if elect_one: bars.mb_cg0_acc_ready[member1_acc_idx].arrive(cta_group=1) - # ---- K*state = state(T) @ K^T (GEMM 3) ---------------------------------------- + # ---- k state (GEMM 3) = state(T) @ K^T ----------------------------------- if have_state: - bars.mb_state_inp_ready[state_inp_idx].wait(state_inp_index.phase) - state_inp_index = advance(state_inp_index, cfg.tmem_state_inp_stages) + bars.mb_state_input_ready[state_input_idx].wait(state_input_index.phase) + state_input_index = advance(state_input_index, cfg.tmem_state_input_stages) for k in cutlass.range_constexpr(KQ_HALF_K): - mma_ts_step(bmm_k_state_desc, state_a_ptr, desc_k, k_state_acc_ptr, k, cutlass.Boolean(k > 0)) + mma_ts_step(bmm_state_k_desc, state_a_ptr, desc_k, k_state_acc_ptr, k, cutlass.Boolean(k > 0)) for k in cutlass.range_constexpr(KQ_HALF_K): - mma_ts_step(bmm_k_state_desc, state_a_ptr.subview(KQ_A_HALF), desc_k + KQ_SEG, k_state_acc_ptr, k, cutlass.Boolean(True)) + mma_ts_step(bmm_state_k_desc, state_a_ptr.subview(KQ_A_HALF), desc_k + KQ_SEG, k_state_acc_ptr, k, cutlass.Boolean(True)) if elect_one: bars.mb_k_state_acc_ready[0].arrive(cta_group=1) - # ---- U = Y(T) @ T_inv (GEMM 5) --------------------------------------------- + # ---- U (GEMM 5) = Y(T) @ T^-1 -------------------------------------------- bars.mb_t_inv_ready[tinv_idx].wait(tinv_index.phase) tinv_index = advance(tinv_index, cfg.smem_t_inv_stages) - bars.mb_y_inp_ready[0].wait(y_inp_ready.phase) - y_inp_ready = advance(y_inp_ready, 1) + bars.mb_y_input_ready[0].wait(y_input_ready.phase) + y_input_ready = advance(y_input_ready, 1) for k in cutlass.range_constexpr(cfg.b_t // 16): - mma_ts_step(bmm_u_ts_desc, y_inp_ptr, desc_tinv, u_acc_ptr, k, cutlass.Boolean(k > 0)) + mma_ts_step(bmm_y_t_inv_desc, y_input_ptr, desc_tinv, u_acc_ptr, k, cutlass.Boolean(k > 0)) if elect_one: bars.mb_u_acc_ready[0].arrive(cta_group=1) bars.mb_t_inv_done[tinv_idx].arrive(cta_group=1) - # ---- KK pair lookahead (member 0) = K(S) @ K^T ----------------------- + # ---- KK pair lookahead (member 0) = K(S) @ K^T --------------------------- if member == 0: if local_idx + 2 < n_local: member0_acc_idx = cg0_acc_index.idx bars.mb_cg0_acc_done[member0_acc_idx].wait(cg0_acc_index.phase) cg0_acc_index = advance(cg0_acc_index, cfg.tmem_cg0_acc_stages) - kqf_idx = kq_fused_index.idx - bars.mb_kq_ready[kqf_idx].wait(kq_fused_index.phase) - kq_fused_index = advance(kq_fused_index, cfg.smem_kq_stages) - desc_kqf = sKQ[kqf_idx].desc() - acc_cg0 = nvvm.make_tmem_ptr(tmem_cg0_acc_col_f + member0_acc_idx * cfg.b_t, cutlass.Float32) - mma_ss(bmm_kk_desc, desc_kqf, desc_kqf, acc_cg0, accumulate=False, k_count=KQ_HALF_K) - mma_ss(bmm_kk_desc, desc_kqf + KQ_SEG, desc_kqf + KQ_SEG, acc_cg0, accumulate=True, k_count=KQ_HALF_K) + kq_cg0_idx = kq_cg0_index.idx + bars.mb_kq_ready[kq_cg0_idx].wait(kq_cg0_index.phase) + kq_cg0_index = advance(kq_cg0_index, cfg.smem_kq_stages) + desc_kq_cg0 = sKQ[kq_cg0_idx].desc() + acc_cg0 = nvvm.make_tmem_ptr(tmem_cg0_acc_col + member0_acc_idx * cfg.b_t, cutlass.Float32) + mma_ss(bmm_k_k_desc, desc_kq_cg0, desc_kq_cg0, acc_cg0, accumulate=False, k_count=KQ_HALF_K) + mma_ss(bmm_k_k_desc, desc_kq_cg0 + KQ_SEG, desc_kq_cg0 + KQ_SEG, acc_cg0, accumulate=True, k_count=KQ_HALF_K) if elect_one: bars.mb_cg0_acc_ready[member0_acc_idx].arrive(cta_group=1) - # ---- state += decayed U(T) @ K (GEMM 7) ---------------------------------- - bars.mb_decay_u_inp_ready[0].wait(decay_u_inp_ready.phase) - decay_u_inp_ready = advance(decay_u_inp_ready, 1) + # ---- state (GEMM 7) += decayed U(T) @ K ---------------------------------- + bars.mb_decay_u_input_ready[0].wait(decay_u_input_ready.phase) + decay_u_input_ready = advance(decay_u_input_ready, 1) kv_acc_index = advance(kv_acc_index, cfg.tmem_state_acc_stages) for k in cutlass.range_constexpr(cfg.b_t // 16): - mma_ts_step(bmm_kv_desc, decay_u_inp_ptr, desc_kt, state_acc_ptr, k, cutlass.Boolean(True) if cutlass.const_expr(k > 0) else have_state) + mma_ts_step( + bmm_decay_u_k_desc, decay_u_input_ptr, desc_k_trans, state_acc_ptr, k, cutlass.Boolean(True) if cutlass.const_expr(k > 0) else have_state + ) if elect_one: bars.mb_state_acc_ready[kv_acc_idx].arrive(cta_group=1) bars.mb_kq_done[kq_idx].arrive(cta_group=1) - tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one) bars.mb_tmem_done[0].wait(0) nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_1) nvvm.tcgen05_dealloc( - nvvm.make_tmem_ptr(tmem_base, cutlass.Int8), + nvvm.make_tmem_ptr(tmem_col, cutlass.Int8), cutlass.Int32(512), group=nvvm.CTAGroup.CTA_1, ) @@ -929,34 +859,35 @@ def tmaldg_warp( sV_raw, desc_k_base, desc_v_base, - mSched, - sSched, + mScheduler, + sScheduler, bars, ): """TMA-LDG warp role (warp 9): persistent scheduler loop + per-chunk K/V G->S TMA loads.""" - elect_one = nvvm.elect_sync() nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) + kq_index = PipelineState.start(phase=1) v_index = PipelineState.start(phase=1) - sched_state = PipelineState.start(phase=1) - tile_idx = cutlass.Int32(bidx) + scheduler_state = PipelineState.start(phase=1) + elect_one = nvvm.elect_sync() + tile_idx = cutlass.Int32(bidx) bpe = cfg.io_dtype.width // 8 - granu = 128 // bpe + granule = 128 // bpe bt = cfg.b_t - kq_stage_elems = cfg.kq_cosize // cfg.smem_kq_stages - kq_box_elems = kq_stage_elems // 4 + kq_stage_elements = cfg.kq_cosize // cfg.smem_kq_stages + kq_box_elements = kq_stage_elements // 4 sKQ_lo_tma = SmemTile( base=sKQ_raw, - elems_per_stage=kq_stage_elems, + elems_per_stage=kq_stage_elements, stages=cfg.smem_kq_stages, leading_byte_offset=0, stride_byte_offset=0, layout=0, tma_loads_per_tile=2, - tma_granu_elems=granu, - tma_subtile_stride_elems=2 * bt * granu, + tma_granu_elems=granule, + tma_subtile_stride_elems=2 * bt * granule, ) sV_tma = SmemTile( base=sV_raw, @@ -966,14 +897,16 @@ def tmaldg_warp( stride_byte_offset=0, layout=0, tma_loads_per_tile=2, - tma_granu_elems=granu, + tma_granu_elems=granule, tma_subtile_stride_elems=4096, ) heads_out = cutlass.Int32(cfg.n_heads_out) desc_qwords = cutlass.Int32(TENSOR_MAP_QWORDS) while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) head_k = head_idx if cfg.k_ratio == 1 else head_idx // cutlass.Int32(cfg.k_ratio) head_v = head_idx if cfg.v_ratio == 1 else head_idx // cutlass.Int32(cfg.v_ratio) @@ -984,34 +917,34 @@ def tmaldg_warp( tma_tensormap_acquire(desc_k_slot) tma_tensormap_acquire(desc_v_slot) - if wend > cstart: + if write_end > compute_start: kq_idx = kq_index.idx bars.mb_kq_done[kq_idx].wait(kq_index.phase) kq_index = advance(kq_index, cfg.smem_kq_stages) if elect_one: bars.mb_kq_ready[kq_idx].arrive(n_bytes=cfg.tma_kq_bytes) - tok_coord = cstart * cutlass.Int32(cfg.b_t) + tok_coord = compute_start * cutlass.Int32(cfg.b_t) k_slice = tma_slice_runtime_desc(desc_k_slot, cutlass.Int32(0), head_k, tok_coord) kq_tile = sKQ_lo_tma[kq_idx] tma_load_tile(kq_tile, k_slice, bars.mb_kq_ready[kq_idx].smem_ptr, acquire=False) - for chunk_idx in cutlass.range(cstart + 1, wend): + for chunk_idx in cutlass.range(compute_start + 1, write_end): tok_coord = chunk_idx * cutlass.Int32(cfg.b_t) - # ---- K load ------------------------------------------------------ + # ---- K load ---------------------------------------------------------- kq_idx = kq_index.idx bars.mb_kq_done[kq_idx].wait(kq_index.phase) kq_index = advance(kq_index, cfg.smem_kq_stages) if elect_one: bars.mb_kq_ready[kq_idx].arrive(n_bytes=cfg.tma_kq_bytes) - member = (chunk_idx - cstart) & 1 + member = (chunk_idx - compute_start) & 1 k_slice = tma_slice_runtime_desc(desc_k_slot, cutlass.Int32(0), head_k, tok_coord) kq_tile = sKQ_lo_tma[kq_idx] if member == 0: tma_load_tile(kq_tile, k_slice, bars.mb_kq_ready[kq_idx].smem_ptr, acquire=False) else: - tma_load_tile(kq_tile.shifted(kq_box_elems), k_slice, bars.mb_kq_ready[kq_idx].smem_ptr, acquire=False) + tma_load_tile(kq_tile.shifted(kq_box_elements), k_slice, bars.mb_kq_ready[kq_idx].smem_ptr, acquire=False) - # ---- V load ------------------------------------------------------ + # ---- V load ---------------------------------------------------------- v_idx = v_index.idx bars.mb_v_done[v_idx].wait(v_index.phase) v_index = advance(v_index, cfg.smem_v_stages) @@ -1026,11 +959,11 @@ def tmaldg_warp( v_index = advance(v_index, cfg.smem_v_stages) if elect_one: bars.mb_v_ready[v_idx].arrive(n_bytes=cfg.tma_v_bytes) - v_tok = (wend - cutlass.Int32(1)) * cutlass.Int32(cfg.b_t) + v_tok = (write_end - cutlass.Int32(1)) * cutlass.Int32(cfg.b_t) v_slice = tma_slice_runtime_desc(desc_v_slot, cutlass.Int32(0), head_v, v_tok) tma_load_tile(sV_tma[v_idx], v_slice, bars.mb_v_ready[v_idx].smem_ptr, acquire=False) - tile_idx, sched_state = sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_publish_next(cfg, bars, sScheduler, mScheduler, scheduler_state, tile_idx, num_ctas, elect_one) for _ in range(cfg.smem_kq_stages): bars.mb_kq_done[kq_index.idx].wait(kq_index.phase) @@ -1055,59 +988,60 @@ def compute0_warp_group( sTinv, sCheckpoint_raw, checkpoint_every_n_tokens, - sSched, + sScheduler, bars, ): """Compute warp-group 0 role (warps 0-3): persistent scheduler loop building the per-chunk beta-scaled T_inv operands.""" nvvm.setmaxregister(cfg.num_regs_compute_group_0, nvvm.SetMaxRegisterAction.INCREASE) + elect_one = nvvm.elect_sync() + gate_index = PipelineState.start(phase=0) beta_index = PipelineState.start(phase=0) cg0_acc_ready = PipelineState.start(phase=0) tinv_index = PipelineState.start(phase=1) - - nvvm.barrier_cta_sync_aligned( - cfg.tmem_alloc_barrier_id, - thread_count=cfg.tmem_alloc_barrier_threads, - ) - tmem_base = tmem_hold.load() + scheduler_state = PipelineState.start(phase=0) num_threads_cg0 = cfg.threads_per_warp * len(cfg.compute_group_0_warp_ids) cg0_tidx = tidx % num_threads_cg0 warp_id = cg0_tidx // cfg.threads_per_warp - lane_id = cg0_tidx % cfg.threads_per_warp + lane_idx = cg0_tidx % cfg.threads_per_warp inverse_local_warp = warp_id % 2 + store_row = warp_id * 16 + lane_idx % 16 + store_col = (lane_idx // 16) * 8 + pair_half = warp_id // 2 half_row_base = inverse_local_warp * 32 - bpe = cfg.io_dtype.width // 8 num_vals = 32 FRAG_COLS = 16 ACC_N_FRAGS = cfg.b_t // FRAG_COLS - store_row = warp_id * 16 + lane_id % 16 - store_row_frag = lane_id % 16 - store_col = (lane_id // 16) * 8 - tmem_warp_row = warp_id * cfg.threads_per_warp - tmem_cg0_acc_col = tmem_base + cfg.tmem_cg0_acc_offset + store_row_frag = lane_idx % 16 ACC_STAGE_COLS = cfg.b_t mask_zero = opaque_f32_zero() - crow_lo = warp_id * 16 + lane_id // 4 - crow_hi = crow_lo + 8 - - sched_state = PipelineState.start(phase=0) + chunk_row_lo = warp_id * 16 + lane_idx // 4 + chunk_row_hi = chunk_row_lo + 8 tile_idx = cutlass.Int32(bidx) + + nvvm.barrier_cta_sync_aligned(cfg.tmem_lifecycle_barrier_id, thread_count=cfg.tmem_user_threads) + tmem_base = tmem_hold.load() + tmem_col = tmem_base & 0xFFFF + tmem_row = tmem_base >> 16 + row_lo_addr = tmem_row << 16 + row_hi_addr = (tmem_row + 16) << 16 + tmem_cg0_acc_col = tmem_col + cfg.tmem_cg0_acc_offset while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) - n_local = wend - cstart + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) + n_local = write_end - compute_start n_pairs = (n_local + 1) // 2 for pair_i in cutlass.range(n_pairs): - # An odd chunk count leaves the last pair with member 0 only; have_m1 - # is uniform across CG0, so the shared barriers below stay aligned. have_m1 = pair_i * 2 + 1 < n_local do_kk = have_m1 or pair_half == 0 - # ---- Gate rows for this warp's KK member role ------------------------ + # ---- Gate rows for this warp's KK member role ---------------------------- gate0_idx = gate_index.idx bars.mb_gate_ready[gate0_idx].wait(gate_index.phase) gate_index = advance(gate_index, cfg.smem_gate_stages) @@ -1118,7 +1052,7 @@ def compute0_warp_group( gate_index = advance(gate_index, cfg.smem_gate_stages) kk_gate_idx = gate1_idx if pair_half == 1 else gate0_idx - row_u0_lo = half_row_base + lane_id // 4 + row_u0_lo = half_row_base + lane_idx // 4 row_u0_hi = row_u0_lo + 8 row_u1_lo = row_u0_lo + 16 row_u1_hi = row_u0_lo + 24 @@ -1129,18 +1063,18 @@ def compute0_warp_group( kk_col_cumsumlog = [] for g in cutlass.range_constexpr(8): for b in cutlass.range_constexpr(2): - ccol = (lane_id % 4) * 2 + g * 8 + b - kk_col_cumsumlog.append(sCumsumlog[ccol, 0, kk_gate_idx]) + chunk_col = (lane_idx % 4) * 2 + g * 8 + b + kk_col_cumsumlog.append(sCumsumlog[chunk_col, 0, kk_gate_idx]) decay_t_kk = [] for u in cutlass.range_constexpr(2): for k in cutlass.range_constexpr(num_vals): hi_row = ((k // 2) % 2) == 1 - crow_u0 = row_u0_hi if cutlass.const_expr(hi_row) else row_u0_lo - crow_u1 = row_u1_hi if cutlass.const_expr(hi_row) else row_u1_lo - crow = crow_u1 if cutlass.const_expr(u == 1) else crow_u0 - ccol = (lane_id % 4) * 2 + ((k // 4) * 8 + k % 2) - is_lower = crow >= ccol + chunk_row_u0 = row_u0_hi if cutlass.const_expr(hi_row) else row_u0_lo + chunk_row_u1 = row_u1_hi if cutlass.const_expr(hi_row) else row_u1_lo + chunk_row = chunk_row_u1 if cutlass.const_expr(u == 1) else chunk_row_u0 + chunk_col = (lane_idx % 4) * 2 + ((k // 4) * 8 + k % 2) + is_lower = chunk_row >= chunk_col row_cumsumlog = kk_row_cumsumlog[u * 2 + (1 if hi_row else 0)] col = (k // 4) * 2 + (k % 2) decay_t_kk.append(cute.math.exp2(row_cumsumlog - kk_col_cumsumlog[col], fastmath=True) if is_lower else mask_zero) @@ -1161,7 +1095,7 @@ def compute0_warp_group( for r in (row_u0_lo, row_u0_hi, row_u1_lo, row_u1_hi): kk_beta.append(sBeta[r, 0, kk_beta_idx]) - # ---- KK_epi (each warp pair stages its own member) ------------------- + # ---- KK epilogue (each warp pair stages its own member) ------------------ acc0_idx = cg0_acc_ready.idx acc0_phase = cg0_acc_ready.phase cg0_acc_ready = advance(cg0_acc_ready, cfg.tmem_cg0_acc_stages) @@ -1191,17 +1125,13 @@ def compute0_warp_group( kk_base = tinv1_base if pair_half == 1 else tinv0_base if do_kk: bars.mb_cg0_acc_ready[kk_acc_idx].wait(kk_acc_phase) - kk_vec0 = nvvm.tcgen05_ld( - "16x256b", nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_cg0_acc_col + kk_acc_idx * ACC_STAGE_COLS, cutlass.Float32), num=8 - ) - kk_vec1 = nvvm.tcgen05_ld( - "16x256b", nvvm.make_tmem_ptr(((tmem_warp_row + 16) << 16) + tmem_cg0_acc_col + kk_acc_idx * ACC_STAGE_COLS, cutlass.Float32), num=8 - ) + kk_vec_lo = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_lo_addr + tmem_cg0_acc_col + kk_acc_idx * ACC_STAGE_COLS, cutlass.Float32), num=8) + kk_vec_hi = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_hi_addr + tmem_cg0_acc_col + kk_acc_idx * ACC_STAGE_COLS, cutlass.Float32), num=8) nvvm.tcgen05_wait("load") bars.mb_cg0_acc_done[kk_acc_idx].arrive() bars.mb_t_inv_done[kk_tinv_idx].wait(kk_tinv_phase) for u in cutlass.range_constexpr(2): - kk_vec = kk_vec1 if cutlass.const_expr(u == 1) else kk_vec0 + kk_vec = kk_vec_hi if cutlass.const_expr(u == 1) else kk_vec_lo kk_pack = [] for k in cutlass.range_constexpr(num_vals // 2): row_beta = kk_beta[u * 2 + 1] if cutlass.const_expr((k % 2) == 1) else kk_beta[u * 2] @@ -1211,18 +1141,12 @@ def compute0_warp_group( st_row = half_row_base + u * 16 + store_row_frag for c in cutlass.range_constexpr(ACC_N_FRAGS): nvvm.stmatrix( - cutlass.inttoptr( - kk_base + (st_row * cfg.b_t + swizzle_xor_128b(st_row, store_col + c * FRAG_COLS)) * bpe, - cutlass.AddressSpace.smem, - cutlass.BFloat16, - ), + kk_base + st_row * cfg.b_t + swizzle_xor_128b(st_row, store_col + c * FRAG_COLS), [kk_pack[c * 4 + 0], kk_pack[c * 4 + 1], kk_pack[c * 4 + 2], kk_pack[c * 4 + 3]], nvvm.MMALayout.ROW, ) - # ---- pair inverse: warps 0-1 own matrix 0, warps 2-3 matrix 1 -------- - # With no member 1 there is one matrix, so warps 2-3 idle through the - # per-warp steps; the barriers below stay unconditional. + # ---- pair inverse: warps 0-1 own matrix 0, warps 2-3 matrix 1 ------------ inv_base = tinv0_base if have_m1: inv_base = tinv1_base if warp_id >= 2 else tinv0_base @@ -1234,16 +1158,16 @@ def compute0_warp_group( thread_count=cfg.inverse_barrier_threads, ) if do_inv: - invert_diagonal_NxN(cfg, inv_base, (inverse_local_warp * cfg.threads_per_warp + lane_id) // 8, cg0_tidx, 8) + invert_diagonal_NxN(cfg, inv_base, (inverse_local_warp * cfg.threads_per_warp + lane_idx) // 8, cg0_tidx, 8) nvvm.barrier_cta_sync_aligned( cfg.inverse_barrier_id, thread_count=cfg.inverse_barrier_threads, ) # 8x8 -> 16x16 (both matrices per warp) - blockwise_diagonal_8x8_to_16x16(cfg, tinv0_base, warp_id * 16, lane_id) + blockwise_diagonal_8x8_to_16x16(cfg, tinv0_base, warp_id * 16, lane_idx) if have_m1: - blockwise_diagonal_8x8_to_16x16(cfg, tinv1_base, warp_id * 16, lane_id) + blockwise_diagonal_8x8_to_16x16(cfg, tinv1_base, warp_id * 16, lane_idx) nvvm.barrier_cta_sync_aligned( cfg.inverse_barrier_id, thread_count=cfg.inverse_barrier_threads, @@ -1251,34 +1175,28 @@ def compute0_warp_group( # 16x16 -> 32x32, one tile per warp within the group if do_inv: - blockwise_diagonal_16x16_to_32x32(cfg, inv_base, inverse_local_warp * 32, lane_id) + blockwise_diagonal_16x16_to_32x32(cfg, inv_base, inverse_local_warp * 32, lane_idx) nvvm.barrier_cta_sync_aligned( cfg.inverse_barrier_id, thread_count=cfg.inverse_barrier_threads, ) # 32x32 -> 64x64, two warps per matrix - # all four warps: this step carries its own CG0 barrier, and on the tail - # inv_base is matrix 0 for every warp so the duplicate band is identical - blockwise_diagonal_32x32_to_64x64(cfg, inv_base, inverse_local_warp, lane_id) + blockwise_diagonal_32x32_to_64x64(cfg, inv_base, inverse_local_warp, lane_idx) nvvm.barrier_cta_sync_aligned( cfg.inverse_barrier_id, thread_count=cfg.inverse_barrier_threads, ) - # ---- Beta column-scaling + publish, stage 0 -------------------------- + # ---- Beta column-scaling + publish, stage 0 ------------------------------ beta_col = [] for k in cutlass.range_constexpr(num_vals): - beta_col.append(sBeta[(lane_id % 4) * 2 + ((k // 4) * 8 + k % 2), 0, beta0_idx]) + beta_col.append(sBeta[(lane_idx % 4) * 2 + ((k // 4) * 8 + k % 2), 0, beta0_idx]) tinv_frags = [] for c in cutlass.range_constexpr(ACC_N_FRAGS): tinv_frags += list( nvvm.ldmatrix( - cutlass.inttoptr( - tinv0_base + (store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS)) * bpe, - cutlass.AddressSpace.smem, - cutlass.BFloat16, - ), + tinv0_base + store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS), 4, nvvm.MMALayout.ROW, ) @@ -1290,11 +1208,7 @@ def compute0_warp_group( tinv_pack.append(fp32_to_fp16(s0, s1, dtype=cfg.io_dtype)) for c in cutlass.range_constexpr(ACC_N_FRAGS): nvvm.stmatrix( - cutlass.inttoptr( - tinv0_base + (store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS)) * bpe, - cutlass.AddressSpace.smem, - cutlass.BFloat16, - ), + tinv0_base + store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS), [tinv_pack[c * 4 + 0], tinv_pack[c * 4 + 1], tinv_pack[c * 4 + 2], tinv_pack[c * 4 + 3]], nvvm.MMALayout.ROW, ) @@ -1306,16 +1220,12 @@ def compute0_warp_group( # ---- Beta column-scaling + publish, stage 1 -------------------------- beta_col = [] for k in cutlass.range_constexpr(num_vals): - beta_col.append(sBeta[(lane_id % 4) * 2 + ((k // 4) * 8 + k % 2), 0, beta1_idx]) + beta_col.append(sBeta[(lane_idx % 4) * 2 + ((k // 4) * 8 + k % 2), 0, beta1_idx]) tinv_frags = [] for c in cutlass.range_constexpr(ACC_N_FRAGS): tinv_frags += list( nvvm.ldmatrix( - cutlass.inttoptr( - tinv1_base + (store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS)) * bpe, - cutlass.AddressSpace.smem, - cutlass.BFloat16, - ), + tinv1_base + store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS), 4, nvvm.MMALayout.ROW, ) @@ -1327,11 +1237,7 @@ def compute0_warp_group( tinv_pack.append(fp32_to_fp16(s0, s1, dtype=cfg.io_dtype)) for c in cutlass.range_constexpr(ACC_N_FRAGS): nvvm.stmatrix( - cutlass.inttoptr( - tinv1_base + (store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS)) * bpe, - cutlass.AddressSpace.smem, - cutlass.BFloat16, - ), + tinv1_base + store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS), [tinv_pack[c * 4 + 0], tinv_pack[c * 4 + 1], tinv_pack[c * 4 + 2], tinv_pack[c * 4 + 3]], nvvm.MMALayout.ROW, ) @@ -1339,7 +1245,7 @@ def compute0_warp_group( bars.mb_t_inv_ready[tinv1_idx].arrive() bars.mb_beta_done[beta1_idx].arrive() - tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one) for _ in range(cfg.smem_t_inv_stages): bars.mb_t_inv_done[tinv_index.idx].wait(tinv_index.phase) tinv_index = advance(tinv_index, cfg.smem_t_inv_stages) @@ -1356,7 +1262,7 @@ def compute1_warp_group( tidx, warp_idx, tmem_hold, - sV, + sV_trans, sCumsumlog, sCumprod, sBeta, @@ -1364,112 +1270,97 @@ def compute1_warp_group( mState_init, mState_out, checkpoint_every_n_tokens, - sSched, + sScheduler, bars, ): """Compute warp-group 1 role (warps 4-7): persistent scheduler loop owning the recurrent state from seed to final store.""" - elect_one = nvvm.elect_sync() + nvvm.setmaxregister(cfg.num_regs_compute_group_1, nvvm.SetMaxRegisterAction.INCREASE) v_index = PipelineState.start(phase=0) gate_index = PipelineState.start(phase=0) kv_acc_index = PipelineState.start(phase=0) k_state_ready_index = PipelineState.start(phase=0) u_acc_ready_index = PipelineState.start(phase=0) - state_acc_seed_index = PipelineState.start(phase=1) - state_inp_cnt = cutlass.Int32(0) - kv_done_idx = cutlass.Int32(0) - - nvvm.setmaxregister(cfg.num_regs_compute_group_1, nvvm.SetMaxRegisterAction.INCREASE) - nvvm.barrier_cta_sync_aligned( - cfg.tmem_alloc_barrier_id, - thread_count=cfg.tmem_alloc_barrier_threads, - ) - tmem_base = tmem_hold.load() num_threads_cg1 = cfg.threads_per_warp * len(cfg.compute_group_1_warp_ids) cg1_tidx = tidx % num_threads_cg1 - lane_id = cg1_tidx % cfg.threads_per_warp - tmem_warp_row = (cg1_tidx // cfg.threads_per_warp) * cfg.threads_per_warp + lane_idx = cg1_tidx % cfg.threads_per_warp + + elect_one = nvvm.elect_sync() + state_input_cnt = cutlass.Int32(0) ldtm_width = 32 - sttm_width = ldtm_width // 2 - num_state_subs = cutlass.const_expr(cfg.d_v // ldtm_width) - tmem_state_col = tmem_base + cfg.tmem_state_acc_offset - tmem_state_inp_col = tmem_base + cfg.tmem_state_inp_offset - tmem_inp_col = tmem_base + cfg.tmem_y_decay_u_inp_offset + sttm_width = 16 + num_ldtms = cutlass.const_expr(cfg.d_v // ldtm_width) INP_SLOT_COLS = cfg.b_t // 2 - tmem_k_state_col = tmem_base + cfg.tmem_cg1_acc_offset - tmem_u_acc_col = tmem_k_state_col - tmem_y_inp_col = tmem_inp_col - tmem_decay_v_col = tmem_inp_col + INP_SLOT_COLS - v_frag_tok = cg1_tidx % 8 + (cg1_tidx // 16 % 2) * 8 - v_frag_col = (cg1_tidx // 8 % 2) * 8 + (cg1_tidx // 32 % 2) * 32 - v_frag_slab = (cg1_tidx // 64) * 4096 - v_stage_elems = cfg.v_cosize // cfg.smem_v_stages - sV_base = cute.make_ptr(cfg.io_dtype, sV[0].base, mem_space=cute.AddressSpace.smem, assumed_align=cfg.buffer_align_bytes) + v_row = cg1_tidx % 8 + (cg1_tidx // 16 % 2) * 8 + v_col = (cg1_tidx // 8 % 2) * 8 + (cg1_tidx // 32 % 2) * 32 + v_segment = (cg1_tidx // 64) * 4096 + v_stage_elements = cfg.v_cosize // cfg.smem_v_stages + sV_base = sV_trans[0].base num_vals = 32 + + nvvm.barrier_cta_sync_aligned(cfg.tmem_lifecycle_barrier_id, thread_count=cfg.tmem_user_threads) + tmem_base = tmem_hold.load() + tmem_col = tmem_base & 0xFFFF + tmem_row = tmem_base >> 16 + row_lo_addr = tmem_row << 16 + row_hi_addr = (tmem_row + 16) << 16 + tmem_state_col = tmem_col + cfg.tmem_state_acc_offset + tmem_state_input_col = tmem_col + cfg.tmem_state_input_offset + tmem_input_col = tmem_col + cfg.tmem_y_decay_u_input_offset + tmem_k_state_col = tmem_col + cfg.tmem_cg1_acc_offset + tmem_u_acc_col = tmem_k_state_col + tmem_y_input_col = tmem_input_col + tmem_decay_v_col = tmem_input_col + INP_SLOT_COLS if cutlass.const_expr(cfg.enable_checkpoints): - sCheckpoint_base_int = sCheckpoint_raw.data_ptr().toint() + sCheckpoint_base = sCheckpoint_raw.data_ptr() checkpoint_cnt = cutlass.Int32(0) checkpoint_frag_row = cg1_tidx % 8 + (cg1_tidx // 16 % 2) * 8 checkpoint_frag_col = (cg1_tidx // 8 % 2) * 8 + (cg1_tidx // 32 % 2) * 32 - sched_state = PipelineState.start(phase=0) + scheduler_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) - n_local = wend - cstart + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) + n_local = write_end - compute_start if cutlass.const_expr(cfg.enable_checkpoints): - ckpt_chunks = checkpoint_every_n_tokens // cutlass.Int32(cfg.b_t) - checkpoint_mod = cstart % ckpt_chunks + checkpoint_chunks = checkpoint_every_n_tokens // cutlass.Int32(cfg.b_t) + checkpoint_mod = compute_start % checkpoint_chunks if n_local > 0: if cutlass.const_expr(cfg.use_initial_state): - # ---- initial-state seed: initial_state GMEM -> state TMEM --------------- gState_init = mState_init[None, None, head_idx, batch_idx] - kv_init_idx = state_acc_seed_index.idx - bars.mb_state_acc_scale_done[kv_init_idx].wait(state_acc_seed_index.phase) - state_acc_seed_index = advance(state_acc_seed_index, cfg.tmem_state_acc_stages) - seed_from_initial_state = cstart == 0 - if seed_from_initial_state: - for sub in cutlass.range_constexpr(num_state_subs): + seed_state = compute_start == 0 + if seed_state: + for i in cutlass.range_constexpr(num_ldtms): words = [] for k in cutlass.range_constexpr(32): - v = gState_init[cg1_tidx, sub * ldtm_width + k] + v = gState_init[cg1_tidx, i * ldtm_width + k] if cutlass.const_expr(cfg.state_dtype != cfg.acc_dtype): v = v.to(cfg.acc_dtype) words.append(v) nvvm.tcgen05_st( "32x32b", - nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_state_col + sub * ldtm_width, cutlass.Float32), + nvvm.make_tmem_ptr(row_lo_addr + tmem_state_col + i * ldtm_width, cutlass.Float32), cutlass.Vector.from_elements(tuple(words), cutlass.Float32), ) - else: - for sub in cutlass.range_constexpr(num_state_subs): - nvvm.tcgen05_st( - "32x32b", - nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_state_col + sub * ldtm_width, cutlass.Float32), - cutlass.Vector.from_elements(tuple(cutlass.Float32(0.0) for _ in range(32)), cutlass.Float32), - ) - nvvm.tcgen05_wait("store") - - nvvm.barrier_cta_sync_aligned( - cfg.init_state_store_barrier_id, - thread_count=cfg.init_state_store_barrier_threads, - ) + nvvm.tcgen05_wait("store") for local_idx in cutlass.range(n_local): # noqa: B007 - chunk_idx = cstart + local_idx + chunk_idx = compute_start + local_idx if cutlass.const_expr(cfg.enable_checkpoints): do_checkpoint_now = checkpoint_mod == 0 checkpoint_mod = checkpoint_mod + cutlass.Int32(1) - checkpoint_mod = cutlass.Int32(0) if checkpoint_mod == ckpt_chunks else checkpoint_mod + checkpoint_mod = cutlass.Int32(0) if checkpoint_mod == checkpoint_chunks else checkpoint_mod if cutlass.const_expr(cfg.enable_checkpoints and not cfg.use_initial_state): - if chunk_idx == 0 and wstart == 0: + if chunk_idx == 0 and write_start == 0: checkpoint_stage = checkpoint_cnt % cfg.smem_checkpoint_stages checkpoint_phase_done = cutlass.Int32(1) ^ ((checkpoint_cnt // cfg.smem_checkpoint_stages) & cutlass.Int32(1)) bars.mb_checkpoint_tmastg_done[checkpoint_stage].wait(checkpoint_phase_done) checkpoint_zero_ptr = cutlass.inttoptr( - sCheckpoint_base_int + checkpoint_stage * cfg.d_k * cfg.d_v * 2, cutlass.AddressSpace.smem, cutlass.Int32 + (sCheckpoint_base + checkpoint_stage * cfg.d_k * cfg.d_v).toint(), cutlass.AddressSpace.smem, cutlass.Int32 ) for z in cutlass.range_constexpr(cfg.d_k * cfg.d_v // 2 // num_threads_cg1): (checkpoint_zero_ptr + cg1_tidx + z * num_threads_cg1).store(cutlass.Int32(0)) @@ -1479,61 +1370,51 @@ def compute1_warp_group( checkpoint_cnt = checkpoint_cnt + 1 valid_state = local_idx > 0 if cutlass.const_expr(cfg.use_initial_state): - valid_state = cutlass.Boolean(True) - state_acc_seed_index = advance(state_acc_seed_index, cfg.tmem_state_acc_stages) + valid_state = local_idx > 0 or seed_state gate_idx = gate_index.idx bars.mb_gate_ready[gate_idx].wait(gate_index.phase) gate_index = advance(gate_index, cfg.smem_gate_stages) cumprod_total = sCumprod[sCumprod.shape[0] - 1, 0, gate_idx] - # ---- state restage + rescale ------------------------------------- + # ---- state stage + rescale ------------------------------------------- if valid_state: - kv_idx = kv_acc_index.idx - bars.mb_state_acc_ready[kv_idx].wait(kv_acc_index.phase) - kv_acc_index = advance(kv_acc_index, cfg.tmem_state_acc_stages) - kv_done_idx = kv_idx - - state_regs = [[cutlass.Float32(0.0) for _ in range(num_state_subs)] for _ in range(32)] - state_inp_stage_idx = state_inp_cnt % cfg.tmem_state_inp_stages - state_vecs = [] - for sub in cutlass.range_constexpr(num_state_subs): - state_vecs.append( - nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_state_col + sub * ldtm_width, cutlass.Float32), num=32) - ) - for sub in cutlass.range_constexpr(num_state_subs): - for k in cutlass.range_constexpr(32): - state_regs[k][sub] = state_vecs[sub][k] - state_pack = [fp32_to_fp16(state_regs[2 * j][sub], state_regs[2 * j + 1][sub], dtype=cfg.io_dtype) for j in range(16)] + if local_idx > 0: + bars.mb_state_acc_ready[kv_acc_index.idx].wait(kv_acc_index.phase) + kv_acc_index = advance(kv_acc_index, cfg.tmem_state_acc_stages) + + state_input_stage_idx = state_input_cnt % cfg.tmem_state_input_stages + state_vecs = [ + nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_lo_addr + tmem_state_col + i * ldtm_width, cutlass.Float32), num=32) + for i in range(num_ldtms) + ] + state_regs = [[state_vecs[i][k] for i in range(num_ldtms)] for k in range(32)] + for i in cutlass.range_constexpr(num_ldtms): + state_pack = [fp32_to_fp16(state_regs[2 * j][i], state_regs[2 * j + 1][i], dtype=cfg.io_dtype) for j in range(16)] nvvm.tcgen05_st( "32x32b", - nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_state_inp_col + sub * sttm_width, cutlass.Int32), + nvvm.make_tmem_ptr(row_lo_addr + tmem_state_input_col + i * sttm_width, cutlass.Int32), cutlass.Vector.from_elements(tuple(state_pack), cutlass.Int32), ) nvvm.tcgen05_wait("store") - bars.mb_state_inp_ready[state_inp_stage_idx].arrive() - state_inp_cnt = state_inp_cnt + 1 + bars.mb_state_input_ready[state_input_stage_idx].arrive() + state_input_cnt = state_input_cnt + 1 if cutlass.const_expr(cfg.enable_checkpoints): # ---- state checkpoint ---------------------------------------- - do_checkpoint = do_checkpoint_now and chunk_idx < wend - do_checkpoint = do_checkpoint and chunk_idx >= wstart + do_checkpoint = do_checkpoint_now and chunk_idx < write_end + do_checkpoint = do_checkpoint and chunk_idx >= write_start if do_checkpoint: checkpoint_stage = checkpoint_cnt % cfg.smem_checkpoint_stages checkpoint_phase_done = cutlass.Int32(1) ^ ((checkpoint_cnt // cfg.smem_checkpoint_stages) & cutlass.Int32(1)) bars.mb_checkpoint_tmastg_done[checkpoint_stage].wait(checkpoint_phase_done) checkpoint_stage_base = checkpoint_stage * cfg.d_k * cfg.d_v - for sub in cutlass.range_constexpr(num_state_subs): - checkpoint_vec = nvvm.tcgen05_ld( - "32x32b", - nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_state_col + sub * ldtm_width, cutlass.Float32), - num=32, - ) + for i in cutlass.range_constexpr(num_ldtms): for g in cutlass.range_constexpr(ldtm_width // 8): packs = tuple( - fp32_to_fp16(checkpoint_vec[g * 8 + 2 * t], checkpoint_vec[g * 8 + 2 * t + 1], dtype=cfg.io_dtype) for t in range(4) + fp32_to_fp16(state_regs[g * 8 + 2 * t][i], state_regs[g * 8 + 2 * t + 1][i], dtype=cfg.io_dtype) for t in range(4) ) - dk = sub * ldtm_width + g * 8 + dk = i * ldtm_width + g * 8 checkpoint_addr = checkpoint_stage_base + (dk // 64) * (cfg.d_v * 64) + cg1_tidx * 64 + swizzle_xor_128b(cg1_tidx, dk % 64) (sCheckpoint_raw.data_ptr() + checkpoint_addr).store( cutlass.Vector.from_elements(packs, cutlass.Int32).bitcast(cfg.io_dtype), alignment=16 @@ -1543,27 +1424,26 @@ def compute1_warp_group( bars.mb_checkpoint_tmastg_ready[checkpoint_stage].arrive() checkpoint_cnt = checkpoint_cnt + 1 - for sub in cutlass.range_constexpr(num_state_subs): + for i in cutlass.range_constexpr(num_ldtms): state_scaled = [] for j in cutlass.range_constexpr(16): - s0, s1 = fmul2(state_regs[2 * j][sub], state_regs[2 * j + 1][sub], cumprod_total, cumprod_total) + s0, s1 = fmul2(state_regs[2 * j][i], state_regs[2 * j + 1][i], cumprod_total, cumprod_total) state_scaled += [s0, s1] nvvm.tcgen05_st( "32x32b", - nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_state_col + sub * ldtm_width, cutlass.Float32), + nvvm.make_tmem_ptr(row_lo_addr + tmem_state_col + i * ldtm_width, cutlass.Float32), cutlass.Vector.from_elements(tuple(state_scaled), cutlass.Float32), ) nvvm.tcgen05_wait("store") - bars.mb_state_acc_scale_done[kv_idx].arrive() - # ---- per-row Gate register builds -------------------------------- + # ---- per-row Gate register builds ------------------------------------ cumprod_vals = [] for k in cutlass.range_constexpr(num_vals): - cumprod_vals.append(sCumprod[(lane_id % 4) * 2 + ((k // 4) * 8 + k % 2), 0, gate_idx]) + cumprod_vals.append(sCumprod[(lane_idx % 4) * 2 + ((k // 4) * 8 + k % 2), 0, gate_idx]) last_cumsumlog = sCumsumlog[cfg.b_t - 1, 0, gate_idx] cumsumlog_vals = [] for k in cutlass.range_constexpr(num_vals): - cumsumlog_vals.append(sCumsumlog[(lane_id % 4) * 2 + ((k // 4) * 8 + k % 2), 0, gate_idx]) + cumsumlog_vals.append(sCumsumlog[(lane_idx % 4) * 2 + ((k // 4) * 8 + k % 2), 0, gate_idx]) decay_scale_vals = [] for k in cutlass.range_constexpr(0, num_vals, 2): d0, d1 = fadd2(last_cumsumlog, last_cumsumlog, -cumsumlog_vals[k], -cumsumlog_vals[k + 1]) @@ -1571,119 +1451,112 @@ def compute1_warp_group( decay_scale_vals.append(cute.math.exp2(d1, fastmath=True)) bars.mb_gate_done[gate_idx].arrive() - # ---- Y = V - K*state (packed 16-bit) ----------------------------- + # ---- Y = V - k state (packed 16-bit) --------------------------------- v_idx = v_index.idx bars.mb_v_ready[v_idx].wait(v_index.phase) v_index = advance(v_index, cfg.smem_v_stages) - v_frags = [[cutlass.Int32(0), cutlass.Int32(0)] for _ in range(16)] - for c in cutlass.range_constexpr(8): - tok_block = cutlass.const_expr(c % 4) - sub = cutlass.const_expr(c // 4) - v_frag = nvvm.ldmatrix( - ( - sV_base - + v_idx * v_stage_elems - + v_frag_slab - + (v_frag_tok + tok_block * 16) * 64 - + swizzle_xor_128b(v_frag_tok + tok_block * 16, v_frag_col + sub * 16) - ).raw_ptr(), - 4, - nvvm.MMALayout.COL, - ) - for i in cutlass.range_constexpr(4): - v_frags[4 * tok_block + i][sub] = v_frag[i] + v_frag = [] + for half in cutlass.range_constexpr(2): + v_words = [] + for block in cutlass.range_constexpr(4): + v_raw = nvvm.ldmatrix( + ( + sV_base + + v_idx * v_stage_elements + + v_segment + + (v_row + block * 16) * 64 + + swizzle_xor_128b(v_row + block * 16, v_col + half * 16) + ), + 4, + nvvm.MMALayout.COL, + ) + for i in cutlass.range_constexpr(4): + v_words.append(v_raw[i]) + v_frag.append(v_words) if valid_state: bars.mb_k_state_acc_ready[0].wait(k_state_ready_index.phase) k_state_ready_index = advance(k_state_ready_index, 1) - for sub in cutlass.range_constexpr(2): + for half in cutlass.range_constexpr(2): k_state_vec = nvvm.tcgen05_ld( "16x256b", - nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_k_state_col, cutlass.Float32), + nvvm.make_tmem_ptr(((tmem_row + half * 16) << 16) + tmem_k_state_col, cutlass.Float32), num=8, ) for j in cutlass.range_constexpr(16): s0, s1 = fmul2(k_state_vec[2 * j], k_state_vec[2 * j + 1], cumprod_vals[2 * j], cumprod_vals[2 * j + 1]) k_state_pack = fp32_to_fp16(s0, s1, dtype=cfg.io_dtype) - v_frags[j][sub] = sub_f16x2(v_frags[j][sub], k_state_pack, cfg.io_dtype) - for sub in cutlass.range_constexpr(2): + v_frag[half][j] = sub_f16x2(v_frag[half][j], k_state_pack, cfg.io_dtype) + for half in cutlass.range_constexpr(2): nvvm.tcgen05_st( "16x128b", - nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_y_inp_col, cutlass.Int32), - cutlass.Vector.from_elements(tuple(v_frags[j][sub] for j in range(16)), cutlass.Int32), + nvvm.make_tmem_ptr(((tmem_row + half * 16) << 16) + tmem_y_input_col, cutlass.Int32), + cutlass.Vector.from_elements(tuple(v_frag[half]), cutlass.Int32), ) nvvm.tcgen05_wait("store") - bars.mb_y_inp_ready[0].arrive() + bars.mb_y_input_ready[0].arrive() - # ---- U epilogue + decayed-U publish ------------------------------ + # ---- U epilogue + decayed-U publish ---------------------------------- bars.mb_u_acc_ready[0].wait(u_acc_ready_index.phase) u_acc_ready_index = advance(u_acc_ready_index, 1) bars.mb_v_done[v_idx].arrive() - u_acc_regs = [[cutlass.Float32(0.0), cutlass.Float32(0.0)] for _ in range(32)] u_acc_vecs = [] - for sub in cutlass.range_constexpr(2): + for half in cutlass.range_constexpr(2): u_acc_vecs.append( nvvm.tcgen05_ld( "16x256b", - nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_u_acc_col, cutlass.Float32), + nvvm.make_tmem_ptr(((tmem_row + half * 16) << 16) + tmem_u_acc_col, cutlass.Float32), num=8, ) ) - for sub in cutlass.range_constexpr(2): - for k in cutlass.range_constexpr(32): - u_acc_regs[k][sub] = u_acc_vecs[sub][k] + u_acc_regs = [[u_acc_vecs[0][k], u_acc_vecs[1][k]] for k in range(32)] - for sub in cutlass.range_constexpr(2): + for half in cutlass.range_constexpr(2): for j in cutlass.range_constexpr(16): - u_acc_regs[2 * j][sub], u_acc_regs[2 * j + 1][sub] = fmul2( - u_acc_regs[2 * j][sub], u_acc_regs[2 * j + 1][sub], decay_scale_vals[2 * j], decay_scale_vals[2 * j + 1] + u_acc_regs[2 * j][half], u_acc_regs[2 * j + 1][half] = fmul2( + u_acc_regs[2 * j][half], u_acc_regs[2 * j + 1][half], decay_scale_vals[2 * j], decay_scale_vals[2 * j + 1] ) - decay_pack = [fp32_to_fp16(u_acc_regs[2 * j][sub], u_acc_regs[2 * j + 1][sub], dtype=cfg.io_dtype) for j in range(16)] + decay_pack = [fp32_to_fp16(u_acc_regs[2 * j][half], u_acc_regs[2 * j + 1][half], dtype=cfg.io_dtype) for j in range(16)] nvvm.tcgen05_st( "16x128b", - nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_decay_v_col, cutlass.Int32), + nvvm.make_tmem_ptr(((tmem_row + half * 16) << 16) + tmem_decay_v_col, cutlass.Int32), cutlass.Vector.from_elements(tuple(decay_pack), cutlass.Int32), ) nvvm.tcgen05_wait("store") - bars.mb_decay_u_inp_ready[0].arrive() + bars.mb_decay_u_input_ready[0].arrive() - # ---- final state: state TMEM -> GMEM ----------------------------------- + # ---- final state store: TMEM -> GMEM ----------------------------------------- if n_local > 0: kv_last_idx = kv_acc_index.idx bars.mb_state_acc_ready[kv_last_idx].wait(kv_acc_index.phase) kv_acc_index = advance(kv_acc_index, cfg.tmem_state_acc_stages) if cutlass.const_expr(cfg.store_final_state): - if wend == num_chunks_b: + if write_end == batch_num_chunks: gState_out = mState_out[None, None, head_idx, batch_idx] - for sub in cutlass.range_constexpr(num_state_subs): - state_vec = nvvm.tcgen05_ld( - "32x32b", nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_state_col + sub * ldtm_width, cutlass.Float32), num=32 - ) + for i in cutlass.range_constexpr(num_ldtms): + state_vec = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_lo_addr + tmem_state_col + i * ldtm_width, cutlass.Float32), num=32) for k in cutlass.range_constexpr(32): val = state_vec[k] if cutlass.const_expr(cfg.state_dtype != cfg.acc_dtype): val = val.to(cfg.state_dtype) - gState_out[cg1_tidx, sub * ldtm_width + k] = val - bars.mb_state_acc_scale_done[kv_last_idx].arrive() - else: - bars.mb_state_acc_scale_done[kv_last_idx].arrive() + gState_out[cg1_tidx, i * ldtm_width + k] = val else: if cutlass.const_expr(cfg.store_final_state): - write_passthrough = wend == num_chunks_b + write_passthrough = write_end == batch_num_chunks if write_passthrough: gState_out = mState_out[None, None, head_idx, batch_idx] if cutlass.const_expr(cfg.use_initial_state): gState_in = mState_init[None, None, head_idx, batch_idx] - for r in cutlass.range(num_state_subs * ldtm_width): + for r in cutlass.range(num_ldtms * ldtm_width): gState_out[cg1_tidx, r] = gState_in[cg1_tidx, r] else: - for sub in cutlass.range_constexpr(num_state_subs): + for i in cutlass.range_constexpr(num_ldtms): for k in cutlass.range_constexpr(32): - gState_out[cg1_tidx, sub * ldtm_width + k] = cutlass.Float32(0.0).to(cfg.state_dtype) + gState_out[cg1_tidx, i * ldtm_width + k] = cutlass.Float32(0.0).to(cfg.state_dtype) - tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one) bars.mb_tmem_done[0].arrive() @@ -1701,7 +1574,7 @@ def build_descs_body( base_k, base_v, base_checkpoint, - desc_ws: cute.Tensor, + desc_workspace: cute.Tensor, cu_seqlens: cute.Tensor, k: cute.Tensor, v: cute.Tensor, @@ -1716,9 +1589,9 @@ def build_descs_body( prologue kernel after its order pass; warps past the array count fall through the widx guards.""" arr_words = n_batch * cutlass.Int32(TENSOR_MAP_QWORDS) - desc_k_arr = cute.make_tensor(desc_ws.iterator, cute.make_layout((arr_words,), stride=(1,))) - desc_v_arr = cute.make_tensor(desc_ws.iterator + arr_words, cute.make_layout((arr_words,), stride=(1,))) - desc_checkpoint_arr = cute.make_tensor(desc_ws.iterator + 2 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + desc_k_arr = cute.make_tensor(desc_workspace.iterator, cute.make_layout((arr_words,), stride=(1,))) + desc_v_arr = cute.make_tensor(desc_workspace.iterator + arr_words, cute.make_layout((arr_words,), stride=(1,))) + desc_checkpoint_arr = cute.make_tensor(desc_workspace.iterator + 2 * arr_words, cute.make_layout((arr_words,), stride=(1,))) if widx == 0: if nvvm.elect_sync(): @@ -1745,7 +1618,7 @@ def prologue_kernel( base_k: cutlass.GridConstant[tma.TensorMap], base_v: cutlass.GridConstant[tma.TensorMap], base_checkpoint: cutlass.GridConstant[tma.TensorMap], - desc_ws: cute.Tensor, + desc_workspace: cute.Tensor, cu_seqlens: cute.Tensor, k: cute.Tensor, v: cute.Tensor, @@ -1754,7 +1627,7 @@ def prologue_kernel( mStaging: Optional[cute.Tensor], mCount: cute.Tensor, mWorkItems: cute.Tensor, - mSched: Optional[cute.Tensor], + mScheduler: Optional[cute.Tensor], n_batch: cutlass.Int32, k_row_stride: cutlass.Int32, v_row_stride: cutlass.Int32, @@ -1763,7 +1636,7 @@ def prologue_kernel( ) -> None: """Single-CTA prologue. Under ``run_order`` this kernel is the first work-item-table consumer, so it LPT-orders the table and zeroes both - consumers' sched rings via :func:`order_body`; it then builds the + consumers' scheduler rings via :func:`order_body`; it then builds the per-batch TMA-descriptor arrays via :func:`build_descs_body`, one warp per array (the extra warps only take part in the order phase).""" tidx, _, _ = cute.arch.thread_idx() @@ -1779,7 +1652,7 @@ def prologue_kernel( True, b_t, ORDER_THREADS, - ORDER_ELEMS, + ORDER_ELEMENTS, tidx, n_heads_out, n_heads_out * n_batch, @@ -1787,7 +1660,7 @@ def prologue_kernel( mStaging, mCount, mWorkItems, - mSched, + mScheduler, sKey, sIdx, sSpread, @@ -1797,7 +1670,7 @@ def prologue_kernel( base_k, base_v, base_checkpoint, - desc_ws, + desc_workspace, cu_seqlens, k, v, @@ -1824,7 +1697,7 @@ def prologue( work_item_staging: Optional[cute.Tensor], work_count: cute.Tensor, work_items: cute.Tensor, - sched_all: Optional[cute.Tensor], + scheduler_all: Optional[cute.Tensor], checkpoint_every_n: cutlass.Int32, tensormap_workspace: cute.Tensor, stream: cuda.CUstream, @@ -1838,7 +1711,7 @@ def prologue( batch_size = cu_seqlens.shape[0] - 1 d_v = v.shape[2] bpe = io_dtype.width // 8 - granu = 128 // bpe + granule = 128 // bpe bt = b_t k_row_stride, k_head_stride = k.stride[0], k.stride[1] @@ -1849,14 +1722,14 @@ def prologue( k_headed = cute.make_tensor(k.iterator, cute.make_layout((seqlen, h_k, d_k), stride=(k_row_stride, k_head_stride, 1))) v_headed = cute.make_tensor(v.iterator, cute.make_layout((d_v, h_v, seqlen), stride=(1, v_head_stride, v_row_stride))) swz128 = tma.TensorMapSwizzle.s128b - base_desc_k = tma.create_tensor_map_tiled_from_view(k_headed, box_dims=(bt, 1, granu), stride_order=(2, 1, 0), swizzle=swz128) - base_desc_v = tma.create_tensor_map_tiled_from_view(v_headed, box_dims=(granu, 1, bt), stride_order=(0, 1, 2), swizzle=swz128) + base_desc_k = tma.create_tensor_map_tiled_from_view(k_headed, box_dims=(bt, 1, granule), stride_order=(2, 1, 0), swizzle=swz128) + base_desc_v = tma.create_tensor_map_tiled_from_view(v_headed, box_dims=(granule, 1, bt), stride_order=(0, 1, 2), swizzle=swz128) base_desc_checkpoint = base_desc_v if cutlass.const_expr(state_checkpoints_out is not None): d_k_state = state_checkpoints_out.shape[2] d_v_state = state_checkpoints_out.shape[3] - checkpoint_granu = 128 // (state_checkpoints_out.element_type.width // 8) + checkpoint_granule = 128 // (state_checkpoints_out.element_type.width // 8) checkpoint_view = cute.make_tensor( state_checkpoints_out.iterator, cute.make_layout( @@ -1865,7 +1738,7 @@ def prologue( ), ) base_desc_checkpoint = tma.create_tensor_map_tiled_from_view( - checkpoint_view, box_dims=(checkpoint_granu, d_k_state, 1, 1), stride_order=(0, 1, 2, 3), swizzle=swz128 + checkpoint_view, box_dims=(checkpoint_granule, d_k_state, 1, 1), stride_order=(0, 1, 2, 3), swizzle=swz128 ) prologue_kernel( @@ -1884,7 +1757,7 @@ def prologue( work_item_staging, work_count, work_items, - sched_all, + scheduler_all, cutlass.Int32(batch_size), cutlass.Int32(k_row_stride), cutlass.Int32(v_row_stride), @@ -1907,17 +1780,17 @@ def host( state_out: Optional[cute.Tensor], work_items: Optional[cute.Tensor], work_count: Optional[cute.Tensor], - sched_ctr: Optional[cute.Tensor], + scheduler_counter: Optional[cute.Tensor], checkpoint_every_n_tokens: cutlass.Int32, tensormap_workspace: cute.Tensor, stream: cuda.CUstream, ): - h_k = k.shape[1] - h_v = v.shape[1] + h_k = cfg.h_k + h_v = cfg.h_v batch_size = cu_seqlens.shape[0] - 1 - heads_out = gate.shape[1] + heads_out = cfg.n_heads_out - # ---- GQA reshapes: fold the head group into a -------------------------------- + # ---- GQA reshapes: fold the head group into the Q head axis ---------------------- if cutlass.const_expr(cfg.is_GQA): h_ratio = heads_out // h_v h_native = h_v @@ -1994,28 +1867,28 @@ def host( ), ) - # ---- SMEM sizing: per-buffer element cosizes --------------------------------- + # ---- SMEM sizing: per-buffer element cosizes ------------------------------------- bpe = cfg.io_dtype.width // 8 - kq_tile_elems = 2 * cfg.b_t * cfg.d_k - v_tile_elems = cfg.d_v * cfg.b_t - tinv_tile_elems = cfg.b_t * cfg.b_t - cfg.kq_cosize = kq_tile_elems * cfg.smem_kq_stages - cfg.v_cosize = v_tile_elems * cfg.smem_v_stages - cfg.t_inv_cosize = tinv_tile_elems * cfg.smem_t_inv_stages + kq_tile_elements = 2 * cfg.b_t * cfg.d_k + v_tile_elements = cfg.d_v * cfg.b_t + tinv_tile_elements = cfg.b_t * cfg.b_t + cfg.kq_cosize = kq_tile_elements * cfg.smem_kq_stages + cfg.v_cosize = v_tile_elements * cfg.smem_v_stages + cfg.t_inv_cosize = tinv_tile_elements * cfg.smem_t_inv_stages cfg.checkpoint_cosize = cfg.d_k * cfg.d_v * cfg.smem_checkpoint_stages cumsumlog_smem_layout_staged = cute.make_layout((cfg.b_t, 1, cfg.smem_gate_stages)) beta_smem_layout_staged = cute.make_layout((cfg.b_t, 1, cfg.smem_beta_stages)) - cfg.tma_kq_bytes = (kq_tile_elems // 2) * bpe - cfg.tma_v_bytes = v_tile_elems * bpe + cfg.tma_kq_bytes = (kq_tile_elements // 2) * bpe + cfg.tma_v_bytes = v_tile_elements * bpe cfg.n_heads_out = heads_out cfg.k_ratio = heads_out // h_k cfg.v_ratio = heads_out // h_v num_descs = batch_size - # ---- launch ------------------------------------------------------------------ + # ---- launch ---------------------------------------------------------------------- grid_shape = (cfg.max_active_clusters, 1, 1) kernel( @@ -2029,7 +1902,7 @@ def host( state_out, work_items, work_count, - sched_ctr, + scheduler_counter, checkpoint_every_n_tokens, cumsumlog_smem_layout_staged, beta_smem_layout_staged, @@ -2058,7 +1931,7 @@ def kernel( mState_out: Optional[cute.Tensor], mWorkItems: cute.Tensor, mCount: cute.Tensor, - mSched: Optional[cute.Tensor], + mScheduler: Optional[cute.Tensor], checkpoint_every_n_tokens: cutlass.Int32, cumsumlog_smem_layout_staged: cute.Layout, beta_smem_layout_staged: cute.Layout, @@ -2075,8 +1948,8 @@ def kernel( num_ctas = cute.arch.grid_dim()[0] total_tiles = mCount[0] - if cutlass.const_expr(cfg.dyn_sched): - assert mSched is not None, "mSched must be provided if dyn_sched is True" + if cutlass.const_expr(cfg.dynamic_scheduling): + assert mScheduler is not None, "mScheduler must be provided if dynamic_scheduling is True" if cutlass.const_expr(cfg.use_initial_state): assert mState_init is not None, "mState_init must be provided if use_initial_state is True" @@ -2096,7 +1969,6 @@ def kernel( SMEM = cutlass.AddressSpace.smem - bpe = cfg.io_dtype.width // 8 SWZ = 2 LEAD = 16 STRIDE = 8 * 128 @@ -2118,16 +1990,16 @@ def kernel( alignment=cfg.buffer_align_bytes, ) sKQ = SmemTile( - base=sKQ_raw.data_ptr().toint(), - elems_per_stage=(cfg.kq_cosize // cfg.smem_kq_stages) * bpe, + base=sKQ_raw.data_ptr(), + elems_per_stage=(cfg.kq_cosize // cfg.smem_kq_stages), stages=cfg.smem_kq_stages, leading_byte_offset=LEAD, stride_byte_offset=STRIDE, layout=SWZ, ) sKQ_trans = SmemTile( - base=sKQ_raw.data_ptr().toint(), - elems_per_stage=(cfg.kq_cosize // cfg.smem_kq_stages) * bpe, + base=sKQ_raw.data_ptr(), + elems_per_stage=(cfg.kq_cosize // cfg.smem_kq_stages), stages=cfg.smem_kq_stages, leading_byte_offset=2 * KT_LEAD, stride_byte_offset=STRIDE, @@ -2135,7 +2007,7 @@ def kernel( ) bars = make_gdn_bars(cfg) tmem_hold = cutlass.Array(cutlass.Int32, 1, space=SMEM, alignment=16) - sSched = cutlass.Array(cutlass.Int32, cfg.sched_stages, space=SMEM, alignment=16) + sScheduler = cutlass.Array(cutlass.Int32, cfg.scheduler_stages, space=SMEM, alignment=16) cumsumlog_raw = cutlass.Array(cutlass.Float32, cute.cosize(cumsumlog_smem_layout_staged), space=SMEM, alignment=128) cumprod_raw = cutlass.Array(cutlass.Float32, cute.cosize(cumsumlog_smem_layout_staged), space=SMEM, alignment=128) beta_raw = cutlass.Array(cutlass.Float32, cute.cosize(beta_smem_layout_staged), space=SMEM, alignment=128) @@ -2146,8 +2018,8 @@ def kernel( alignment=cfg.buffer_align_bytes, ) sTinv = SmemTile( - base=sTinv_raw.data_ptr().toint(), - elems_per_stage=(cfg.t_inv_cosize // cfg.smem_t_inv_stages) * bpe, + base=sTinv_raw.data_ptr(), + elems_per_stage=(cfg.t_inv_cosize // cfg.smem_t_inv_stages), stages=cfg.smem_t_inv_stages, leading_byte_offset=LEAD, stride_byte_offset=STRIDE, @@ -2159,9 +2031,9 @@ def kernel( space=cutlass.AddressSpace.smem, alignment=cfg.buffer_align_bytes, ) - sV = SmemTile( - base=sV_raw.data_ptr().toint(), - elems_per_stage=(cfg.v_cosize // cfg.smem_v_stages) * bpe, + sV_trans = SmemTile( + base=sV_raw.data_ptr(), + elems_per_stage=(cfg.v_cosize // cfg.smem_v_stages), stages=cfg.smem_v_stages, leading_byte_offset=V_LEAD, stride_byte_offset=STRIDE, @@ -2180,7 +2052,7 @@ def kernel( beta_smem_layout_staged, ) - # ---- mbarrier init (all threads) --------------------------------------------- + # ---- mbarrier init (all threads) ------------------------------------------------- for s in range(cfg.smem_kq_stages): bars.mb_kq_ready[s].init() bars.mb_kq_done[s].init() @@ -2195,7 +2067,6 @@ def kernel( bars.mb_beta_done[s].init() for s in range(cfg.tmem_state_acc_stages): bars.mb_state_acc_ready[s].init() - bars.mb_state_acc_scale_done[s].init() for s in range(cfg.tmem_cg0_acc_stages): bars.mb_cg0_acc_ready[s].init() bars.mb_cg0_acc_done[s].init() @@ -2204,22 +2075,22 @@ def kernel( for s in range(cfg.smem_t_inv_stages): bars.mb_t_inv_ready[s].init() bars.mb_t_inv_done[s].init() - for s in range(cfg.tmem_state_inp_stages): - bars.mb_state_inp_ready[s].init() - for b in (bars.mb_y_inp_ready, bars.mb_decay_u_inp_ready): + for s in range(cfg.tmem_state_input_stages): + bars.mb_state_input_ready[s].init() + for b in (bars.mb_y_input_ready, bars.mb_decay_u_input_ready): b[0].init() for s in range(cfg.smem_checkpoint_stages): bars.mb_checkpoint_tmastg_ready[s].init() bars.mb_checkpoint_tmastg_done[s].init() - for s_ in range(cfg.sched_stages): - bars.mb_sched_ready[s_].init() - bars.mb_sched_done[s_].init() + for s in range(cfg.scheduler_stages): + bars.mb_scheduler_ready[s].init() + bars.mb_scheduler_done[s].init() bars.mb_tmem_done[0].init() nvvm.fence_mbarrier_init() nvvm.barrier_cta_sync() - # ---- warp specialization ----------------------------------------------------- + # ---- warp specialization --------------------------------------------------------- if warp_idx >= cfg.compute_group_0_warp_ids[0] and warp_idx <= cfg.compute_group_0_warp_ids[-1]: compute0_warp_group( @@ -2236,7 +2107,7 @@ def kernel( sTinv=sTinv, sCheckpoint_raw=sCheckpoint_raw, checkpoint_every_n_tokens=checkpoint_every_n_tokens, - sSched=sSched, + sScheduler=sScheduler, bars=bars, ) @@ -2251,7 +2122,7 @@ def kernel( tidx, warp_idx=warp_idx, tmem_hold=tmem_hold, - sV=sV, + sV_trans=sV_trans, sCumsumlog=sCumsumlog, sCumprod=sCumprod, sBeta=sBeta, @@ -2259,7 +2130,7 @@ def kernel( mState_init=mState_init, mState_out=mState_out, checkpoint_every_n_tokens=checkpoint_every_n_tokens, - sSched=sSched, + sScheduler=sScheduler, bars=bars, ) @@ -2279,12 +2150,12 @@ def kernel( sCumsumlog=sCumsumlog, sCumprod=sCumprod, sBeta=sBeta, - sSched=sSched, + sScheduler=sScheduler, bars=bars, ) - elif warp_idx == cfg.mma_warp_id: - mma_warp( + elif warp_idx == cfg.tcgen05_mma_warp_id: + tcgen05_mma_warp( cfg, total_tiles, bidx, @@ -2295,7 +2166,7 @@ def kernel( sKQ=sKQ, sKQ_trans=sKQ_trans, sTinv=sTinv, - sSched=sSched, + sScheduler=sScheduler, bars=bars, ) @@ -2311,8 +2182,8 @@ def kernel( sV_raw=sV_raw, desc_k_base=desc_k_base, desc_v_base=desc_v_base, - mSched=mSched, - sSched=sSched, + mScheduler=mScheduler, + sScheduler=sScheduler, bars=bars, ) @@ -2328,7 +2199,7 @@ def kernel( tidx=tidx, sCheckpoint_raw=sCheckpoint_raw, desc_checkpoint_base=desc_checkpoint_base, - sSched=sSched, + sScheduler=sScheduler, bars=bars, ) @@ -2355,10 +2226,10 @@ class GdnRecomputeCfg: log_gate: bool = False safe_gate: bool = False beta_sigmoid: bool = False - dyn_sched: bool = False - sched_stages: int = CFG.SMEM_SCHED_STAGES + dynamic_scheduling: bool = False + scheduler_stages: int = CFG.SMEM_SCHEDULER_STAGES - # ---- fixed constants stamped from CFG by build_cfg --------------------------- + # ---- fixed constants stamped from CFG at build time ------------------------------ b_t: int = CFG.B_T d_k: int = CFG.D_K d_v: int = CFG.D_V @@ -2366,7 +2237,7 @@ class GdnRecomputeCfg: compute_group_1_warp_ids: Tuple[int, ...] = CFG.COMPUTE_GROUP_1_WARP_IDS load_gate_beta_warp_id: int = CFG.LOAD_GATE_BETA_WARP_ID tma_kv_warp_id: int = CFG.TMA_KV_WARP_ID - mma_warp_id: int = CFG.MMA_WARP_ID + tcgen05_mma_warp_id: int = CFG.TCGEN05_MMA_WARP_ID epilogue_warp_id: int = CFG.EPILOGUE_WARP_ID num_regs_compute_group_0: int = CFG.NUM_REGS_COMPUTE_GROUP_0 num_regs_compute_group_1: int = CFG.NUM_REGS_COMPUTE_GROUP_1 @@ -2375,15 +2246,13 @@ class GdnRecomputeCfg: threads_per_cta: int = 0 cluster_shape_mnk: Tuple[int, int, int] = CFG.CLUSTER_SHAPE_MNK - # ---- named barrier slots (ids 1-4; 0 is the CTA-wide sync) ------------------- - tmem_alloc_barrier_id: int = 1 - tmem_alloc_barrier_threads: int = 0 + # ---- named barrier slots (ids 1-4; 0 is the CTA-wide sync) ----------------------- + tmem_lifecycle_barrier_id: int = 1 + tmem_user_threads: int = 0 inverse_barrier_id: int = 2 inverse_barrier_threads: int = 0 - init_state_store_barrier_id: int = 4 - init_state_store_barrier_threads: int = 0 - # ---- SMEM / TMEM stage counts + TMEM column offsets -------------------------- + # ---- SMEM / TMEM stage counts + TMEM column offsets ------------------------------ smem_kq_stages: int = CFG.SMEM_KQ_STAGES smem_v_stages: int = CFG.SMEM_V_STAGES smem_t_inv_stages: int = CFG.SMEM_T_INV_STAGES @@ -2391,17 +2260,17 @@ class GdnRecomputeCfg: smem_gate_stages: int = CFG.SMEM_GATE_STAGES smem_beta_stages: int = CFG.SMEM_BETA_STAGES tmem_state_acc_stages: int = CFG.TMEM_KV_ACC_STAGES - tmem_state_inp_stages: int = CFG.TMEM_STATE_INP_STAGES + tmem_state_input_stages: int = CFG.TMEM_STATE_INP_STAGES tmem_cg0_acc_stages: int = CFG.TMEM_CG0_ACC_STAGES tmem_cg1_acc_stages: int = CFG.TMEM_CG1_ACC_STAGES tmem_state_acc_offset: int = 0 - tmem_state_inp_offset: int = 0 + tmem_state_input_offset: int = 0 tmem_cg0_acc_offset: int = 0 tmem_cg1_acc_offset: int = 0 - tmem_y_decay_u_inp_offset: int = 0 + tmem_y_decay_u_input_offset: int = 0 buffer_align_bytes: int = CFG.BUFFER_ALIGN_BYTES - # ---- stamped by host at trace time (shape-derived) -------------------------- + # ---- stamped by host at trace time (shape-derived) ------------------------------- kq_cosize: int = 0 v_cosize: int = 0 t_inv_cosize: int = 0 @@ -2425,7 +2294,7 @@ def build_cfg( log_gate: bool = False, safe_gate: bool = False, beta_sigmoid: bool = False, - dyn_sched: bool = False, + dynamic_scheduling: bool = False, ) -> GdnRecomputeCfg: """Build the per-compile ``GdnRecomputeCfg`` (io_dtype ∈ {Float16, BFloat16}; acc is always Float32).""" @@ -2443,7 +2312,7 @@ def build_cfg( log_gate=log_gate, safe_gate=safe_gate, beta_sigmoid=beta_sigmoid, - dyn_sched=dyn_sched, + dynamic_scheduling=dynamic_scheduling, ) cfg.smem_checkpoint_stages = 1 if enable_checkpoints: @@ -2454,14 +2323,13 @@ def build_cfg( n_cg0 = len(cfg.compute_group_0_warp_ids) n_cg1 = len(cfg.compute_group_1_warp_ids) cfg.threads_per_cta = cfg.threads_per_warp * (4 + n_cg0 + n_cg1) - cfg.tmem_alloc_barrier_threads = cfg.threads_per_warp * (1 + n_cg0 + n_cg1) + cfg.tmem_user_threads = cfg.threads_per_warp * (1 + n_cg0 + n_cg1) cfg.inverse_barrier_threads = cfg.threads_per_warp * n_cg0 - cfg.init_state_store_barrier_threads = cfg.threads_per_warp * n_cg1 cfg.tmem_state_acc_offset = 0 - cfg.tmem_state_inp_offset = cfg.tmem_state_acc_offset + cfg.tmem_state_acc_stages * 128 - cfg.tmem_cg0_acc_offset = cfg.tmem_state_inp_offset + cfg.tmem_state_inp_stages * 64 + cfg.tmem_state_input_offset = cfg.tmem_state_acc_offset + cfg.tmem_state_acc_stages * 128 + cfg.tmem_cg0_acc_offset = cfg.tmem_state_input_offset + cfg.tmem_state_input_stages * 64 cfg.tmem_cg1_acc_offset = cfg.tmem_cg0_acc_offset + cfg.tmem_cg0_acc_stages * 64 - cfg.tmem_y_decay_u_inp_offset = cfg.tmem_cg1_acc_offset + cfg.tmem_cg1_acc_stages * 64 + cfg.tmem_y_decay_u_input_offset = cfg.tmem_cg1_acc_offset + cfg.tmem_cg1_acc_stages * 64 return cfg @@ -2487,7 +2355,7 @@ def get_compiled_cache( log_gate: bool, safe_gate: bool, beta_sigmoid: bool, - dyn_sched: bool, + dynamic_scheduling: bool, run_order: bool, order_gen: bool, ): @@ -2505,9 +2373,12 @@ def compile( log_gate: bool = False, safe_gate: bool = False, beta_sigmoid: bool = False, - dyn_sched: bool = False, + dynamic_scheduling: bool = False, *, num_sm: int, + h_k: int, + h_v: int, + n_heads_out: int, k_cute, v_cute, gate_cute, @@ -2519,7 +2390,7 @@ def compile( state_out_cute, work_items_cute=None, work_count_cute=None, - sched_ctr_cute=None, + scheduler_counter_cute=None, checkpoint_every_n_tokens, workspace_cute, stream, @@ -2536,8 +2407,11 @@ def compile( log_gate=log_gate, safe_gate=safe_gate, beta_sigmoid=beta_sigmoid, - dyn_sched=dyn_sched, + dynamic_scheduling=dynamic_scheduling, ) + cfg.h_k = h_k + cfg.h_v = h_v + cfg.n_heads_out = n_heads_out return cute.compile( host, @@ -2553,7 +2427,7 @@ def compile( state_out_cute, work_items_cute, work_count_cute, - sched_ctr_cute, + scheduler_counter_cute, checkpoint_every_n_tokens, workspace_cute, stream, @@ -2573,8 +2447,8 @@ def chunk_gdn_recompute_sm100( output_state_checkpoints=None, work_items=None, work_count=None, - sched_ctr=None, - sched_all=None, + scheduler_counter=None, + scheduler_all=None, work_item_scratch=None, order_in_prologue: bool = False, log_gate: bool = False, @@ -2589,8 +2463,8 @@ def chunk_gdn_recompute_sm100( """Execute the Blackwell chunked GDN recompute kernel (state/checkpoint-only, THD / varlen entry). - All tensors are contiguous, DLPack-compatible CUDA tensors on the same - device. Compile-cache-and-replay: the kernel is compiled once per static + All tensors are DLPack-compatible CUDA tensors on the same device with a + stride-1 innermost dim (outer strides are runtime arguments). Compile-cache-and-replay: the kernel is compiled once per static config (dtypes, head counts, state flags) and replayed afterwards. Args: @@ -2613,8 +2487,8 @@ def chunk_gdn_recompute_sm100( series the backward pass consumes. work_items: ``(max_items, 8)`` int32 work-item table from ``common/split_k.py`` (REQUIRED; an uncut table row is the whole - (b, h) sequence). Each item computes chunks ``[cstart, wend)`` - and writes checkpoints only for ``[wstart, wend)``. + (b, h) sequence). Each item computes chunks ``[compute_start, write_end)`` + and writes checkpoints only for ``[write_start, write_end)``. work_count: ``(1,)`` int32 device-side item count (REQUIRED) log_gate: ``gate`` holds natural-log decay values; the gate warp skips its log2 (rescales by 1/ln2) instead of exponentiating @@ -2638,12 +2512,14 @@ def chunk_gdn_recompute_sm100( store_final_state = output_state is not None enable_checkpoints = checkpoint_every_n_tokens > 0 if work_items is None or work_count is None: - raise ValueError("work_items/work_count are required (the split-table stage builds them for every launch)") - dyn_sched = sched_ctr is not None + raise ValueError( + "work_items/work_count are required (built by the split-table stage, or by an order-generating prologue when work_item_scratch is None)" + ) + dynamic_scheduling = scheduler_counter is not None run_order = bool(order_in_prologue) order_gen = work_item_scratch is None - if run_order and sched_all is None: - raise ValueError("order_in_prologue requires sched_all (the prologue zeroes both consumers' sched rings)") + if run_order and scheduler_all is None: + raise ValueError("order_in_prologue requires scheduler_all (the prologue zeroes both consumers' scheduler rings)") if not (enable_checkpoints or store_final_state): raise ValueError("output_state_checkpoints or output_state is required") if safe_gate and (a_log is None or dt_bias is None): @@ -2677,22 +2553,18 @@ def chunk_gdn_recompute_sm100( log_gate, safe_gate, use_beta_sigmoid, - dyn_sched, + dynamic_scheduling, run_order, order_gen, ) if "compiled" not in cache: - k_cute = from_dlpack(k, assumed_align=16) - k_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) - v_cute = from_dlpack(v, assumed_align=16) - v_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) - gate_cute = from_dlpack(gate, assumed_align=16) - gate_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1), divisibility=1) + k_cute = from_dlpack(k, assumed_align=16).mark_layout_dynamic(leading_dim=2) + v_cute = from_dlpack(v, assumed_align=16).mark_layout_dynamic(leading_dim=2) + gate_cute = from_dlpack(gate, assumed_align=16).mark_layout_dynamic(leading_dim=1) a_log_cute = from_dlpack(a_log, assumed_align=4) if a_log is not None else None dt_bias_cute = from_dlpack(dt_bias, assumed_align=4) if dt_bias is not None else None - beta_cute = from_dlpack(beta, assumed_align=16) - beta_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1), divisibility=1) + beta_cute = from_dlpack(beta, assumed_align=16).mark_layout_dynamic(leading_dim=1) cu_seqlens_cute = from_dlpack(cu_seqlens, assumed_align=8 if str(cu_seqlens.dtype).endswith("int64") else 4).mark_layout_dynamic() state_in_cute = None @@ -2711,9 +2583,9 @@ def chunk_gdn_recompute_sm100( work_items_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1), divisibility=1) work_count_cute = from_dlpack(work_count, assumed_align=4).mark_layout_dynamic() - sched_ctr_cute = None - if dyn_sched: - sched_ctr_cute = from_dlpack(sched_ctr, assumed_align=4).mark_layout_dynamic() + scheduler_counter_cute = None + if dynamic_scheduling: + scheduler_counter_cute = from_dlpack(scheduler_counter, assumed_align=4).mark_layout_dynamic() cache["compiled"] = compile( io_dtype, @@ -2725,8 +2597,11 @@ def chunk_gdn_recompute_sm100( log_gate, safe_gate, use_beta_sigmoid, - dyn_sched, + dynamic_scheduling, num_sm=multiprocessor_count(current_device()), + h_k=HK, + h_v=HV, + n_heads_out=HO, k_cute=k_cute, v_cute=v_cute, gate_cute=gate_cute, @@ -2738,7 +2613,7 @@ def chunk_gdn_recompute_sm100( state_out_cute=state_out_cute, work_items_cute=work_items_cute, work_count_cute=work_count_cute, - sched_ctr_cute=sched_ctr_cute, + scheduler_counter_cute=scheduler_counter_cute, checkpoint_every_n_tokens=checkpoint_every_n_tokens, workspace_cute=workspace_cute, stream=cu_stream, @@ -2747,17 +2622,13 @@ def chunk_gdn_recompute_sm100( compiled = cache["compiled"] if "prologue" not in cache: - k_pl = from_dlpack(k, assumed_align=16) - k_pl.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) - v_pl = from_dlpack(v, assumed_align=16) - v_pl.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) - gate_pl = from_dlpack(gate, assumed_align=16) - gate_pl.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1), divisibility=1) + k_pl = from_dlpack(k, assumed_align=16).mark_layout_dynamic(leading_dim=2) + v_pl = from_dlpack(v, assumed_align=16).mark_layout_dynamic(leading_dim=2) + gate_pl = from_dlpack(gate, assumed_align=16).mark_layout_dynamic(leading_dim=1) cu_pl = from_dlpack(cu_seqlens, assumed_align=8 if str(cu_seqlens.dtype).endswith("int64") else 4).mark_layout_dynamic() checkpoints_pl = None if enable_checkpoints: - checkpoints_pl = from_dlpack(output_state_checkpoints, assumed_align=16) - checkpoints_pl.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2, 3), divisibility=1) + checkpoints_pl = from_dlpack(output_state_checkpoints, assumed_align=16).mark_layout_dynamic(leading_dim=3) staging_pl = None if not order_gen: staging_pl = from_dlpack(work_item_scratch, assumed_align=16) @@ -2765,10 +2636,10 @@ def chunk_gdn_recompute_sm100( work_count_pl = from_dlpack(work_count, assumed_align=4).mark_layout_dynamic() work_items_pl = from_dlpack(work_items, assumed_align=16) work_items_pl.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1), divisibility=1) - sched_all_pl = None + scheduler_all_pl = None if run_order: - sched_all_pl = from_dlpack(sched_all, assumed_align=4).mark_layout_dynamic() - ws_pl = from_dlpack(workspace, assumed_align=128).mark_layout_dynamic() + scheduler_all_pl = from_dlpack(scheduler_all, assumed_align=4).mark_layout_dynamic() + workspace_pl = from_dlpack(workspace, assumed_align=128).mark_layout_dynamic() cache["prologue"] = cute.compile( prologue, io_dtype, @@ -2783,9 +2654,9 @@ def chunk_gdn_recompute_sm100( staging_pl, work_count_pl, work_items_pl, - sched_all_pl, + scheduler_all_pl, cutlass.Int32(checkpoint_every_n_tokens), - ws_pl, + workspace_pl, cu_stream, options="--enable-tvm-ffi", ) @@ -2798,7 +2669,7 @@ def chunk_gdn_recompute_sm100( work_item_scratch if not order_gen else None, work_count, work_items, - sched_all if run_order else None, + scheduler_all if run_order else None, checkpoint_every_n_tokens, workspace, cu_stream, @@ -2815,7 +2686,7 @@ def chunk_gdn_recompute_sm100( output_state, work_items, work_count, - sched_ctr, + scheduler_counter, checkpoint_every_n_tokens, workspace, cu_stream, @@ -2835,8 +2706,8 @@ def run_recompute( output_state_checkpoints, work_items, work_count, - sched_ctr, - sched_all, + scheduler_counter, + scheduler_all, work_item_scratch, tensormap_workspace, checkpoint_every_n_tokens, @@ -2857,7 +2728,7 @@ def run_recompute( work_item_scratch, work_count, work_items, - sched_all, + scheduler_all, checkpoint_every_n_tokens, tensormap_workspace, cu_stream, @@ -2874,7 +2745,7 @@ def run_recompute( output_state, work_items, work_count, - sched_ctr, + scheduler_counter, checkpoint_every_n_tokens, tensormap_workspace, cu_stream, diff --git a/python/cudnn/linear_attention/frost/kernel/kda_bprop_f16.py b/python/cudnn/linear_attention/frost/kernel/kda_bprop_f16.py index ea64a186c..f9a8467d7 100644 --- a/python/cudnn/linear_attention/frost/kernel/kda_bprop_f16.py +++ b/python/cudnn/linear_attention/frost/kernel/kda_bprop_f16.py @@ -1,58 +1,3 @@ -"""Chunked Kimi Delta Attention (KDA) BPROP kernel for Blackwell SM100/SM103 -(Cutlass DSL), BT=16 tiling with a per-key-channel decay. Framework-neutral -entry ``chunk_kda_bwd_sm100``. - -Algorithm overview (per chunk c, iterated c = NT-1 .. 0; within-chunk -log2-domain gate cumsum G[t,d], eG = 2^G, eGl = 2^G[BT-1]): - Inputs : Q/K[BT,DK], V/dO[BT,DV], S = state_checkpoints[c-1] (state ENTERING chunk c, KV), - Gate[BT,DK], Beta[BT] - State : dH[DV,DK] (state gradient, fp32 TMEM, accumulated backward) - - Operands (WG0, prefill recompute): K_decay = eG.K, K_inv = K/eG, - K_restore = (eGl/eG).K, Q_decay = eG.Q, diag(eGl). - - Forward recompute: A_kk = K_decay@K_inv^T; T_inv = (I + Beta.tril(A_kk,-1))^-1 - (register Neumann); A = tril_incl(Q_decay@K_inv^T); Y = Beta.(V - S^T K_decay); - U = T_inv@Y. - - Backward math: - dU = K_restore@dH + A^T@dO (Q_decay carries scale, so A does too) - dY = T_inv^T@dU dV = Beta.dY - dA = tril_incl(dO@U^T) (unscaled) dM = dY@U^T dM_strict = Beta_row.strict(dM) - dQ = eG.scale.(dO@S^T + dA@K_inv) - dK = eG.dK_decay + dK_inv/eG + (eGl/eG).dK_restore where (sign-flipped parts) - dK_decay part = (Beta.dY)@S^T + dM_strict@K_inv (= -dK_decay) - dK_inv part = dA^T@(scale.Q_decay) - dM_strict^T@K_decay (= dK_inv; one TMEM - acc, the minus rides the staged -dM_strict tile) - dK_restore part = U@dH^T (= +dK_restore) - dBeta = sum_v dY.Y / Beta - sum_{j KdaBwdBars: - """Bars factory. MUST be called from inside ``kernel`` (allocates the - mbarrier rings in SMEM ahead of the data buffers).""" + """KdaBwdBars factory.""" def alloc(n): return cutlass.Array(cutlass.Int64, n, space=cutlass.AddressSpace.smem, alignment=8) @@ -224,12 +159,12 @@ def alloc(n): mb_dqk_acc_done=MBarrier(alloc(1), stages=1, init_count=CG2, producer=Producer.THREAD), mb_qk_raw_ready=MBarrier(alloc(cfg.tmem_qk_raw_stages), stages=cfg.tmem_qk_raw_stages, init_count=CG0, producer=Producer.THREAD), mb_qk_raw_done=MBarrier(alloc(cfg.tmem_qk_raw_stages), stages=cfg.tmem_qk_raw_stages, init_count=CG2, producer=Producer.THREAD), - mb_state_inp_ready=MBarrier(alloc(2), stages=2, init_count=CG0, producer=Producer.THREAD), - mb_state_inp_done=MBarrier(alloc(2), stages=2, init_count=MMA, producer=Producer.MMA_COMMIT), - mb_state_inp_cg2_done=MBarrier(alloc(2), stages=2, init_count=CG2, producer=Producer.THREAD), - mb_y_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1, producer=Producer.THREAD), - mb_du_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1, producer=Producer.THREAD), - mb_neg_beta_dy_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1, producer=Producer.THREAD), + mb_state_input_ready=MBarrier(alloc(2), stages=2, init_count=CG0, producer=Producer.THREAD), + mb_state_input_done=MBarrier(alloc(2), stages=2, init_count=MMA, producer=Producer.MMA_COMMIT), + mb_state_input_cg2_done=MBarrier(alloc(2), stages=2, init_count=CG2, producer=Producer.THREAD), + mb_y_input_ready=MBarrier(alloc(1), stages=1, init_count=CG1, producer=Producer.THREAD), + mb_du_input_ready=MBarrier(alloc(1), stages=1, init_count=CG1, producer=Producer.THREAD), + mb_neg_beta_dy_input_ready=MBarrier(alloc(1), stages=1, init_count=CG1, producer=Producer.THREAD), mb_k_decay_inv_ready=MBarrier(alloc(cfg.smem_decay_stages), stages=cfg.smem_decay_stages, init_count=CG0, producer=Producer.THREAD), mb_q_decay_k_restore_ready=MBarrier(alloc(cfg.smem_decay_stages), stages=cfg.smem_decay_stages, init_count=CG0, producer=Producer.THREAD), mb_decay_done=MBarrier(alloc(cfg.smem_decay_stages), stages=cfg.smem_decay_stages, init_count=MMA, producer=Producer.MMA_COMMIT), @@ -245,7 +180,7 @@ def alloc(n): mb_dy_smem_ready=MBarrier(alloc(1), stages=1, init_count=CG1, producer=Producer.THREAD), mb_dbeta_m_ready=MBarrier(alloc(1), stages=1, init_count=WARP, producer=Producer.THREAD), mb_dstate_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA, producer=Producer.MMA_COMMIT), - mb_dstate_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1, producer=Producer.THREAD), + mb_dstate_input_ready=MBarrier(alloc(1), stages=1, init_count=CG1, producer=Producer.THREAD), mb_dstate_smem_ready=MBarrier(alloc(1), stages=1, init_count=CG1, producer=Producer.THREAD), mb_dstate_smem_done=MBarrier(alloc(1), stages=1, init_count=MMA, producer=Producer.MMA_COMMIT), mb_dstate_smem_cg2_done=MBarrier(alloc(1), stages=1, init_count=CG2, producer=Producer.THREAD), @@ -259,43 +194,46 @@ def alloc(n): mb_dgate_tmastg_done=MBarrier(alloc(cfg.smem_dgate_stages), stages=cfg.smem_dgate_stages, init_count=WARP, producer=Producer.THREAD), mb_dstate0_acc_stored=MBarrier(alloc(1), stages=1, init_count=CG1, producer=Producer.THREAD), mb_tmem_done=MBarrier(alloc(1), stages=1, init_count=CG1 + CG2, producer=Producer.THREAD), - mb_sched_ready=MBarrier(alloc(cfg.sched_stages), stages=cfg.sched_stages, init_count=1, producer=Producer.THREAD), - mb_sched_done=MBarrier(alloc(cfg.sched_stages), stages=cfg.sched_stages, init_count=15, producer=Producer.THREAD), + mb_scheduler_ready=MBarrier(alloc(cfg.scheduler_stages), stages=cfg.scheduler_stages, init_count=1, producer=Producer.THREAD), + mb_scheduler_done=MBarrier(alloc(cfg.scheduler_stages), stages=cfg.scheduler_stages, init_count=15, producer=Producer.THREAD), ) -# ---- Dynamic tile scheduler ------------------------------------------------------ +# ---- Dynamic tile scheduler ---------------------------------------------------------- @cute.jit -def sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ctas): +def scheduler_publish_next(cfg, bars, sScheduler, mScheduler, scheduler_state, tile_idx, num_ctas, tail_base, tail_row, elect_one): """TMA-warp side: pull the next tile off the global ticket, publish it.""" - if cutlass.const_expr(cfg.dyn_sched): - bars.mb_sched_done[sched_state.idx].wait(sched_state.phase) - if nvvm.elect_sync(): - fetched = cutlass.Int32(nvvm.atomicrmw("add", mSched.iterator, cutlass.Int32(1), mem_order="relaxed", syncscope="gpu")) - sSched[sched_state.idx] = num_ctas + fetched + if cutlass.const_expr(cfg.dynamic_scheduling): + sentinel = cutlass.Int32(1 << 28) + pinned = tail_row if tile_idx < tail_base else sentinel + bars.mb_scheduler_done[scheduler_state.idx].wait(scheduler_state.phase) + if elect_one: + fetched = cutlass.Int32(nvvm.atomicrmw("add", mScheduler.iterator, cutlass.Int32(1), mem_order="relaxed", syncscope="gpu")) + granted = num_ctas + fetched + sScheduler[scheduler_state.idx] = granted if granted < tail_base else pinned nvvm.bar_warp_sync(cute.arch.FULL_MASK) - next_tile = sSched[sched_state.idx] - if nvvm.elect_sync(): - bars.mb_sched_ready[sched_state.idx].arrive() - return next_tile, advance(sched_state, cfg.sched_stages) - return tile_idx + num_ctas, sched_state + next_tile = sScheduler[scheduler_state.idx] + if elect_one: + bars.mb_scheduler_ready[scheduler_state.idx].arrive() + return next_tile, advance(scheduler_state, cfg.scheduler_stages) + return tile_idx + num_ctas, scheduler_state @cute.jit -def sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas): +def scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one): """Consumer side: read the TMA warp's published next tile.""" - if cutlass.const_expr(cfg.dyn_sched): - bars.mb_sched_ready[sched_state.idx].wait(sched_state.phase) - next_tile = sSched[sched_state.idx] - if nvvm.elect_sync(): - bars.mb_sched_done[sched_state.idx].arrive() - return next_tile, advance(sched_state, cfg.sched_stages) - return tile_idx + num_ctas, sched_state + if cutlass.const_expr(cfg.dynamic_scheduling): + bars.mb_scheduler_ready[scheduler_state.idx].wait(scheduler_state.phase) + next_tile = sScheduler[scheduler_state.idx] + if elect_one: + bars.mb_scheduler_done[scheduler_state.idx].arrive() + return next_tile, advance(scheduler_state, cfg.scheduler_stages) + return tile_idx + num_ctas, scheduler_state -# ---- Warp bodies ----------------------------------------------------------------- +# ---- Warp bodies --------------------------------------------------------------------- @cute.jit @@ -306,8 +244,8 @@ def epilogue_warp( num_ctas, cu_seqlens, mWorkItems, - sSched, - lane, + sScheduler, + lane_idx, sK_inv_raw, sQ_decay_raw, sDo_raw, @@ -326,29 +264,28 @@ def epilogue_warp( """Epilogue warp role (warp 15): the register-MMA A/dA tiles and the dQ/dK/dV/dGate TMA stores, in chunk order with a one-behind store ladder.""" - elect_one = nvvm.elect_sync() - nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) - # ---- ldmatrix/stmatrix lane decode ------------------------------------------- - rhs_row_coord = lane % 8 + (cutlass.Int32(8) if (lane // 16) else cutlass.Int32(0)) - rhs_col_offset = cutlass.Int32(8) if ((lane // 8) % 2) else cutlass.Int32(0) - lhs_row_coord = lane % 8 + (cutlass.Int32(8) if ((lane // 8) % 2) else cutlass.Int32(0)) - lhs_col_offset = cutlass.Int32(8) if ((lane // 8) // 2) else cutlass.Int32(0) - stsm_row_coord = lane & 7 - stsm_col_coord = cutlass.Int32(0) - if (lane // 8) & 1: - stsm_row_coord = stsm_row_coord + cutlass.Int32(8) - if lane // 8 >= 2: - stsm_col_coord = cutlass.Int32(8) - stsm_idx = swizzle_lin_S(stsm_row_coord * cfg.b_t + stsm_col_coord, bbits=1, mbase=3, sshift=3) - row_lo = lane // 4 + elect_one = nvvm.elect_sync() + # ---- ldmatrix/stmatrix lane decode ----------------------------------------------- + b_row_coord = lane_idx % 8 + (cutlass.Int32(8) if (lane_idx // 16) else cutlass.Int32(0)) + b_col_offset = cutlass.Int32(8) if ((lane_idx // 8) % 2) else cutlass.Int32(0) + a_row_coord = lane_idx % 8 + (cutlass.Int32(8) if ((lane_idx // 8) % 2) else cutlass.Int32(0)) + a_col_offset = cutlass.Int32(8) if ((lane_idx // 8) // 2) else cutlass.Int32(0) + intermediate_row_coord = lane_idx & 7 + intermediate_col_coord = cutlass.Int32(0) + if (lane_idx // 8) & 1: + intermediate_row_coord = intermediate_row_coord + cutlass.Int32(8) + if lane_idx // 8 >= 2: + intermediate_col_coord = cutlass.Int32(8) + intermediate_idx = intermediate_row_coord * cfg.b_t + swizzle_xor_32b(intermediate_row_coord, intermediate_col_coord) + row_lo = lane_idx // 4 row_hi = row_lo + cutlass.Int32(8) tril_incl_mask = cutlass.Int32(0) for accum_idx in cutlass.range_constexpr(8): row_coord = row_hi if cutlass.const_expr(accum_idx % 4 >= 2) else row_lo - col_coord = (accum_idx // 4) * 8 + 2 * (lane % 4) + col_coord = (accum_idx // 4) * 8 + 2 * (lane_idx % 4) if cutlass.const_expr(accum_idx % 2 == 1): col_coord = col_coord + cutlass.Int32(1) tril_incl_mask = tril_incl_mask | (cutlass.Int32(1 << accum_idx) if row_coord >= col_coord else cutlass.Int32(0)) @@ -404,11 +341,13 @@ def epilogue_warp( dk_index = PipelineState.start(phase=0) dv_index = PipelineState.start(phase=0) dgate_index = PipelineState.start(phase=0) - sched_state = PipelineState.start(phase=0) + scheduler_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) FIRST_STATE_CHUNK = 0 if cfg.use_initial_state else 1 while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) head_o = head_idx slot = batch_idx * cutlass.Int32(TENSOR_MAP_QWORDS) if elect_one: @@ -420,13 +359,13 @@ def epilogue_warp( tma_tensormap_acquire(desc_dk_slot) tma_tensormap_acquire(desc_dv_slot) tma_tensormap_acquire(desc_dgate_slot) - num_compute_chunks = cend - wstart + num_compute_chunks = compute_end - write_start pend_start = cutlass.Int32(0) pend_writes = cutlass.Boolean(False) for rev_idx in cutlass.range(num_compute_chunks, unroll=1): - chunk_idx = cend - cutlass.Int32(1) - rev_idx + chunk_idx = compute_end - cutlass.Int32(1) - rev_idx chunk_start = chunk_idx * cfg.b_t - writes = chunk_idx < wend + writes = chunk_idx < write_end chunk_serial = chunk_serial_base + rev_idx decay_stage = chunk_serial % cfg.smem_decay_stages intermediate_stage = chunk_serial % cfg.smem_intermediate_stages @@ -436,31 +375,31 @@ def epilogue_warp( sDo_ptr = sDo_raw.data_ptr() + raw_stage * (cfg.d_v * cfg.b_t) sIntermediate_ptr = sIntermediate_raw.data_ptr() + intermediate_stage * (cfg.intermediate_tiles * cfg.b_t * cfg.b_t) - # ---- A = tril_incl(Q_decay @ K_inv^T) -------------------------------- + # ---- A = tril(Q decay @ K inv^T, 0) -------------------------------------- bars.mb_a_done[intermediate_stage].wait(((chunk_serial // cfg.smem_intermediate_stages) + 1) % 2) bars.mb_q_decay_k_restore_ready[decay_stage].wait((chunk_serial // cfg.smem_decay_stages) % 2) a_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): a_acc[accum_idx] = cutlass.Float32(0.0) - for k_block in cutlass.range_constexpr(cfg.d_k // 16): - a_col = k_block * 16 + lhs_col_offset + for i in cutlass.range_constexpr(cfg.d_k // 16): + a_col = i * 16 + a_col_offset a_seg = a_col // 64 - a_frag = nvvm.ldmatrix( - sQ_decay_ptr + a_seg * (cfg.b_t * 64) + lhs_row_coord * 64 + swizzle_xor_128b(lhs_row_coord, a_col - a_seg * 64, elem_bytes=2), + q_decay_frag = nvvm.ldmatrix( + sQ_decay_ptr + a_seg * (cfg.b_t * 64) + a_row_coord * 64 + swizzle_xor_128b(a_row_coord, a_col - a_seg * 64, elem_bytes=2), 4, nvvm.MMALayout.ROW, ) - b_col = k_block * 16 + rhs_col_offset + b_col = i * 16 + b_col_offset b_seg = b_col // 64 - b_frag = nvvm.ldmatrix( - sK_inv_ptr + b_seg * (cfg.b_t * 64) + rhs_row_coord * 64 + swizzle_xor_128b(rhs_row_coord, b_col - b_seg * 64, elem_bytes=2), + k_inv_frag = nvvm.ldmatrix( + sK_inv_ptr + b_seg * (cfg.b_t * 64) + b_row_coord * 64 + swizzle_xor_128b(b_row_coord, b_col - b_seg * 64, elem_bytes=2), 4, nvvm.MMALayout.ROW, ) mma_step( a_acc, - (a_frag[0], a_frag[1], a_frag[2], a_frag[3]), - (b_frag[0], b_frag[1], b_frag[2], b_frag[3]), + (q_decay_frag[0], q_decay_frag[1], q_decay_frag[2], q_decay_frag[3]), + (k_inv_frag[0], k_inv_frag[1], k_inv_frag[2], k_inv_frag[3]), k_step=0, M=16, N=16, @@ -469,7 +408,7 @@ def epilogue_warp( for accum_idx in cutlass.range_constexpr(8): a_acc[accum_idx] = a_acc[accum_idx] if (tril_incl_mask >> accum_idx) & 1 else cutlass.Float32(0.0) nvvm.stmatrix( - sIntermediate_ptr + stsm_idx, + sIntermediate_ptr + intermediate_idx, [ fp32_to_fp16(a_acc[0], a_acc[1], dtype=cfg.io_dtype), fp32_to_fp16(a_acc[2], a_acc[3], dtype=cfg.io_dtype), @@ -482,45 +421,43 @@ def epilogue_warp( nvvm.fence_proxy("async.shared", space="cta") bars.mb_a_ready[intermediate_stage].arrive() - # ---- dA = tril_incl(dO @ U^T) ---------------------------------------- + # ---- dA = tril(dO @ U^T, 0) ---------------------------------------------- bars.mb_u_smem_ready.wait(u_index.phase) u_index = advance(u_index, 1) da_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): da_acc[accum_idx] = cutlass.Float32(0.0) - for k_block in cutlass.range_constexpr(cfg.d_v // 16): - a_col = k_block * 16 + lhs_col_offset + for i in cutlass.range_constexpr(cfg.d_v // 16): + a_col = i * 16 + a_col_offset a_seg = a_col // 64 - a_frag = nvvm.ldmatrix( - sDo_ptr + a_seg * (cfg.b_t * 64) + lhs_row_coord * 64 + swizzle_xor_128b(lhs_row_coord, a_col - a_seg * 64, elem_bytes=2), + do_frag = nvvm.ldmatrix( + sDo_ptr + a_seg * (cfg.b_t * 64) + a_row_coord * 64 + swizzle_xor_128b(a_row_coord, a_col - a_seg * 64, elem_bytes=2), 4, nvvm.MMALayout.ROW, ) - b_col = k_block * 16 + rhs_col_offset + b_col = i * 16 + b_col_offset b_seg = b_col // 64 - b_frag = nvvm.ldmatrix( - sU_raw.data_ptr() + b_seg * (cfg.b_t * 64) + rhs_row_coord * 64 + swizzle_xor_128b(rhs_row_coord, b_col - b_seg * 64, elem_bytes=2), + u_frag = nvvm.ldmatrix( + sU_raw.data_ptr() + b_seg * (cfg.b_t * 64) + b_row_coord * 64 + swizzle_xor_128b(b_row_coord, b_col - b_seg * 64, elem_bytes=2), 4, nvvm.MMALayout.ROW, ) mma_step( da_acc, - (a_frag[0], a_frag[1], a_frag[2], a_frag[3]), - (b_frag[0], b_frag[1], b_frag[2], b_frag[3]), + (do_frag[0], do_frag[1], do_frag[2], do_frag[3]), + (u_frag[0], u_frag[1], u_frag[2], u_frag[3]), k_step=0, M=16, N=16, ab_dtype=cfg.io_dtype, ) - # fence: the dO/U ldmatrix reads must complete before this release - # licenses the TMA reload (sDo) nvvm.fence_proxy("async.shared", space="cta") bars.mb_do_done[raw_stage].arrive() for accum_idx in cutlass.range_constexpr(8): da_acc[accum_idx] = da_acc[accum_idx] if (tril_incl_mask >> accum_idx) & 1 else cutlass.Float32(0.0) bars.mb_da_done[intermediate_stage].wait(((chunk_serial // cfg.smem_intermediate_stages) + 1) % 2) nvvm.stmatrix( - sIntermediate_ptr + 2 * (cfg.b_t * cfg.b_t) + stsm_idx, + sIntermediate_ptr + 2 * (cfg.b_t * cfg.b_t) + intermediate_idx, [ fp32_to_fp16(da_acc[0], da_acc[1], dtype=cfg.io_dtype), fp32_to_fp16(da_acc[2], da_acc[3], dtype=cfg.io_dtype), @@ -534,7 +471,7 @@ def epilogue_warp( bars.mb_da_ready[intermediate_stage].arrive() raw_index = advance(raw_index, cfg.smem_raw_stages) - # ---- dQ/dK/dGate/dV: previous chunk, one-behind store ladder --------- + # ---- dQ/dK/dGate/dV: previous chunk, one-behind store ladder ------------- if rev_idx > 0: bars.mb_dq_tmastg_ready[dq_index.idx].wait(dq_index.phase) if pend_writes: @@ -575,7 +512,7 @@ def epilogue_warp( pend_start = chunk_start pend_writes = writes - # ---- tile tail: drain the last chunk's dQ/dK/dGate/dV -------------------- + # ---- tile tail: store last chunk dQ/dK/dGate/dV ------------------------------ if num_compute_chunks > 0: bars.mb_dq_tmastg_ready[dq_index.idx].wait(dq_index.phase) if pend_writes: @@ -615,7 +552,7 @@ def epilogue_warp( dv_index = advance(dv_index, cfg.smem_dv_stages) chunk_serial_base += num_compute_chunks - tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one) @cute.jit @@ -626,8 +563,8 @@ def super_mma_warp( num_ctas, cu_seqlens, mWorkItems, - sSched, - lane, + sScheduler, + lane_idx, sK_decay_raw, sK_inv_raw, sU_raw, @@ -640,41 +577,44 @@ def super_mma_warp( """Super-MMA warp role (warp 12): the Neumann T_inv and dM register MMAs plus the dBeta M-term row sums, in chunk order.""" nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) + elect_one = nvvm.elect_sync() sdy_index = PipelineState.start(phase=0) - # ---- ldmatrix lane decode ---------------------------------------------------- - rhs_row_coord = lane % 8 + (cutlass.Int32(8) if (lane // 16) else cutlass.Int32(0)) - rhs_col_offset = cutlass.Int32(8) if ((lane // 8) % 2) else cutlass.Int32(0) - lhs_row_coord = lane % 8 + (cutlass.Int32(8) if ((lane // 8) % 2) else cutlass.Int32(0)) - lhs_col_offset = cutlass.Int32(8) if ((lane // 8) // 2) else cutlass.Int32(0) - stsm_row_coord = lane & 7 - stsm_col_coord = cutlass.Int32(0) - if (lane // 8) & 1: - stsm_row_coord = stsm_row_coord + cutlass.Int32(8) - if lane // 8 >= 2: - stsm_col_coord = cutlass.Int32(8) - stsm_idx = swizzle_lin_S(stsm_row_coord * cfg.b_t + stsm_col_coord, bbits=1, mbase=3, sshift=3) - row_lo = lane // 4 + # ---- ldmatrix/stmatrix lane decode ----------------------------------------------- + b_row_coord = lane_idx % 8 + (cutlass.Int32(8) if (lane_idx // 16) else cutlass.Int32(0)) + b_col_offset = cutlass.Int32(8) if ((lane_idx // 8) % 2) else cutlass.Int32(0) + a_row_coord = lane_idx % 8 + (cutlass.Int32(8) if ((lane_idx // 8) % 2) else cutlass.Int32(0)) + a_col_offset = cutlass.Int32(8) if ((lane_idx // 8) // 2) else cutlass.Int32(0) + intermediate_row_coord = lane_idx & 7 + intermediate_col_coord = cutlass.Int32(0) + if (lane_idx // 8) & 1: + intermediate_row_coord = intermediate_row_coord + cutlass.Int32(8) + if lane_idx // 8 >= 2: + intermediate_col_coord = cutlass.Int32(8) + intermediate_idx = intermediate_row_coord * cfg.b_t + swizzle_xor_32b(intermediate_row_coord, intermediate_col_coord) + row_lo = lane_idx // 4 row_hi = row_lo + cutlass.Int32(8) - # hoisted tril bitmasks: bit i = row > col / row == col for accum index i + # tril bitmasks: bit i = row > col / row == col for accum index i tril_strict_mask = cutlass.Int32(0) eye_mask = cutlass.Int32(0) for accum_idx in cutlass.range_constexpr(8): row_coord = row_hi if cutlass.const_expr(accum_idx % 4 >= 2) else row_lo - col_coord = (accum_idx // 4) * 8 + 2 * (lane % 4) + col_coord = (accum_idx // 4) * 8 + 2 * (lane_idx % 4) if cutlass.const_expr(accum_idx % 2 == 1): col_coord = col_coord + cutlass.Int32(1) tril_strict_mask = tril_strict_mask | (cutlass.Int32(1 << accum_idx) if row_coord > col_coord else cutlass.Int32(0)) eye_mask = eye_mask | (cutlass.Int32(1 << accum_idx) if row_coord == col_coord else cutlass.Int32(0)) chunk_serial_base = cutlass.Int32(0) - sched_state = PipelineState.start(phase=0) + scheduler_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) FIRST_STATE_CHUNK = 0 if cfg.use_initial_state else 1 while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) - num_compute_chunks = cend - wstart + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) + num_compute_chunks = compute_end - write_start for rev_idx in cutlass.range(num_compute_chunks, unroll=1): chunk_serial = chunk_serial_base + rev_idx decay_stage = chunk_serial % cfg.smem_decay_stages @@ -686,38 +626,39 @@ def super_mma_warp( bars.mb_t_inv_done[intermediate_stage].wait(((chunk_serial // cfg.smem_intermediate_stages) + 1) % 2) - # ---- KK = K_decay @ K_inv^T ------------------------------------------ + # ---- KK = K decay @ K inv^T ---------------------------------------------- bars.mb_k_decay_inv_ready[decay_stage].wait((chunk_serial // cfg.smem_decay_stages) % 2) - kk_lhs_row = lhs_row_coord + kk_a_row = a_row_coord kk_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): kk_acc[accum_idx] = cutlass.Float32(0.0) - for k_block in cutlass.range_constexpr(cfg.d_k // 16): - a_col = k_block * 16 + lhs_col_offset + for i in cutlass.range_constexpr(cfg.d_k // 16): + a_col = i * 16 + a_col_offset a_seg = a_col // 64 - a_frag = nvvm.ldmatrix( - sK_decay_ptr + a_seg * (cfg.b_t * 64) + kk_lhs_row * 64 + swizzle_xor_128b(kk_lhs_row, a_col - a_seg * 64, elem_bytes=2), + k_decay_frag = nvvm.ldmatrix( + sK_decay_ptr + a_seg * (cfg.b_t * 64) + kk_a_row * 64 + swizzle_xor_128b(kk_a_row, a_col - a_seg * 64, elem_bytes=2), 4, nvvm.MMALayout.ROW, ) - b_col = k_block * 16 + rhs_col_offset + b_col = i * 16 + b_col_offset b_seg = b_col // 64 - b_frag = nvvm.ldmatrix( - sK_inv_ptr + b_seg * (cfg.b_t * 64) + rhs_row_coord * 64 + swizzle_xor_128b(rhs_row_coord, b_col - b_seg * 64, elem_bytes=2), + k_inv_frag = nvvm.ldmatrix( + sK_inv_ptr + b_seg * (cfg.b_t * 64) + b_row_coord * 64 + swizzle_xor_128b(b_row_coord, b_col - b_seg * 64, elem_bytes=2), 4, nvvm.MMALayout.ROW, ) + mma_step( kk_acc, - (a_frag[0], a_frag[1], a_frag[2], a_frag[3]), - (b_frag[0], b_frag[1], b_frag[2], b_frag[3]), + (k_decay_frag[0], k_decay_frag[1], k_decay_frag[2], k_decay_frag[3]), + (k_inv_frag[0], k_inv_frag[1], k_inv_frag[2], k_inv_frag[3]), k_step=0, M=16, N=16, ab_dtype=cfg.io_dtype, ) - # ---- L = Beta * tril(KK, -1) ----------------------------------------- + # ---- L = Beta * tril(KK, -1) --------------------------------------------- bars.mb_beta_ready[chunk_serial % cfg.smem_beta_stages].wait((chunk_serial // cfg.smem_beta_stages) % 2) beta_lo = (sBeta_ptr + row_lo).load().to(cutlass.Float32) beta_hi = (sBeta_ptr + row_hi).load().to(cutlass.Float32) @@ -733,6 +674,7 @@ def super_mma_warp( l_a3 = fp32_to_fp16(l_regs[6], l_regs[7], dtype=cfg.io_dtype) l_values = cutlass.Vector.from_elements((l_a0, l_a1, l_a2, l_a3), cutlass.Int32).bitcast(cfg.io_dtype).to(cutlass.Float32) + # ---- T^-1 = I - L, then three Neumann doubling rounds -------------------- tinv_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): eye = cutlass.Float32(1.0) if (eye_mask >> accum_idx) & 1 else cutlass.Float32(0.0) @@ -740,7 +682,8 @@ def super_mma_warp( lpow_a0, lpow_a1, lpow_a2, lpow_a3 = l_a0, l_a1, l_a2, l_a3 mov_lpow0, mov_lpow1, mov_lpow2, mov_lpow3 = movmatrix_16b(l_a0), movmatrix_16b(l_a1), movmatrix_16b(l_a2), movmatrix_16b(l_a3) - for _round in cutlass.range_constexpr(3): + for neumann_round in cutlass.range_constexpr(3): + # ---- Lpow = Lpow @ Lpow ---------------------------------------------- sq_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): sq_acc[accum_idx] = cutlass.Float32(0.0) @@ -758,6 +701,7 @@ def super_mma_warp( lpow_a2 = fp32_to_fp16(sq_acc[4], sq_acc[5], dtype=cfg.io_dtype) lpow_a3 = fp32_to_fp16(sq_acc[6], sq_acc[7], dtype=cfg.io_dtype) mov_lpow0, mov_lpow1, mov_lpow2, mov_lpow3 = movmatrix_16b(lpow_a0), movmatrix_16b(lpow_a1), movmatrix_16b(lpow_a2), movmatrix_16b(lpow_a3) + # ---- T^-1 += T^-1 @ Lpow --------------------------------------------- upd_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): upd_acc[accum_idx] = cutlass.Float32(0.0) @@ -784,7 +728,7 @@ def super_mma_warp( tinv_acc[6], tinv_acc[7] = fadd2(tinv_lo3, tinv_hi3, upd_acc[6], upd_acc[7]) nvvm.stmatrix( - sIntermediate_ptr + 1 * (cfg.b_t * cfg.b_t) + stsm_idx, + sIntermediate_ptr + 1 * (cfg.b_t * cfg.b_t) + intermediate_idx, [ fp32_to_fp16(tinv_acc[0], tinv_acc[1], dtype=cfg.io_dtype), fp32_to_fp16(tinv_acc[2], tinv_acc[3], dtype=cfg.io_dtype), @@ -797,38 +741,38 @@ def super_mma_warp( nvvm.fence_proxy("async.shared", space="cta") bars.mb_t_inv_ready[intermediate_stage].arrive() - # ---- dM = dY @ U^T --------------------------------------------------- + # ---- dM = dY @ U^T ------------------------------------------------------- bars.mb_dm_done[intermediate_stage].wait(((chunk_serial // cfg.smem_intermediate_stages) + 1) % 2) bars.mb_dy_smem_ready.wait(sdy_index.phase) sdy_index = advance(sdy_index, 1) dm_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): dm_acc[accum_idx] = cutlass.Float32(0.0) - for k_block in cutlass.range_constexpr(cfg.d_v // 16): - a_col = k_block * 16 + lhs_col_offset + for i in cutlass.range_constexpr(cfg.d_v // 16): + a_col = i * 16 + a_col_offset a_seg = a_col // 64 - a_frag = nvvm.ldmatrix( - sDy_raw.data_ptr() + a_seg * (cfg.b_t * 64) + lhs_row_coord * 64 + swizzle_xor_128b(lhs_row_coord, a_col - a_seg * 64, elem_bytes=2), + dy_frag = nvvm.ldmatrix( + sDy_raw.data_ptr() + a_seg * (cfg.b_t * 64) + a_row_coord * 64 + swizzle_xor_128b(a_row_coord, a_col - a_seg * 64, elem_bytes=2), 4, nvvm.MMALayout.ROW, ) - b_col = k_block * 16 + rhs_col_offset + b_col = i * 16 + b_col_offset b_seg = b_col // 64 - b_frag = nvvm.ldmatrix( - sU_raw.data_ptr() + b_seg * (cfg.b_t * 64) + rhs_row_coord * 64 + swizzle_xor_128b(rhs_row_coord, b_col - b_seg * 64, elem_bytes=2), + u_frag = nvvm.ldmatrix( + sU_raw.data_ptr() + b_seg * (cfg.b_t * 64) + b_row_coord * 64 + swizzle_xor_128b(b_row_coord, b_col - b_seg * 64, elem_bytes=2), 4, nvvm.MMALayout.ROW, ) mma_step( dm_acc, - (a_frag[0], a_frag[1], a_frag[2], a_frag[3]), - (b_frag[0], b_frag[1], b_frag[2], b_frag[3]), + (dy_frag[0], dy_frag[1], dy_frag[2], dy_frag[3]), + (u_frag[0], u_frag[1], u_frag[2], u_frag[3]), k_step=0, M=16, N=16, ab_dtype=cfg.io_dtype, ) - # ---- dM_strict = Beta_row . strict(dM) ------------------------------- + # ---- dM strict = Beta row . strict(dM) ----------------------------------- dm_strict_regs = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): beta_scale = beta_lo if accum_idx % 4 < 2 else beta_hi @@ -838,16 +782,16 @@ def super_mma_warp( w1 = fp32_to_fp16(dm_strict_regs[2], dm_strict_regs[3], dtype=cfg.io_dtype) w2 = fp32_to_fp16(dm_strict_regs[4], dm_strict_regs[5], dtype=cfg.io_dtype) w3 = fp32_to_fp16(dm_strict_regs[6], dm_strict_regs[7], dtype=cfg.io_dtype) - nvvm.stmatrix(sIntermediate_ptr + 3 * (cfg.b_t * cfg.b_t) + stsm_idx, [w0, w1, w2, w3], nvvm.MMALayout.ROW, shape=nvvm.StoreShape.M8N8) + nvvm.stmatrix(sIntermediate_ptr + 3 * (cfg.b_t * cfg.b_t) + intermediate_idx, [w0, w1, w2, w3], nvvm.MMALayout.ROW, shape=nvvm.StoreShape.M8N8) nw0 = fp32_to_fp16(-dm_strict_regs[0], -dm_strict_regs[1], dtype=cfg.io_dtype) nw1 = fp32_to_fp16(-dm_strict_regs[2], -dm_strict_regs[3], dtype=cfg.io_dtype) nw2 = fp32_to_fp16(-dm_strict_regs[4], -dm_strict_regs[5], dtype=cfg.io_dtype) nw3 = fp32_to_fp16(-dm_strict_regs[6], -dm_strict_regs[7], dtype=cfg.io_dtype) - nvvm.stmatrix(sIntermediate_ptr + 4 * (cfg.b_t * cfg.b_t) + stsm_idx, [nw0, nw1, nw2, nw3], nvvm.MMALayout.ROW, shape=nvvm.StoreShape.M8N8) + nvvm.stmatrix(sIntermediate_ptr + 4 * (cfg.b_t * cfg.b_t) + intermediate_idx, [nw0, nw1, nw2, nw3], nvvm.MMALayout.ROW, shape=nvvm.StoreShape.M8N8) nvvm.fence_proxy("async.shared", space="cta") bars.mb_dm_ready[intermediate_stage].arrive() - # ---- M-term: bsum = sum strict(dM . KK) ------------------------------ + # ---- M-term: bsum = sum strict(dM . KK) ---------------------------------- bsum_lo = cutlass.Float32(0.0) bsum_hi = cutlass.Float32(0.0) for accum_idx in cutlass.range_constexpr(8): @@ -860,13 +804,13 @@ def super_mma_warp( bsum_lo = bsum_lo + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, bsum_lo, 2, 31, kind=nvvm.Shfl.BFLY)) bsum_hi = bsum_hi + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, bsum_hi, 1, 31, kind=nvvm.Shfl.BFLY)) bsum_hi = bsum_hi + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, bsum_hi, 2, 31, kind=nvvm.Shfl.BFLY)) - if lane % 4 == 0: + if lane_idx % 4 == 0: sBetaM_raw[row_lo] = -bsum_lo sBetaM_raw[row_hi] = -bsum_hi nvvm.fence_proxy("async.shared", space="cta") bars.mb_dbeta_m_ready.arrive() chunk_serial_base += num_compute_chunks - tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one) @cute.jit @@ -877,21 +821,21 @@ def tcgen05_mma_warp( num_ctas, cu_seqlens, mWorkItems, - sSched, + sScheduler, tmem_base_holder, - sState_alt, - sState_direct, - sK_decay_lead16, - sK_inv_lead16, - sK_inv_amaj, - sK_restore_lead16, - sDo_lead16, - sDo_amaj, + sState_trans, + sState, + sK_decay, + sK_inv, + sK_inv_trans, + sK_restore, + sDo, + sDo_trans, sQ_decay_trans, sK_decay_trans, - sU_lead16, - sDv_lead16, - sDstate_alt, + sU, + sDv, + sDstate_trans, sIntermediate, sState_scale_diag, bars, @@ -906,7 +850,7 @@ def tcgen05_mma_warp( tmem_base = tmem_base_holder.load() bpe = cfg.io_dtype.width // 8 - # ---- chunk-invariant GEMM descriptors ---------------------------------------- + # ---- chunk-invariant GEMM descriptors -------------------------------------------- idesc_mv_nt = nvvm.Tcgen05InstrDesc.build( c_dtype=cutlass.Float32, a_dtype=cfg.io_dtype, @@ -922,7 +866,7 @@ def tcgen05_mma_warp( m_dim=cfg.d_v, a_major=1, ) - bmm_state_k_desc = MmaDesc( + bmm_state_k_decay_desc = MmaDesc( M=cfg.d_v, N=cfg.b_t, K=cfg.d_k, @@ -935,14 +879,14 @@ def tcgen05_mma_warp( idesc=idesc_state_k_at, kind=nvvm.Tcgen05MMAKind.F16, ) - idesc_state_k_kmaj = nvvm.Tcgen05InstrDesc.build( + idesc_state_k = nvvm.Tcgen05InstrDesc.build( c_dtype=cutlass.Float32, a_dtype=cfg.io_dtype, b_dtype=cfg.io_dtype, n_dim=cfg.b_t, m_dim=cfg.d_v, ) - bmm_state_k_kmaj_desc = MmaDesc( + bmm_state_k_decay_desc = MmaDesc( M=cfg.d_v, N=cfg.b_t, K=cfg.d_k, @@ -952,10 +896,10 @@ def tcgen05_mma_warp( btranspose=False, atranspose=False, cta_group=1, - idesc=idesc_state_k_kmaj, + idesc=idesc_state_k, kind=nvvm.Tcgen05MMAKind.F16, ) - bmm_dvinter_desc = MmaDesc( + bmm_dstate_k_restore_desc = MmaDesc( M=cfg.d_v, N=cfg.b_t, K=cfg.d_k, @@ -976,7 +920,7 @@ def tcgen05_mma_warp( a_major=1, b_major=1, ) - bmm_du_at_desc = MmaDesc( + bmm_do_a_desc = MmaDesc( M=cfg.d_v, N=cfg.b_t, K=cfg.b_t, @@ -998,7 +942,7 @@ def tcgen05_mma_warp( a_major=1, b_major=1, ) - bmm_dstate_q_at_desc = MmaDesc( + bmm_do_q_decay_desc = MmaDesc( M=cfg.d_v, N=cfg.d_k, K=cfg.b_t, @@ -1011,7 +955,7 @@ def tcgen05_mma_warp( idesc=idesc_dstate_q_at, kind=nvvm.Tcgen05MMAKind.F16, ) - bmm_qk_ts_desc = MmaDesc( + bmm_y_t_inv_desc = MmaDesc( M=cfg.d_v, N=cfg.b_t, K=cfg.b_t, @@ -1031,7 +975,7 @@ def tcgen05_mma_warp( m_dim=cfg.d_v, b_major=1, ) - bmm_qk_ts_t_desc = MmaDesc( + bmm_du_t_inv_trans_desc = MmaDesc( M=cfg.d_v, N=cfg.b_t, K=cfg.b_t, @@ -1050,7 +994,7 @@ def tcgen05_mma_warp( n_dim=16, m_dim=cfg.d_v, ) - bmm_diag_desc = MmaDesc( + bmm_dstate_diag_desc = MmaDesc( M=cfg.d_v, N=16, K=16, @@ -1070,7 +1014,7 @@ def tcgen05_mma_warp( m_dim=cfg.d_v, b_major=1, ) - bmm_dstate_k_desc = MmaDesc( + bmm_dy_k_decay_desc = MmaDesc( M=cfg.d_v, N=cfg.d_k, K=cfg.b_t, @@ -1089,7 +1033,7 @@ def tcgen05_mma_warp( n_dim=cfg.b_t, m_dim=cfg.d_k, ) - bmm_dstate_ts_desc = MmaDesc( + bmm_state_dy_desc = MmaDesc( M=cfg.d_k, N=cfg.b_t, K=cfg.d_v, @@ -1101,6 +1045,7 @@ def tcgen05_mma_warp( idesc=idesc_dstate, kind=nvvm.Tcgen05MMAKind.F16, ) + bmm_state_do_desc = bmm_state_dy_desc idesc_dstate_at = nvvm.Tcgen05InstrDesc.build( c_dtype=cutlass.Float32, a_dtype=cfg.io_dtype, @@ -1109,7 +1054,7 @@ def tcgen05_mma_warp( m_dim=cfg.d_k, a_major=1, ) - bmm_dstate_at_desc = MmaDesc( + bmm_dstate_u_desc = MmaDesc( M=cfg.d_k, N=cfg.b_t, K=cfg.d_v, @@ -1129,18 +1074,6 @@ def tcgen05_mma_warp( n_dim=cfg.b_t, m_dim=cfg.d_k, ) - bmm_dgrad_ts_desc = MmaDesc( - M=cfg.d_k, - N=cfg.b_t, - K=cfg.b_t, - bpe_a=bpe, - bpe_b=bpe, - tile_k_hw=16, - btranspose=False, - cta_group=1, - idesc=idesc_dgp, - kind=nvvm.Tcgen05MMAKind.F16, - ) idesc_dgp_at = nvvm.Tcgen05InstrDesc.build( c_dtype=cutlass.Float32, a_dtype=cfg.io_dtype, @@ -1149,7 +1082,7 @@ def tcgen05_mma_warp( m_dim=cfg.d_k, a_major=1, ) - bmm_dgrad_at_desc = MmaDesc( + bmm_k_inv_dm_desc = MmaDesc( M=cfg.d_k, N=cfg.b_t, K=cfg.b_t, @@ -1162,6 +1095,7 @@ def tcgen05_mma_warp( idesc=idesc_dgp_at, kind=nvvm.Tcgen05MMAKind.F16, ) + bmm_k_inv_da_desc = bmm_k_inv_dm_desc idesc_dgp_at_t = nvvm.Tcgen05InstrDesc.build( c_dtype=cutlass.Float32, a_dtype=cfg.io_dtype, @@ -1171,7 +1105,7 @@ def tcgen05_mma_warp( a_major=1, b_major=1, ) - bmm_dgrad_at_t_desc = MmaDesc( + bmm_k_decay_dm_desc = MmaDesc( M=cfg.d_k, N=cfg.b_t, K=cfg.b_t, @@ -1184,11 +1118,12 @@ def tcgen05_mma_warp( idesc=idesc_dgp_at_t, kind=nvvm.Tcgen05MMAKind.F16, ) + bmm_q_decay_da_desc = bmm_k_decay_dm_desc state_index = PipelineState.start(phase=0) - y_inp_index = PipelineState.start(phase=0) - dstate_inp_index = PipelineState.start(phase=0) - du_inp_index = PipelineState.start(phase=0) + y_input_index = PipelineState.start(phase=0) + dstate_input_index = PipelineState.start(phase=0) + du_input_index = PipelineState.start(phase=0) neg_beta_dy_index = PipelineState.start(phase=0) u_smem_index = PipelineState.start(phase=0) dstate_smem_index = PipelineState.start(phase=0) @@ -1200,30 +1135,32 @@ def tcgen05_mma_warp( intermediate_slot = (cfg.b_t * cfg.b_t * (cfg.io_dtype.width // 8)) >> 4 diag_seg = ((cfg.d_k // 16) * 256 * (cfg.io_dtype.width // 8)) >> 4 dv_seg = (cfg.b_t * cfg.d_v * (cfg.io_dtype.width // 8)) >> 4 - d_do_amaj0 = sDo_amaj[0].desc() + d_do_trans0 = sDo_trans[0].desc() d_qd_trans0 = sQ_decay_trans[0].desc() d_kd_trans0 = sK_decay_trans[0].desc() - d_ki_amaj0 = sK_inv_amaj[0].desc() + d_ki_trans0 = sK_inv_trans[0].desc() d_int0 = sIntermediate[0].desc() - d_kd_lead0 = sK_decay_lead16[0].desc() - d_do_lead0 = sDo_lead16[0].desc() - d_kr_lead0 = sK_restore_lead16[0].desc() + d_kd0 = sK_decay[0].desc() + d_do0 = sDo[0].desc() + d_kr0 = sK_restore[0].desc() d_diag0 = sState_scale_diag[0].desc() - d_dv_lead0 = sDv_lead16[0].desc() - d_dstate_alt0 = sDstate_alt[0].desc() - d_u_lead0 = sU_lead16[0].desc() + d_dv0 = sDv[0].desc() + d_dstate_trans0 = sDstate_trans[0].desc() + d_u0 = sU[0].desc() assert cfg.smem_state_stages == 1 - d_state_alt0 = sState_alt[0].desc() - d_state_direct0 = sState_direct[0].desc() + d_state_trans0 = sState_trans[0].desc() + d_state0 = sState[0].desc() dstate0_index = PipelineState.start(phase=0) chunk_serial_base = cutlass.Int32(0) - sched_state = PipelineState.start(phase=0) + scheduler_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) FIRST_STATE_CHUNK = 0 if cfg.use_initial_state else 1 while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) - num_compute_chunks = cend - wstart + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) + num_compute_chunks = compute_end - write_start for rev_idx in cutlass.range(num_compute_chunks, unroll=1): chunk_serial = chunk_serial_base + rev_idx decay_stage = chunk_serial % cfg.smem_decay_stages @@ -1235,27 +1172,27 @@ def tcgen05_mma_warp( has_dstate = cutlass.Boolean(True) raw_stage_idx = chunk_serial % cfg.smem_raw_stages - # ---- stage-derived operand descriptors ------------------------------- + # ---- stage-derived operand descriptors ----------------------------------- decay_op_off = decay_stage * op_seg - d_do_amaj = d_do_amaj0 + raw_stage_idx * do_seg + d_do_trans = d_do_trans0 + raw_stage_idx * do_seg d_qd_trans = d_qd_trans0 + decay_op_off d_kd_trans = d_kd_trans0 + decay_op_off - d_ki_amaj = d_ki_amaj0 + decay_op_off + d_ki_trans = d_ki_trans0 + decay_op_off d_int = d_int0 + intermediate_stage * intermediate_seg d_int_tinv = d_int + intermediate_slot d_int_da = d_int + 2 * intermediate_slot d_int_dm = d_int + 3 * intermediate_slot d_int_ndm = d_int + 4 * intermediate_slot - chunk_idx = cend - cutlass.Int32(1) - rev_idx + chunk_idx = compute_end - cutlass.Int32(1) - rev_idx - # ---- state_k = state(S) @ K_decay^T -------------------------------------- + # ---- k state = state(S) @ K decay^T -------------------------------------- bars.mb_k_decay_inv_ready[decay_stage].wait(decay_phase) if chunk_idx >= FIRST_STATE_CHUNK: bars.mb_state_ready[state_index.idx].wait(state_index.phase) mma_ss( - bmm_state_k_kmaj_desc, - d_state_direct0, - d_kd_lead0 + decay_op_off, + bmm_state_k_decay_desc, + d_state0, + d_kd0 + decay_op_off, nvvm.make_tmem_ptr((tmem_base + cfg.tmem_state_k_acc_offset), cutlass.Float32), accumulate=False, ) @@ -1267,65 +1204,65 @@ def tcgen05_mma_warp( # ---- dQ inter = state(T) @ dO^T ------------------------------------------ bars.mb_dqk_acc_done.wait(parts_done_index.phase) parts_done_index = advance(parts_done_index, 1) - bars.mb_state_inp_ready[chunk_serial % 2].wait((chunk_serial // 2) % 2) + bars.mb_state_input_ready[chunk_serial % 2].wait((chunk_serial // 2) % 2) bars.mb_do_ready[raw_stage_idx].wait((chunk_serial // cfg.smem_raw_stages) % 2) if chunk_idx >= FIRST_STATE_CHUNK: - a_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_state_inp_offset + (chunk_serial % 2) * (cfg.d_v // 2)), cutlass.Int8) - b_desc = d_do_lead0 + raw_stage_idx * do_seg + a_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_state_input_offset + (chunk_serial % 2) * (cfg.d_v // 2)), cutlass.Int8) + b_desc = d_do0 + raw_stage_idx * do_seg c_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dq_acc_offset), cutlass.Float32) - for sub in cutlass.range_constexpr(bmm_dstate_ts_desc.num_subtiles_B): - for k in cutlass.range_constexpr(bmm_dstate_ts_desc.sps_B): + for i in cutlass.range_constexpr(bmm_state_do_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_state_do_desc.sps_B): mma_ts_step( - bmm_dstate_ts_desc, - a_ptr.subview(sub * bmm_dstate_ts_desc.sps_B * bmm_dstate_ts_desc.tmem_advance_A), - b_desc + sub * (bmm_dstate_ts_desc.smem_subtile_B >> 4), + bmm_state_do_desc, + a_ptr.subview(i * bmm_state_do_desc.sps_B * bmm_state_do_desc.tmem_advance_A), + b_desc + i * (bmm_state_do_desc.smem_subtile_B >> 4), c_ptr, k, - cutlass.Boolean(sub + k > 0), + cutlass.Boolean(i + k > 0), ) - # ---- dU inter = dstate(T) @ K_restore ------------------------------------ + # ---- dU inter = dstate input(T) @ K restore ------------------------------ bars.mb_q_decay_k_restore_ready[decay_stage].wait(decay_phase) if has_dstate: - bars.mb_dstate_inp_ready.wait(dstate_inp_index.phase) - a_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dstate_inp_offset), cutlass.Int8) - b_desc = d_kr_lead0 + decay_op_off + bars.mb_dstate_input_ready.wait(dstate_input_index.phase) + a_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dstate_input_offset), cutlass.Int8) + b_desc = d_kr0 + decay_op_off c_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_du_acc_offset), cutlass.Float32) - for sub in cutlass.range_constexpr(bmm_dvinter_desc.num_subtiles_B): - for k in cutlass.range_constexpr(bmm_dvinter_desc.sps_B): + for i in cutlass.range_constexpr(bmm_dstate_k_restore_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_dstate_k_restore_desc.sps_B): mma_ts_step( - bmm_dvinter_desc, - a_ptr.subview(sub * bmm_dvinter_desc.sps_B * bmm_dvinter_desc.tmem_advance_A), - b_desc + sub * (bmm_dvinter_desc.smem_subtile_B >> 4), + bmm_dstate_k_restore_desc, + a_ptr.subview(i * bmm_dstate_k_restore_desc.sps_B * bmm_dstate_k_restore_desc.tmem_advance_A), + b_desc + i * (bmm_dstate_k_restore_desc.smem_subtile_B >> 4), c_ptr, k, - cutlass.Boolean(sub + k > 0), + cutlass.Boolean(i + k > 0), ) - # ---- dstate decay = dstate(T) @ diag(eGl) ------------------------------------ + # ---- dstate decay = dstate input(T) @ diag(eGl) -------------------------- if has_dstate: desc_diag = d_diag0 + decay_stage * diag_seg - for k_block in cutlass.range_constexpr(cfg.d_k // 16): - a_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dstate_inp_offset) + k_block * 8, cutlass.Int8) - b_desc = desc_diag.advance_start_address(k_block * 256 * 2) - c_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dstate_acc_offset) + k_block * 16, cutlass.Float32) - for sub in cutlass.range_constexpr(bmm_diag_desc.num_subtiles_B): - for k in cutlass.range_constexpr(bmm_diag_desc.sps_B): + for i in cutlass.range_constexpr(cfg.d_k // 16): + a_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dstate_input_offset) + i * 8, cutlass.Int8) + b_desc = desc_diag.advance_start_address(i * 256 * 2) + c_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dstate_acc_offset) + i * 16, cutlass.Float32) + for i in cutlass.range_constexpr(bmm_dstate_diag_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_dstate_diag_desc.sps_B): mma_ts_step( - bmm_diag_desc, - a_ptr.subview(sub * bmm_diag_desc.sps_B * bmm_diag_desc.tmem_advance_A), - b_desc + sub * (bmm_diag_desc.smem_subtile_B >> 4), + bmm_dstate_diag_desc, + a_ptr.subview(i * bmm_dstate_diag_desc.sps_B * bmm_dstate_diag_desc.tmem_advance_A), + b_desc + i * (bmm_dstate_diag_desc.smem_subtile_B >> 4), c_ptr, k, - cutlass.Boolean(sub + k > 0), + cutlass.Boolean(i + k > 0), ) - dstate_inp_index = advance(dstate_inp_index, 1) + dstate_input_index = advance(dstate_input_index, 1) - # ---- dU intra += dO^T(S) @ A ----------------------------------------- + # ---- dU intra += dO^T(S) @ A --------------------------------------------- bars.mb_a_ready[intermediate_stage].wait(intermediate_phase) mma_ss( - bmm_du_at_desc, - d_do_amaj, + bmm_do_a_desc, + d_do_trans, d_int, nvvm.make_tmem_ptr((tmem_base + cfg.tmem_du_acc_offset), cutlass.Float32), accumulate=has_dstate, @@ -1334,65 +1271,65 @@ def tcgen05_mma_warp( bars.mb_du_acc_ready.arrive(cta_group=1) bars.mb_a_done[intermediate_stage].arrive(cta_group=1) - # ---- dstate q-term += dO^T(S) @ Q_decay ---------------------------------- + # ---- dstate Q-term += dO^T(S) @ Q decay ---------------------------------- mma_ss( - bmm_dstate_q_at_desc, - d_do_amaj, + bmm_do_q_decay_desc, + d_do_trans, d_qd_trans, nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dstate_acc_offset), cutlass.Float32), accumulate=has_dstate, ) - # ---- U = Y(T) @ T_inv ------------------------------------------------ + # ---- U = Y(T) @ T^-1 ----------------------------------------------------- bars.mb_t_inv_ready[intermediate_stage].wait(intermediate_phase) - bars.mb_y_inp_ready.wait(y_inp_index.phase) - y_inp_index = advance(y_inp_index, 1) - a_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_y_inp_offset), cutlass.Int8) + bars.mb_y_input_ready.wait(y_input_index.phase) + y_input_index = advance(y_input_index, 1) + a_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_y_input_offset), cutlass.Int8) b_desc = d_int_tinv c_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_u_acc_offset), cutlass.Float32) - for sub in cutlass.range_constexpr(bmm_qk_ts_desc.num_subtiles_B): - for k in cutlass.range_constexpr(bmm_qk_ts_desc.sps_B): + for i in cutlass.range_constexpr(bmm_y_t_inv_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_y_t_inv_desc.sps_B): mma_ts_step( - bmm_qk_ts_desc, - a_ptr.subview(sub * bmm_qk_ts_desc.sps_B * bmm_qk_ts_desc.tmem_advance_A), - b_desc + sub * (bmm_qk_ts_desc.smem_subtile_B >> 4), + bmm_y_t_inv_desc, + a_ptr.subview(i * bmm_y_t_inv_desc.sps_B * bmm_y_t_inv_desc.tmem_advance_A), + b_desc + i * (bmm_y_t_inv_desc.smem_subtile_B >> 4), c_ptr, k, - cutlass.Boolean(sub + k > 0), + cutlass.Boolean(i + k > 0), ) if elect_one: bars.mb_u_acc_ready.arrive(cta_group=1) - # ---- dY = dU(T) @ T_inv ---------------------------------------------- - bars.mb_du_inp_ready.wait(du_inp_index.phase) - du_inp_index = advance(du_inp_index, 1) - a_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_du_inp_offset), cutlass.Int8) + # ---- dY = dU(T) @ T^-1 --------------------------------------------------- + bars.mb_du_input_ready.wait(du_input_index.phase) + du_input_index = advance(du_input_index, 1) + a_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_du_input_offset), cutlass.Int8) dy_b_desc = d_int_tinv c_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dy_acc_offset), cutlass.Float32) - for sub in cutlass.range_constexpr(bmm_qk_ts_t_desc.num_subtiles_B): - for k in cutlass.range_constexpr(bmm_qk_ts_t_desc.sps_B): + for i in cutlass.range_constexpr(bmm_du_t_inv_trans_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_du_t_inv_trans_desc.sps_B): mma_ts_step( - bmm_qk_ts_t_desc, - a_ptr.subview(sub * bmm_qk_ts_t_desc.sps_B * bmm_qk_ts_t_desc.tmem_advance_A), - dy_b_desc + sub * (bmm_qk_ts_t_desc.smem_subtile_B >> 4), + bmm_du_t_inv_trans_desc, + a_ptr.subview(i * bmm_du_t_inv_trans_desc.sps_B * bmm_du_t_inv_trans_desc.tmem_advance_A), + dy_b_desc + i * (bmm_du_t_inv_trans_desc.smem_subtile_B >> 4), c_ptr, k, - cutlass.Boolean(sub + k > 0), + cutlass.Boolean(i + k > 0), ) if elect_one: bars.mb_dy_acc_ready.arrive(cta_group=1) bars.mb_t_inv_done[intermediate_stage].arrive(cta_group=1) - # ---- dK_restore part = dstate(S) @ U^T ----------------------------------- + # ---- dK restore part = dstate(S) @ U^T ----------------------------------- bars.mb_u_smem_ready.wait(u_smem_index.phase) u_smem_index = advance(u_smem_index, 1) if has_dstate: bars.mb_dstate_smem_ready.wait(dstate_smem_index.phase) dstate_smem_index = advance(dstate_smem_index, 1) mma_ss( - bmm_dstate_at_desc, - d_dstate_alt0, - d_u_lead0, + bmm_dstate_u_desc, + d_dstate_trans0, + d_u0, nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dk_restore_acc_offset), cutlass.Float32), accumulate=False, ) @@ -1400,18 +1337,18 @@ def tcgen05_mma_warp( bars.mb_dk_restore_part_acc_ready.arrive(cta_group=1) bars.mb_dstate_smem_done.arrive(cta_group=1) - # ---- dstate dY-term += -Beta.dY(T) @ K_decay ----------------------------- - bars.mb_neg_beta_dy_inp_ready.wait(neg_beta_dy_index.phase) + # ---- dstate K-term += -Beta.dY(T) @ K decay ------------------------------ + bars.mb_neg_beta_dy_input_ready.wait(neg_beta_dy_index.phase) neg_beta_dy_index = advance(neg_beta_dy_index, 1) - a_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_neg_beta_dy_inp_offset), cutlass.Int8) + a_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_neg_beta_dy_input_offset), cutlass.Int8) b_desc = d_kd_trans c_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dstate_acc_offset), cutlass.Float32) - for sub in cutlass.range_constexpr(bmm_dstate_k_desc.num_subtiles_B): - for k in cutlass.range_constexpr(bmm_dstate_k_desc.sps_B): + for i in cutlass.range_constexpr(bmm_dy_k_decay_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_dy_k_decay_desc.sps_B): mma_ts_step( - bmm_dstate_k_desc, - a_ptr.subview(sub * bmm_dstate_k_desc.sps_B * bmm_dstate_k_desc.tmem_advance_A), - b_desc + sub * (bmm_dstate_k_desc.smem_subtile_B >> 4), + bmm_dy_k_decay_desc, + a_ptr.subview(i * bmm_dy_k_decay_desc.sps_B * bmm_dy_k_decay_desc.tmem_advance_A), + b_desc + i * (bmm_dy_k_decay_desc.smem_subtile_B >> 4), c_ptr, k, cutlass.Boolean(True), @@ -1419,29 +1356,29 @@ def tcgen05_mma_warp( if elect_one: bars.mb_dstate_acc_ready.arrive(cta_group=1) - # ---- dK_inv part = scale.Q_decay^T(S) @ dA --------------------------- + # ---- dK inv part = scale.Q decay^T(S) @ dA ------------------------------- bars.mb_da_ready[intermediate_stage].wait(intermediate_phase) mma_ss( - bmm_dgrad_at_t_desc, + bmm_q_decay_da_desc, d_qd_trans, d_int_da, nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dk_inv_acc_offset), cutlass.Float32), accumulate=False, ) - # ---- dQ attn += K_inv^T(S) @ dA^T ------------------------------------ + # ---- dQ attn += K inv^T(S) @ dA^T ---------------------------------------- if chunk_idx >= FIRST_STATE_CHUNK: mma_ss( - bmm_dgrad_at_desc, - d_ki_amaj, + bmm_k_inv_da_desc, + d_ki_trans, d_int_da, nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dq_acc_offset), cutlass.Float32), accumulate=True, ) if chunk_idx < FIRST_STATE_CHUNK: mma_ss( - bmm_dgrad_at_desc, - d_ki_amaj, + bmm_k_inv_da_desc, + d_ki_trans, d_int_da, nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dq_acc_offset), cutlass.Float32), accumulate=False, @@ -1450,29 +1387,29 @@ def tcgen05_mma_warp( bars.mb_dq_acc_ready.arrive(cta_group=1) bars.mb_da_done[intermediate_stage].arrive(cta_group=1) - # ---- dK_decay part = state(T) @ (Beta.dY)^T ------------------------------ + # ---- dK decay part = state(T) @ (Beta.dY)^T ------------------------------ bars.mb_dv_tmastg_ready[chunk_serial % cfg.smem_dv_stages].wait((chunk_serial // cfg.smem_dv_stages) % 2) if chunk_idx >= FIRST_STATE_CHUNK: - a_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_state_inp_offset + (chunk_serial % 2) * (cfg.d_v // 2)), cutlass.Int8) - b_desc = d_dv_lead0 + (chunk_serial % cfg.smem_dv_stages) * dv_seg + a_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_state_input_offset + (chunk_serial % 2) * (cfg.d_v // 2)), cutlass.Int8) + b_desc = d_dv0 + (chunk_serial % cfg.smem_dv_stages) * dv_seg c_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dk_decay_acc_offset), cutlass.Float32) - for sub in cutlass.range_constexpr(bmm_dstate_ts_desc.num_subtiles_B): - for k in cutlass.range_constexpr(bmm_dstate_ts_desc.sps_B): + for i in cutlass.range_constexpr(bmm_state_dy_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_state_dy_desc.sps_B): mma_ts_step( - bmm_dstate_ts_desc, - a_ptr.subview(sub * bmm_dstate_ts_desc.sps_B * bmm_dstate_ts_desc.tmem_advance_A), - b_desc + sub * (bmm_dstate_ts_desc.smem_subtile_B >> 4), + bmm_state_dy_desc, + a_ptr.subview(i * bmm_state_dy_desc.sps_B * bmm_state_dy_desc.tmem_advance_A), + b_desc + i * (bmm_state_dy_desc.smem_subtile_B >> 4), c_ptr, k, - cutlass.Boolean(sub + k > 0), + cutlass.Boolean(i + k > 0), ) if elect_one: - bars.mb_state_inp_done[chunk_serial % 2].arrive(cta_group=1) + bars.mb_state_input_done[chunk_serial % 2].arrive(cta_group=1) - # ---- dK_inv part += K_decay^T(S) @ -dM_strict ------------------------ + # ---- dK inv part += K decay^T(S) @ -dM strict ---------------------------- bars.mb_dm_ready[intermediate_stage].wait(intermediate_phase) mma_ss( - bmm_dgrad_at_t_desc, + bmm_k_decay_dm_desc, d_kd_trans, d_int_ndm, nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dk_inv_acc_offset), cutlass.Float32), @@ -1481,19 +1418,19 @@ def tcgen05_mma_warp( if elect_one: bars.mb_dk_inv_part_acc_ready.arrive(cta_group=1) - # ---- dK_decay part += K_inv^T(S) @ dM_strict^T ----------------------- + # ---- dK decay part += K inv^T(S) @ dM strict^T --------------------------- if chunk_idx >= FIRST_STATE_CHUNK: mma_ss( - bmm_dgrad_at_desc, - d_ki_amaj, + bmm_k_inv_dm_desc, + d_ki_trans, d_int_dm, nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dk_decay_acc_offset), cutlass.Float32), accumulate=True, ) if chunk_idx < FIRST_STATE_CHUNK: mma_ss( - bmm_dgrad_at_desc, - d_ki_amaj, + bmm_k_inv_dm_desc, + d_ki_trans, d_int_dm, nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dk_decay_acc_offset), cutlass.Float32), accumulate=False, @@ -1503,11 +1440,12 @@ def tcgen05_mma_warp( bars.mb_dm_done[intermediate_stage].arrive(cta_group=1) bars.mb_decay_done[decay_stage].arrive(cta_group=1) - # ---- tile end: WG1's dstate0 drain gates the next tile's dstate reuse ------------ - bars.mb_dstate0_acc_stored.wait(dstate0_index.phase) - dstate0_index = advance(dstate0_index, 1) + # ---- tile end: WG1's dstate0 store gates the next tile's dstate reuse -------- + if num_compute_chunks > 0: + bars.mb_dstate0_acc_stored.wait(dstate0_index.phase) + dstate0_index = advance(dstate0_index, 1) chunk_serial_base += num_compute_chunks - tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one) bars.mb_tmem_done[0].wait(0) nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_1) nvvm.tcgen05_dealloc( @@ -1525,9 +1463,9 @@ def tmaldg_warp( num_ctas, cu_seqlens, mWorkItems, - mSched, - sSched, - lane, + mScheduler, + sScheduler, + lane_idx, sQ_raw, sK_raw, sV_raw, @@ -1544,10 +1482,17 @@ def tmaldg_warp( ) -> None: """TMA-LDG warp role (warp 14): persistent tile-scheduler loop issuing every G->S TMA load.""" - elect_one = nvvm.elect_sync() - nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) + raw_index = PipelineState.start(phase=1) + state_index = PipelineState.start(phase=1) + scheduler_state = PipelineState.start(phase=1) + tail_count = ((total_tiles - cutlass.Int32(1)) % num_ctas) + cutlass.Int32(1) + tail_base = (total_tiles - tail_count) if tail_count * 2 >= num_ctas else total_tiles + tail_row = tail_base + cute.arch.smid() + tail_row = tail_row if tail_row < total_tiles else cutlass.Int32(1 << 28) + + elect_one = nvvm.elect_sync() sQ_tma = SmemTile( base=sQ_raw, elems_per_stage=(cfg.d_k * cfg.b_t), @@ -1614,14 +1559,15 @@ def tmaldg_warp( tma_granu_elems=64, tma_subtile_stride_elems=cfg.d_k * 64, ) - raw_index = PipelineState.start(phase=1) - state_index = PipelineState.start(phase=1) - sched_state = PipelineState.start(phase=1) tile_idx = cutlass.Int32(bidx) FIRST_STATE_CHUNK = 0 if cfg.use_initial_state else 1 while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) - next_tile, sched_state = sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ctas) + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) + next_tile, scheduler_state = scheduler_publish_next( + cfg, bars, sScheduler, mScheduler, scheduler_state, tile_idx, num_ctas, tail_base, tail_row, elect_one + ) head_o = head_idx head_q = head_idx if cfg.q_ratio == 1 else head_idx // cutlass.Int32(cfg.q_ratio) head_k = head_idx if cfg.k_ratio == 1 else head_idx // cutlass.Int32(cfg.k_ratio) @@ -1640,47 +1586,47 @@ def tmaldg_warp( tma_tensormap_acquire(desc_gate_slot) tma_tensormap_acquire(desc_do_slot) tma_tensormap_acquire(desc_checkpoint_slot) - num_compute_chunks = cend - wstart + num_compute_chunks = compute_end - write_start for rev_idx in cutlass.range(num_compute_chunks, unroll=1): - chunk_idx = cend - cutlass.Int32(1) - rev_idx + chunk_idx = compute_end - cutlass.Int32(1) - rev_idx chunk_start = chunk_idx * cfg.b_t - # ---- Q load ---------------------------------------------------------- + # ---- Q load -------------------------------------------------------------- bars.mb_q_done[raw_index.idx].wait(raw_index.phase) if elect_one: bars.mb_q_ready[raw_index.idx].arrive(n_bytes=cfg.tma_q_bytes) q_slice = tma_slice_runtime_desc(desc_q_slot, cutlass.Int32(0), head_q, chunk_start) tma_load_tile(sQ_tma[raw_index.idx], q_slice, bars.mb_q_ready[raw_index.idx].smem_ptr, acquire=False) - # ---- K load ---------------------------------------------------------- + # ---- K load -------------------------------------------------------------- bars.mb_k_done[raw_index.idx].wait(raw_index.phase) if elect_one: bars.mb_k_ready[raw_index.idx].arrive(n_bytes=cfg.tma_k_bytes) k_slice = tma_slice_runtime_desc(desc_k_slot, cutlass.Int32(0), head_k, chunk_start) tma_load_tile(sK_tma[raw_index.idx], k_slice, bars.mb_k_ready[raw_index.idx].smem_ptr, acquire=False) - # ---- Gate load ------------------------------------------------------- + # ---- Gate load: GMEM -> SMEM --------------------------------------------- bars.mb_gate_done[raw_index.idx].wait(raw_index.phase) if elect_one: bars.mb_gate_ready[raw_index.idx].arrive(n_bytes=cfg.tma_gate_bytes) gate_slice = tma_slice_runtime_desc(desc_gate_slot, cutlass.Int32(0), head_o, chunk_start) tma_load_tile(sGate_tma[raw_index.idx], gate_slice, bars.mb_gate_ready[raw_index.idx].smem_ptr, acquire=False) - # ---- dO load --------------------------------------------------------- + # ---- dO load ------------------------------------------------------------- bars.mb_do_done[raw_index.idx].wait(raw_index.phase) if elect_one: bars.mb_do_ready[raw_index.idx].arrive(n_bytes=cfg.tma_do_bytes) do_slice = tma_slice_runtime_desc(desc_do_slot, cutlass.Int32(0), head_o, chunk_start) tma_load_tile(sDo_tma[raw_index.idx], do_slice, bars.mb_do_ready[raw_index.idx].smem_ptr, acquire=False) - # ---- V load ---------------------------------------------------------- + # ---- V load -------------------------------------------------------------- bars.mb_v_done[raw_index.idx].wait(raw_index.phase) if elect_one: bars.mb_v_ready[raw_index.idx].arrive(n_bytes=cfg.tma_v_bytes) v_slice = tma_slice_runtime_desc(desc_v_slot, cutlass.Int32(0), head_v, chunk_start) tma_load_tile(sV_tma[raw_index.idx], v_slice, bars.mb_v_ready[raw_index.idx].smem_ptr, acquire=False) - # ---- entering state ---- + # ---- entering state ------------------------------------------------------ if chunk_idx >= FIRST_STATE_CHUNK: state_idx = state_index.idx bars.mb_state_cg0_done[state_idx].wait(state_index.phase) @@ -1699,9 +1645,7 @@ def gate_scale(cfg, raw_gate: cutlass.Float32) -> cutlass.Float32: """Map raw gate to the log2-domain decay increment used by KDA.""" if cutlass.const_expr(cfg.safe_gate): - half = cutlass.Float32(0.5) - sigmoid = cute.math.tanh(raw_gate * half, approx=True) * half + half - return cfg.gate_scale_log2 * sigmoid + return cfg.gate_scale_log2 * sigmoid(raw_gate) # Default ABI: Gate arrives in natural-log space return raw_gate * cutlass.Float32(LOG2_E) @@ -1714,8 +1658,8 @@ def compute0_warp_group( num_ctas, cu_seqlens, mWorkItems, - sSched, - lane, + sScheduler, + lane_idx, tmem_base_holder, warp_idx, scale, @@ -1741,8 +1685,9 @@ def compute0_warp_group( stashes the per-row Q/K inverse norms for WG2's dGate assembly and copies H -> TMEM f16 at the chunk tail.""" nvvm.setmaxregister(cfg.num_regs_compute_group_0, nvvm.SetMaxRegisterAction.INCREASE) + elect_one = nvvm.elect_sync() cg0_warp = warp_idx - cfg.compute_group_0_warp_ids[0] - prefix_dim = cg0_warp * cfg.threads_per_warp + lane + channel_dim = cg0_warp * cfg.threads_per_warp + lane_idx cg0_a_log_exp = cutlass.Float32(1.0) cg0_dt_bias_value = cutlass.Float32(0.0) nvvm.barrier_cta_sync(cfg.tmem_lifecycle_barrier_id, thread_count=cfg.tmem_user_threads) @@ -1750,22 +1695,24 @@ def compute0_warp_group( tmem_col = tmem_base & 0xFFFF tmem_row = tmem_base >> 16 tmem_subpartition = warp_idx % (cfg.d_v // cfg.threads_per_warp) - value_dim = tmem_subpartition * cfg.threads_per_warp + lane - row_addr = (tmem_row + tmem_subpartition * cfg.threads_per_warp) << 16 + value_dim = tmem_subpartition * cfg.threads_per_warp + lane_idx + row_lo_addr = tmem_row << 16 state_index = PipelineState.start(phase=0) chunk_serial_base = cutlass.Int32(0) - sched_state = PipelineState.start(phase=0) + scheduler_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) FIRST_STATE_CHUNK = 0 if cfg.use_initial_state else 1 while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) - num_compute_chunks = cend - wstart + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) + num_compute_chunks = compute_end - write_start if cutlass.const_expr(cfg.safe_gate): if num_compute_chunks > 0: cg0_a_log_exp = cute.math.exp2(mA_log[head_idx].to(cutlass.Float32) * LOG2_E, fastmath=True) - cg0_dt_bias_value = mDt_bias[head_idx, prefix_dim].to(cutlass.Float32) + cg0_dt_bias_value = mDt_bias[head_idx, channel_dim].to(cutlass.Float32) for rev_idx in cutlass.range(num_compute_chunks, unroll=1): - chunk_idx = cend - cutlass.Int32(1) - rev_idx + chunk_idx = compute_end - cutlass.Int32(1) - rev_idx chunk_serial = chunk_serial_base + rev_idx chunk_start = chunk_idx * cfg.b_t decay_stage = chunk_serial % cfg.smem_decay_stages @@ -1779,19 +1726,18 @@ def compute0_warp_group( sK_restore_ptr = sK_restore_raw.data_ptr() + decay_stage * (cfg.d_k * cfg.b_t) sState_scale_diag_ptr = sState_scale_diag_raw.data_ptr() + decay_stage * ((cfg.d_k // 16) * 256) - # ---- beta scalars: gathered in the inputs-wait shadow ---------------- + # ---- beta scalars: gathered in the inputs-wait shadow -------------------- if cg0_warp == 0: beta_stage = chunk_serial % cfg.smem_beta_stages bars.mb_beta_done[beta_stage].wait(((chunk_serial // cfg.smem_beta_stages) + 1) % 2) - if lane < cfg.b_t: - token_idx = chunk_start + lane + if lane_idx < cfg.b_t: + token_idx = chunk_start + lane_idx beta_value = cutlass.Float32(0.0) - if token_idx < seqlen_b: + if token_idx < batch_seqlen: beta_value = mBeta[batch_start + token_idx, head_idx].to(cutlass.Float32) if cutlass.const_expr(cfg.beta_sigmoid): - half = cutlass.Float32(0.5) - beta_value = (cute.math.tanh(beta_value * half, approx=True) * half + half).to(mBeta.element_type).to(cutlass.Float32) - sBeta_raw[beta_stage * cfg.b_t + lane] = beta_value + beta_value = sigmoid(beta_value).to(mBeta.element_type).to(cutlass.Float32) + sBeta_raw[beta_stage * cfg.b_t + lane_idx] = beta_value bars.mb_beta_ready[beta_stage].arrive() bars.mb_gate_ready[raw_stage].wait((chunk_serial // cfg.smem_raw_stages) % 2) @@ -1799,45 +1745,37 @@ def compute0_warp_group( bars.mb_k_ready[raw_stage].wait((chunk_serial // cfg.smem_raw_stages) % 2) row_group_start = cg0_warp * (cfg.b_t // len(cfg.compute_group_0_warp_ids)) - lane_row_group = lane // 8 - lane_in_row_group = lane - lane_row_group * 8 + lane_row_group = lane_idx // 8 + lane_in_row_group = lane_idx - lane_row_group * 8 decay_row = row_group_start + lane_row_group g_prefix_ptr = sGate_ptr - prefix_dim = cg0_warp * cfg.threads_per_warp + lane - # ---- gate prefix scan: cumulative log-gate per key channel ----------- + channel_dim = cg0_warp * cfg.threads_per_warp + lane_idx + # ---- gate prefix scan: cumulative log-gate per key channel --------------- gate_raw = cutlass.Array(cutlass.Float32, cfg.b_t, alignment=16) for row in cutlass.range_constexpr(cfg.b_t): - f32_segment = prefix_dim // 32 - f32_segment_dim = prefix_dim - f32_segment * 32 + f32_segment = channel_dim // 32 + f32_segment_dim = channel_dim - f32_segment * 32 prefix_idx = f32_segment * (cfg.b_t * 32) + row * 32 + swizzle_xor_128b(row, f32_segment_dim, elem_bytes=4) gate_raw[row] = (sGate_ptr + prefix_idx).load() g_prefix_regs = cutlass.Array(cutlass.Float32, cfg.b_t, alignment=16) if cutlass.const_expr(cfg.safe_gate): - valid_rows = seqlen_b - chunk_idx * cutlass.Int32(cfg.b_t) - valid_mask = cutlass.vector.create_mask([cfg.b_t], [valid_rows]) - for row_pair in cutlass.range_constexpr(cfg.b_t // 2): - row0 = row_pair * 2 - row1 = row0 + 1 - gate0 = cg0_a_log_exp * (gate_raw[row0] + cg0_dt_bias_value) - gate1 = cg0_a_log_exp * (gate_raw[row1] + cg0_dt_bias_value) - gate0 = gate_scale( - cfg, - gate0, - ) - gate1 = gate_scale( - cfg, - gate1, - ) - gate_pair = cutlass.Vector.from_elements((gate0, gate1), cutlass.Float32) - gate_pair = cutlass.vector.where(valid_mask[row0 : row1 + 1], gate_pair, 0.0) - g_prefix_regs[row0] = gate_pair[0] - g_prefix_regs[row1] = gate_pair[1] + for row in cutlass.range_constexpr(cfg.b_t): + gate = gate_raw[row] + token_idx = chunk_idx * cutlass.Int32(cfg.b_t) + cutlass.Int32(row) + if token_idx < batch_seqlen: + gate = gate_scale( + cfg, + cg0_a_log_exp * (gate + cg0_dt_bias_value), + ) + else: + gate = cutlass.Float32(0.0) + g_prefix_regs[row] = gate else: for row in cutlass.range_constexpr(cfg.b_t): gate = gate_raw[row] token_idx = chunk_idx * cutlass.Int32(cfg.b_t) + cutlass.Int32(row) - if token_idx < seqlen_b: + if token_idx < batch_seqlen: gate = gate_scale( cfg, gate, @@ -1868,28 +1806,27 @@ def compute0_warp_group( g_prefix_regs[row] = cute.math.exp2(g_prefix_regs[row], fastmath=True) exp_g_last = g_prefix_regs[cfg.b_t - 1] - # ---- decay-slot guard: previous use fully consumed ------------------- + # ---- decay-slot guard: previous use fully consumed ----------------------- operand_done_phase = ((chunk_serial // cfg.smem_decay_stages) + 1) % 2 bars.mb_decay_done[decay_stage].wait(operand_done_phase) for row in cutlass.range_constexpr(cfg.b_t): - f32_segment = prefix_dim // 32 - f32_segment_dim = prefix_dim - f32_segment * 32 + f32_segment = channel_dim // 32 + f32_segment_dim = channel_dim - f32_segment * 32 prefix_idx = f32_segment * (cfg.b_t * 32) + row * 32 + swizzle_xor_128b(row, f32_segment_dim, elem_bytes=4) (sGate_ptr + prefix_idx).store(g_prefix_regs[row]) - # ---- state-scale diag: stage exp2(g_last) decay blocks --------------- - block = prefix_dim // cutlass.Int32(16) - coord = prefix_dim - block * cutlass.Int32(16) - linear_idx = block * cutlass.Int32(256) + coord * cutlass.Int32(16) + coord - diag_idx = swizzle_lin_S(linear_idx, bbits=1, mbase=3, sshift=3) + # ---- state-scale diag: stage exp2(g last) decay blocks ------------------- + block = channel_dim // cutlass.Int32(16) + coord = channel_dim - block * cutlass.Int32(16) + diag_idx = block * cutlass.Int32(256) + coord * cutlass.Int32(16) + swizzle_xor_32b(channel_dim, coord) sState_scale_diag_ptr[diag_idx] = exp_g_last.to(cfg.io_dtype) - # ---- raw Q/K: SMEM -> TMEM ring (channel-major, for WG2) ------------- + # ---- raw Q/K: SMEM -> TMEM ring (channel-major, for WG2) ----------------- qk_raw_stage = chunk_serial % cfg.tmem_qk_raw_stages bars.mb_qk_raw_done[qk_raw_stage].wait(((chunk_serial // cfg.tmem_qk_raw_stages) + 1) % 2) - raw_seg = prefix_dim // 64 - raw_dim = prefix_dim - raw_seg * 64 + raw_seg = channel_dim // 64 + raw_dim = channel_dim - raw_seg * 64 q_raw_words = cutlass.Array(cutlass.Int32, cfg.b_t // 2, alignment=16) k_raw_words = cutlass.Array(cutlass.Int32, cfg.b_t // 2, alignment=16) for t2 in cutlass.range_constexpr(cfg.b_t // 2): @@ -1904,12 +1841,12 @@ def compute0_warp_group( k_raw_words[t2] = fp32_to_fp16(k0, k1, dtype=cfg.io_dtype) nvvm.tcgen05_st( "32x32b", - nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_qraw_inp_offset + qk_raw_stage * (cfg.b_t // 2)), cutlass.Int8), + nvvm.make_tmem_ptr(row_lo_addr + (tmem_col + cfg.tmem_qraw_input_offset + qk_raw_stage * (cfg.b_t // 2)), cutlass.Int8), q_raw_words[0 : (cfg.b_t // 2)], ) nvvm.tcgen05_st( "32x32b", - nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_kraw_inp_offset + qk_raw_stage * (cfg.b_t // 2)), cutlass.Int8), + nvvm.make_tmem_ptr(row_lo_addr + (tmem_col + cfg.tmem_kraw_input_offset + qk_raw_stage * (cfg.b_t // 2)), cutlass.Int8), k_raw_words[0 : (cfg.b_t // 2)], ) nvvm.tcgen05_wait("store") @@ -1920,7 +1857,7 @@ def compute0_warp_group( k_inv_pack = cutlass.Array(cutlass.Int32, 2 * 4, alignment=16) raw_q_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) raw_k_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) - # ---- optional Q/K L2-norm -------------------------------------------- + # ---- optional Q/K L2-norm ------------------------------------------------ if cutlass.const_expr(cfg.l2norm): qk0_lo = opaque_f32_zero() qk0_hi = opaque_f32_zero() @@ -1970,7 +1907,7 @@ def compute0_warp_group( sNorm_raw[(chunk_serial % cfg.tmem_qk_raw_stages) * (2 * cfg.b_t) + cfg.b_t + decay_row] = k_inv_norm q_stage_norm = q_inv_norm * scale - # ---- decay/restore operands: exp2(+-g) applied per key channel ------- + # ---- decay/restore operands: exp2(+-g) applied per key channel ----------- exp_g_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) exp_g_last_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) for dim_half in cutlass.range_constexpr(2): @@ -1999,7 +1936,7 @@ def compute0_warp_group( for dim_half in cutlass.range_constexpr(2): dim_base = dim_half * (cfg.d_k // 2) + lane_in_row_group * 8 reg_base = dim_half * 8 - # ---- K decay + K inv operands: exp2(+g) * K / exp2(-g) * K ------- + # ---- K decay + K inv operands: exp2(+g) * K / exp2(-g) * K ----------- k_decay_pack = cutlass.Array(cutlass.Int32, 4, alignment=16) for pair_idx in cutlass.range_constexpr(4): dim0 = pair_idx * 2 @@ -2036,7 +1973,7 @@ def compute0_warp_group( nvvm.fence_proxy("async.shared", space="cta") bars.mb_k_decay_inv_ready[decay_stage].arrive() - # ---- Q decay + K_restore operands ------------------------------------ + # ---- Q decay + K restore operands ---------------------------------------- for dim_half in cutlass.range_constexpr(2): dim_base = dim_half * (cfg.d_k // 2) + lane_in_row_group * 8 reg_base = dim_half * 8 @@ -2071,38 +2008,37 @@ def compute0_warp_group( bars.mb_q_decay_k_restore_ready[decay_stage].arrive() # ---- state copy: SMEM -> TMEM f16 ---------------------------------------- - bars.mb_state_inp_done[chunk_serial % 2].wait(((chunk_serial // 2) + 1) % 2) - bars.mb_state_inp_cg2_done[chunk_serial % 2].wait(((chunk_serial // 2) + 1) % 2) + bars.mb_state_input_done[chunk_serial % 2].wait(((chunk_serial // 2) + 1) % 2) + bars.mb_state_input_cg2_done[chunk_serial % 2].wait(((chunk_serial // 2) + 1) % 2) if chunk_idx >= FIRST_STATE_CHUNK: bars.mb_state_ready[state_index.idx].wait(state_index.phase) state_src = sState_raw.data_ptr() + state_index.idx * (cfg.d_k * cfg.d_v) - for v_seg in cutlass.range_constexpr(2): - for v_col8 in cutlass.range_constexpr(8): - state_frag = cutlass.Vector.from_elements( - tuple( - ( - state_src - + (value_dim // 64) * (cfg.d_v * 64) - + (v_seg * 64 + v_col8 * 8 + e) * 64 - + swizzle_xor_128b(v_seg * 64 + v_col8 * 8 + e, value_dim % 64, elem_bytes=2) - ).load() - for e in range(8) - ), - cfg.io_dtype, - ) - nvvm.tcgen05_st( - "32x32b", - nvvm.make_tmem_ptr( - row_addr + (tmem_col + cfg.tmem_state_inp_offset + (chunk_serial % 2) * (cfg.d_v // 2) + v_seg * 32 + v_col8 * 4), cutlass.Int8 - ), - state_frag.bitcast(cutlass.Int32), - ) + ldm_row_coord = (lane_idx // 16) * 8 + (lane_idx & 7) + ldm_col_offset = ((lane_idx // 8) & 1) * 8 + k_base = tmem_subpartition * cfg.threads_per_warp + k_seg_off = (k_base // 64) * (cfg.d_v * 64) + row_hi_addr = row_lo_addr + (16 << 16) + state_col = tmem_col + cfg.tmem_state_input_offset + (chunk_serial % 2) * (cfg.d_v // 2) + for dv_blk in cutlass.range_constexpr(cfg.d_v // 16): + dv_row = dv_blk * 16 + ldm_row_coord + frag_lo = nvvm.ldmatrix( + state_src + k_seg_off + dv_row * 64 + swizzle_xor_128b(dv_row, (k_base + ldm_col_offset) % 64, elem_bytes=2), + 4, + nvvm.MMALayout.COL, + ) + frag_hi = nvvm.ldmatrix( + state_src + k_seg_off + dv_row * 64 + swizzle_xor_128b(dv_row, (k_base + 16 + ldm_col_offset) % 64, elem_bytes=2), + 4, + nvvm.MMALayout.COL, + ) + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_lo_addr + (state_col + dv_blk * 8), cutlass.Int8), frag_lo) + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_hi_addr + (state_col + dv_blk * 8), cutlass.Int8), frag_hi) nvvm.tcgen05_wait("store") bars.mb_state_cg0_done[state_index.idx].arrive() state_index = advance(state_index, cfg.smem_state_stages) - bars.mb_state_inp_ready[chunk_serial % 2].arrive() + bars.mb_state_input_ready[chunk_serial % 2].arrive() chunk_serial_base += num_compute_chunks - tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one) @cute.jit @@ -2113,8 +2049,8 @@ def compute1_warp_group( num_ctas, cu_seqlens, mWorkItems, - sSched, - lane, + sScheduler, + lane_idx, tmem_base_holder, warp_idx, mDbeta, @@ -2133,16 +2069,17 @@ def compute1_warp_group( ) -> None: """WG1 warp role (warps 4-7): the value-side TMEM staging.""" nvvm.setmaxregister(cfg.num_regs_compute_group_1, nvvm.SetMaxRegisterAction.INCREASE) + elect_one = nvvm.elect_sync() nvvm.barrier_cta_sync(cfg.tmem_lifecycle_barrier_id, thread_count=cfg.tmem_user_threads) tmem_base = tmem_base_holder.load() tmem_col = tmem_base & 0xFFFF tmem_row = tmem_base >> 16 tmem_subpartition = warp_idx % (cfg.d_v // cfg.threads_per_warp) - token_row_coord = (lane // 16) * 8 + (lane & 7) - value_col_offset = ((lane // 8) & 1) * 8 - value_dim = tmem_subpartition * cfg.threads_per_warp + lane + ov_row_coord = (lane_idx // 16) * 8 + (lane_idx & 7) + ov_col_offset = ((lane_idx // 8) & 1) * 8 + value_dim = tmem_subpartition * cfg.threads_per_warp + lane_idx value_dim_base = tmem_subpartition * cfg.threads_per_warp - cg1_tidx = warp_idx % 4 * cfg.threads_per_warp + lane + cg1_tidx = warp_idx % 4 * cfg.threads_per_warp + lane_idx raw_index = PipelineState.start(phase=0) state_k_index = PipelineState.start(phase=0) @@ -2155,32 +2092,34 @@ def compute1_warp_group( dv_done_index = PipelineState.start(phase=1) chunk_serial_base = cutlass.Int32(0) - sched_state = PipelineState.start(phase=0) + scheduler_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) FIRST_STATE_CHUNK = 0 if cfg.use_initial_state else 1 while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) - num_compute_chunks = cend - wstart + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) + num_compute_chunks = compute_end - write_start - # ---- dht seeding: dh acc + dh_inp f16 + sdH ------------------------------ + # ---- dstate seed: GMEM -> TMEM + SMEM ---------------------------------------- if cutlass.const_expr(cfg.use_dstate_in): if num_compute_chunks > 0: - seed_true = cend == num_chunks_b + seed_true = compute_end == batch_num_chunks bars.mb_dstate_smem_done.wait(dstate_smem_done_index.phase) bars.mb_dstate_smem_cg2_done.wait(dstate_smem_done_index.phase) dstate_smem_done_index = advance(dstate_smem_done_index, 1) - row_addr = (tmem_row + tmem_subpartition * cfg.threads_per_warp) << 16 + row_lo_addr = tmem_row << 16 dstate_src = (mDstate_in.iterator + mDstate_in.layout((batch_idx, head_idx, value_dim, 0))).raw_ptr() - for sub in cutlass.range_constexpr(cfg.d_k // 16): + for i in cutlass.range_constexpr(cfg.d_k // 16): seed_block = cutlass.Array(cutlass.Float32, 16, alignment=16) for g in cutlass.range_constexpr(4): - seed_chunk = (dstate_src + sub * 16 + g * 4).load(count=4, alignment=16) + seed_chunk = (dstate_src + i * 16 + g * 4).load(count=4, alignment=16) for t in cutlass.range_constexpr(4): dval = seed_chunk[t].to(cutlass.Float32) seed_block[g * 4 + t] = dval if seed_true else cutlass.Float32(0.0) nvvm.tcgen05_st( "32x32b", - nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_dstate_acc_offset + sub * 16), cutlass.Float32), + nvvm.make_tmem_ptr(row_lo_addr + (tmem_col + cfg.tmem_dstate_acc_offset + i * 16), cutlass.Float32), seed_block[0:16], ) seed_pack = cutlass.Array(cutlass.Int32, 8, alignment=16) @@ -2188,19 +2127,19 @@ def compute1_warp_group( seed_pack[pc] = fp32_to_fp16(seed_block[2 * pc], seed_block[2 * pc + 1], dtype=cfg.io_dtype) nvvm.tcgen05_st( "32x32b", - nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_dstate_inp_offset + sub * 8), cutlass.Int8), + nvvm.make_tmem_ptr(row_lo_addr + (tmem_col + cfg.tmem_dstate_input_offset + i * 8), cutlass.Int8), seed_pack[0:8], ) nvvm.tcgen05_wait("store") - bars.mb_dstate_inp_ready.arrive() + bars.mb_dstate_input_ready.arrive() - # ---- dht seed -> sdH: re-read dh_inp after the TMEM publish ------ - for sub in cutlass.range_constexpr(cfg.d_k // 16): + # ---- dstate seed -> sdH: re-read after the TMEM publish -------------- + for i in cutlass.range_constexpr(cfg.d_k // 16): dstate_words = nvvm.tcgen05_ld( - "32x32b", nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_dstate_inp_offset + sub * 8), cutlass.Float32), num=8 + "32x32b", nvvm.make_tmem_ptr(row_lo_addr + (tmem_col + cfg.tmem_dstate_input_offset + i * 8), cutlass.Float32), num=8 ) for half in cutlass.range_constexpr(2): - d_base = sub * 16 + half * 8 + d_base = i * 16 + half * 8 h_vec = cutlass.Vector.from_elements( (dstate_words[half * 4], dstate_words[half * 4 + 1], dstate_words[half * 4 + 2], dstate_words[half * 4 + 3]), cutlass.Float32, @@ -2211,254 +2150,245 @@ def compute1_warp_group( bars.mb_dstate_smem_ready.arrive() for rev_idx in cutlass.range(num_compute_chunks, unroll=1): - chunk_idx = cend - cutlass.Int32(1) - rev_idx + chunk_idx = compute_end - cutlass.Int32(1) - rev_idx chunk_serial = chunk_serial_base + rev_idx sV_ptr = sV_raw.data_ptr() + raw_index.idx * (cfg.d_v * cfg.b_t) sDo_ptr = sDo_raw.data_ptr() + raw_index.idx * (cfg.d_v * cfg.b_t) sBeta_ptr = sBeta_raw.data_ptr() + (chunk_serial % cfg.smem_beta_stages) * cfg.b_t - row_addr_lo = tmem_row << 16 - row_addr_hi = (tmem_row + 16) << 16 - row_id0 = tmem_row + value_dim_base - row_id1 = row_id0 + 16 + row_lo_addr = tmem_row << 16 + row_hi_addr = (tmem_row + 16) << 16 has_dstate = cutlass.Boolean(rev_idx > 0) if cutlass.const_expr(cfg.use_dstate_in): has_dstate = cutlass.Boolean(True) - # ---- Y staging: Y = Beta * (V - state_k) ----------------------------- + # ---- Y stage: Y = Beta * (V - k state) ----------------------------------- bars.mb_v_ready[raw_index.idx].wait((chunk_serial // cfg.smem_raw_stages) % 2) projection_col_id = tmem_col + cfg.tmem_state_k_acc_offset - input_col_id = tmem_col + cfg.tmem_y_inp_offset - raw_v_frag0 = nvvm.ldmatrix( + input_col_id = tmem_col + cfg.tmem_y_input_offset + raw_v_frag_lo = nvvm.ldmatrix( sV_ptr - + (value_dim_base + value_col_offset) // 64 * (cfg.b_t * 64) - + token_row_coord * 64 - + swizzle_xor_128b(token_row_coord, (value_dim_base + value_col_offset) % 64, elem_bytes=2), + + (value_dim_base + ov_col_offset) // 64 * (cfg.b_t * 64) + + ov_row_coord * 64 + + swizzle_xor_128b(ov_row_coord, (value_dim_base + ov_col_offset) % 64, elem_bytes=2), 4, nvvm.MMALayout.COL, ) - raw_v_frag1 = nvvm.ldmatrix( + raw_v_frag_hi = nvvm.ldmatrix( sV_ptr - + (value_dim_base + 16 + value_col_offset) // 64 * (cfg.b_t * 64) - + token_row_coord * 64 - + swizzle_xor_128b(token_row_coord, (value_dim_base + 16 + value_col_offset) % 64, elem_bytes=2), + + (value_dim_base + 16 + ov_col_offset) // 64 * (cfg.b_t * 64) + + ov_row_coord * 64 + + swizzle_xor_128b(ov_row_coord, (value_dim_base + 16 + ov_col_offset) % 64, elem_bytes=2), 4, nvvm.MMALayout.COL, ) - # fence: the V ldmatrix reads must complete before this release - # licenses the TMA reload (sV) nvvm.fence_proxy("async.shared", space="cta") bars.mb_v_done[raw_index.idx].arrive() bars.mb_beta_ready[chunk_serial % cfg.smem_beta_stages].wait((chunk_serial // cfg.smem_beta_stages) % 2) beta_pairs = cutlass.Array(cutlass.Int32, 2, space=cutlass.AddressSpace.rmem) for half in cutlass.range_constexpr(2): - token0 = ((half * 4 + (lane & 3)) ^ 4) * 2 + token0 = (half * 4 + (lane_idx & 3)) * 2 beta0 = (sBeta_ptr + token0).load().to(cutlass.Float32) beta1 = (sBeta_ptr + token0 + 1).load().to(cutlass.Float32) beta_pairs[half] = fp32_to_fp16(beta0, beta1, dtype=cfg.io_dtype) - diff_w0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) - diff_w1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) - y_inp_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) - y_inp_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + diff_w_lo = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + diff_w_hi = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + y_input_pack_lo = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + y_input_pack_hi = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) if chunk_idx >= FIRST_STATE_CHUNK: bars.mb_state_k_acc_ready.wait(state_k_index.phase) state_k_index = advance(state_k_index, 1) - state_k_vec0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id0 << 16) + projection_col_id, cutlass.Float32), num=2) - state_k_vec1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id1 << 16) + projection_col_id, cutlass.Float32), num=2) + state_k_vec_lo = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_lo_addr + projection_col_id, cutlass.Float32), num=2) + state_k_vec_hi = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_hi_addr + projection_col_id, cutlass.Float32), num=2) for reg_idx in cutlass.range_constexpr(4): - raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1) - frag_pair = (reg_idx ^ 2) * 2 - state_k_pair = fp32_to_fp16(state_k_vec0[frag_pair], state_k_vec0[frag_pair + 1], dtype=cfg.io_dtype) - diff_pair = sub_f16x2(raw_v_frag0[raw_matrix], state_k_pair, cfg.io_dtype) - diff_w0[reg_idx ^ 2] = diff_pair - y_inp_pack0[reg_idx ^ 2] = mul_f16x2(beta_pairs[reg_idx // 2], diff_pair, cfg.io_dtype) + frag_pair = reg_idx * 2 + state_k_pair = fp32_to_fp16(state_k_vec_lo[frag_pair], state_k_vec_lo[frag_pair + 1], dtype=cfg.io_dtype) + diff_pair = sub_f16x2(raw_v_frag_lo[reg_idx], state_k_pair, cfg.io_dtype) + diff_w_lo[reg_idx] = diff_pair + y_input_pack_lo[reg_idx] = mul_f16x2(beta_pairs[reg_idx // 2], diff_pair, cfg.io_dtype) for reg_idx in cutlass.range_constexpr(4): - raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1) - frag_pair = (reg_idx ^ 2) * 2 - state_k_pair = fp32_to_fp16(state_k_vec1[frag_pair], state_k_vec1[frag_pair + 1], dtype=cfg.io_dtype) - diff_pair = sub_f16x2(raw_v_frag1[raw_matrix], state_k_pair, cfg.io_dtype) - diff_w1[reg_idx ^ 2] = diff_pair - y_inp_pack1[reg_idx ^ 2] = mul_f16x2(beta_pairs[reg_idx // 2], diff_pair, cfg.io_dtype) + frag_pair = reg_idx * 2 + state_k_pair = fp32_to_fp16(state_k_vec_hi[frag_pair], state_k_vec_hi[frag_pair + 1], dtype=cfg.io_dtype) + diff_pair = sub_f16x2(raw_v_frag_hi[reg_idx], state_k_pair, cfg.io_dtype) + diff_w_hi[reg_idx] = diff_pair + y_input_pack_hi[reg_idx] = mul_f16x2(beta_pairs[reg_idx // 2], diff_pair, cfg.io_dtype) if chunk_idx < FIRST_STATE_CHUNK: for reg_idx in cutlass.range_constexpr(4): - raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1) - diff_w0[reg_idx ^ 2] = raw_v_frag0[raw_matrix] - y_inp_pack0[reg_idx ^ 2] = mul_f16x2(beta_pairs[reg_idx // 2], raw_v_frag0[raw_matrix], cfg.io_dtype) + diff_w_lo[reg_idx] = raw_v_frag_lo[reg_idx] + y_input_pack_lo[reg_idx] = mul_f16x2(beta_pairs[reg_idx // 2], raw_v_frag_lo[reg_idx], cfg.io_dtype) for reg_idx in cutlass.range_constexpr(4): - raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1) - diff_w1[reg_idx ^ 2] = raw_v_frag1[raw_matrix] - y_inp_pack1[reg_idx ^ 2] = mul_f16x2(beta_pairs[reg_idx // 2], raw_v_frag1[raw_matrix], cfg.io_dtype) - nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_addr_lo + input_col_id, cutlass.Int8), y_inp_pack0[0:4]) - nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_addr_hi + input_col_id, cutlass.Int8), y_inp_pack1[0:4]) + diff_w_hi[reg_idx] = raw_v_frag_hi[reg_idx] + y_input_pack_hi[reg_idx] = mul_f16x2(beta_pairs[reg_idx // 2], raw_v_frag_hi[reg_idx], cfg.io_dtype) + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_lo_addr + input_col_id, cutlass.Int8), y_input_pack_lo[0:4]) + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_hi_addr + input_col_id, cutlass.Int8), y_input_pack_hi[0:4]) nvvm.tcgen05_wait("store") - bars.mb_y_inp_ready.arrive() + bars.mb_y_input_ready.arrive() - # ---- dU restage: dU acc -> TMEM f16 A operand ------------------------ + # ---- dU stage: dU acc -> TMEM f16 A operand ------------------------------ bars.mb_du_acc_ready.wait(du_acc_index.phase) du_acc_index = advance(du_acc_index, 1) du_col_id = tmem_col + cfg.tmem_du_acc_offset - du_vec0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id0 << 16) + du_col_id, cutlass.Float32), num=2) - du_vec1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id1 << 16) + du_col_id, cutlass.Float32), num=2) + du_vec_lo = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_lo_addr + du_col_id, cutlass.Float32), num=2) + du_vec_hi = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_hi_addr + du_col_id, cutlass.Float32), num=2) - du_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) - du_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + du_pack_lo = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + du_pack_hi = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) for reg_idx in cutlass.range_constexpr(4): frag_pair = reg_idx * 2 - du_pack0[reg_idx] = fp32_to_fp16(du_vec0[frag_pair], du_vec0[frag_pair + 1], dtype=cfg.io_dtype) - du_pack1[reg_idx] = fp32_to_fp16(du_vec1[frag_pair], du_vec1[frag_pair + 1], dtype=cfg.io_dtype) - nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_addr_lo + (tmem_col + cfg.tmem_du_inp_offset), cutlass.Int8), du_pack0[0:4]) - nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_addr_hi + (tmem_col + cfg.tmem_du_inp_offset), cutlass.Int8), du_pack1[0:4]) + du_pack_lo[reg_idx] = fp32_to_fp16(du_vec_lo[frag_pair], du_vec_lo[frag_pair + 1], dtype=cfg.io_dtype) + du_pack_hi[reg_idx] = fp32_to_fp16(du_vec_hi[frag_pair], du_vec_hi[frag_pair + 1], dtype=cfg.io_dtype) + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_lo_addr + (tmem_col + cfg.tmem_du_input_offset), cutlass.Int8), du_pack_lo[0:4]) + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_hi_addr + (tmem_col + cfg.tmem_du_input_offset), cutlass.Int8), du_pack_hi[0:4]) nvvm.tcgen05_wait("store") - bars.mb_du_inp_ready.arrive() + bars.mb_du_input_ready.arrive() - # ---- U readback -> sU ------------------------------------------------ + # ---- U read: TMEM -> sU -------------------------------------------------- bars.mb_u_acc_ready.wait(u_acc_index.phase) u_acc_index = advance(u_acc_index, 1) u_col_id = tmem_col + cfg.tmem_u_acc_offset - u_vec0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id0 << 16) + u_col_id, cutlass.Float32), num=2) - u_vec1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id1 << 16) + u_col_id, cutlass.Float32), num=2) + u_vec_lo = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_lo_addr + u_col_id, cutlass.Float32), num=2) + u_vec_hi = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_hi_addr + u_col_id, cutlass.Float32), num=2) - u_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) - u_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + u_pack_lo = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + u_pack_hi = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) for reg_idx in cutlass.range_constexpr(4): - u_pack0[reg_idx] = fp32_to_fp16(u_vec0[2 * reg_idx], u_vec0[2 * reg_idx + 1], dtype=cfg.io_dtype) - u_pack1[reg_idx] = fp32_to_fp16(u_vec1[2 * reg_idx], u_vec1[2 * reg_idx + 1], dtype=cfg.io_dtype) + u_pack_lo[reg_idx] = fp32_to_fp16(u_vec_lo[2 * reg_idx], u_vec_lo[2 * reg_idx + 1], dtype=cfg.io_dtype) + u_pack_hi[reg_idx] = fp32_to_fp16(u_vec_hi[2 * reg_idx], u_vec_hi[2 * reg_idx + 1], dtype=cfg.io_dtype) nvvm.stmatrix( sU_raw.data_ptr() - + (value_dim_base + value_col_offset) // 64 * (cfg.b_t * 64) - + token_row_coord * 64 - + swizzle_xor_128b(token_row_coord, (value_dim_base + value_col_offset) % 64, elem_bytes=2), - u_pack0.data_ptr().load(count=4, alignment=4), + + (value_dim_base + ov_col_offset) // 64 * (cfg.b_t * 64) + + ov_row_coord * 64 + + swizzle_xor_128b(ov_row_coord, (value_dim_base + ov_col_offset) % 64, elem_bytes=2), + u_pack_lo.data_ptr().load(count=4, alignment=4), nvvm.MMALayout.COL, shape=nvvm.StoreShape.M8N8, ) nvvm.stmatrix( sU_raw.data_ptr() - + (value_dim_base + 16 + value_col_offset) // 64 * (cfg.b_t * 64) - + token_row_coord * 64 - + swizzle_xor_128b(token_row_coord, (value_dim_base + 16 + value_col_offset) % 64, elem_bytes=2), - u_pack1.data_ptr().load(count=4, alignment=4), + + (value_dim_base + 16 + ov_col_offset) // 64 * (cfg.b_t * 64) + + ov_row_coord * 64 + + swizzle_xor_128b(ov_row_coord, (value_dim_base + 16 + ov_col_offset) % 64, elem_bytes=2), + u_pack_hi.data_ptr().load(count=4, alignment=4), nvvm.MMALayout.COL, shape=nvvm.StoreShape.M8N8, ) nvvm.fence_proxy("async.shared", space="cta") bars.mb_u_smem_ready.arrive() - # ---- dY readback ------------------------------------------------------- + # ---- dY read ------------------------------------------------------------- bars.mb_dy_acc_ready.wait(dy_acc_index.phase) dy_acc_index = advance(dy_acc_index, 1) dy_col_id = tmem_col + cfg.tmem_dy_acc_offset - dy_vec0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id0 << 16) + dy_col_id, cutlass.Float32), num=2) - dy_vec1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id1 << 16) + dy_col_id, cutlass.Float32), num=2) - - # ---- dY -> sdY: pack + store + publish (super warp's dM operand) ------- - addr_lo0 = ( - (value_dim_base + value_col_offset) // 64 * (cfg.b_t * 64) - + token_row_coord * 64 - + swizzle_xor_128b(token_row_coord, (value_dim_base + value_col_offset) % 64, elem_bytes=2) + dy_vec_lo = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_lo_addr + dy_col_id, cutlass.Float32), num=2) + dy_vec_hi = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_hi_addr + dy_col_id, cutlass.Float32), num=2) + + # ---- dY -> sdY: pack + store + publish (super warp's dM operand) --------- + dy_addr_lo = ( + (value_dim_base + ov_col_offset) // 64 * (cfg.b_t * 64) + + ov_row_coord * 64 + + swizzle_xor_128b(ov_row_coord, (value_dim_base + ov_col_offset) % 64, elem_bytes=2) ) - addr_lo1 = ( - (value_dim_base + 16 + value_col_offset) // 64 * (cfg.b_t * 64) - + token_row_coord * 64 - + swizzle_xor_128b(token_row_coord, (value_dim_base + 16 + value_col_offset) % 64, elem_bytes=2) + dy_addr_hi = ( + (value_dim_base + 16 + ov_col_offset) // 64 * (cfg.b_t * 64) + + ov_row_coord * 64 + + swizzle_xor_128b(ov_row_coord, (value_dim_base + 16 + ov_col_offset) % 64, elem_bytes=2) ) - dy_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) - dy_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + dy_pack_lo = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + dy_pack_hi = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) for reg_idx in cutlass.range_constexpr(4): - dy_pack0[reg_idx] = fp32_to_fp16(dy_vec0[2 * reg_idx], dy_vec0[2 * reg_idx + 1], dtype=cfg.io_dtype) - dy_pack1[reg_idx] = fp32_to_fp16(dy_vec1[2 * reg_idx], dy_vec1[2 * reg_idx + 1], dtype=cfg.io_dtype) - nvvm.stmatrix(sDy_raw.data_ptr() + addr_lo0, dy_pack0.data_ptr().load(count=4, alignment=4), nvvm.MMALayout.COL, shape=nvvm.StoreShape.M8N8) - nvvm.stmatrix(sDy_raw.data_ptr() + addr_lo1, dy_pack1.data_ptr().load(count=4, alignment=4), nvvm.MMALayout.COL, shape=nvvm.StoreShape.M8N8) + dy_pack_lo[reg_idx] = fp32_to_fp16(dy_vec_lo[2 * reg_idx], dy_vec_lo[2 * reg_idx + 1], dtype=cfg.io_dtype) + dy_pack_hi[reg_idx] = fp32_to_fp16(dy_vec_hi[2 * reg_idx], dy_vec_hi[2 * reg_idx + 1], dtype=cfg.io_dtype) + nvvm.stmatrix(sDy_raw.data_ptr() + dy_addr_lo, dy_pack_lo.data_ptr().load(count=4, alignment=4), nvvm.MMALayout.COL, shape=nvvm.StoreShape.M8N8) + nvvm.stmatrix(sDy_raw.data_ptr() + dy_addr_hi, dy_pack_hi.data_ptr().load(count=4, alignment=4), nvvm.MMALayout.COL, shape=nvvm.StoreShape.M8N8) nvvm.fence_proxy("async.shared", space="cta") bars.mb_dy_smem_ready.arrive() - # ---- beta scalars -> beta.dY -> sdV (epilogue TMA + dK_decay MMA operand) --- - beta_c0 = (sBeta_ptr + (lane % 4) * 2).load().to(cutlass.Float32) - beta_c1 = (sBeta_ptr + (lane % 4) * 2 + 1).load().to(cutlass.Float32) - beta_c8 = (sBeta_ptr + (lane % 4) * 2 + 8).load().to(cutlass.Float32) - beta_c9 = (sBeta_ptr + (lane % 4) * 2 + 9).load().to(cutlass.Float32) - # latched before the stage is released; the dbeta store runs after the refill + # ---- beta scalars -> beta.dY -> sdV (epilogue TMA + dK decay MMA + # operand) ----------------------------------------------------------------- + beta_c0 = (sBeta_ptr + (lane_idx % 4) * 2).load().to(cutlass.Float32) + beta_c1 = (sBeta_ptr + (lane_idx % 4) * 2 + 1).load().to(cutlass.Float32) + beta_c8 = (sBeta_ptr + (lane_idx % 4) * 2 + 8).load().to(cutlass.Float32) + beta_c9 = (sBeta_ptr + (lane_idx % 4) * 2 + 9).load().to(cutlass.Float32) beta_self = cutlass.Float32(0.0) if cutlass.const_expr(cfg.beta_sigmoid): if cg1_tidx < cfg.b_t: beta_self = (sBeta_ptr + cg1_tidx).load().to(cutlass.Float32) bars.mb_beta_done[chunk_serial % cfg.smem_beta_stages].arrive() - beta_dy_regs0 = cutlass.Array(cutlass.Float32, 8, alignment=16) - beta_dy_regs1 = cutlass.Array(cutlass.Float32, 8, alignment=16) + beta_dy_regs_lo = cutlass.Array(cutlass.Float32, 8, alignment=16) + beta_dy_regs_hi = cutlass.Array(cutlass.Float32, 8, alignment=16) for e2 in cutlass.range_constexpr(4): e = 2 * e2 b_lo = beta_c8 if cutlass.const_expr(e >= 4) else beta_c0 b_hi = beta_c9 if cutlass.const_expr(e >= 4) else beta_c1 - beta_dy_regs0[e], beta_dy_regs0[e + 1] = fmul2(dy_vec0[e], dy_vec0[e + 1], b_lo, b_hi) - beta_dy_regs1[e], beta_dy_regs1[e + 1] = fmul2(dy_vec1[e], dy_vec1[e + 1], b_lo, b_hi) + beta_dy_regs_lo[e], beta_dy_regs_lo[e + 1] = fmul2(dy_vec_lo[e], dy_vec_lo[e + 1], b_lo, b_hi) + beta_dy_regs_hi[e], beta_dy_regs_hi[e + 1] = fmul2(dy_vec_hi[e], dy_vec_hi[e + 1], b_lo, b_hi) - # ---- -beta.dY -> TMEM: A operand of the dH ds-term --------------------- - neg_beta_dy_regs0 = cutlass.Array(cutlass.Float32, 8, alignment=16) - neg_beta_dy_regs1 = cutlass.Array(cutlass.Float32, 8, alignment=16) + # ---- -beta.dY -> TMEM: A operand of the dH ds-term ----------------------- + neg_beta_dy_regs_lo = cutlass.Array(cutlass.Float32, 8, alignment=16) + neg_beta_dy_regs_hi = cutlass.Array(cutlass.Float32, 8, alignment=16) for e in cutlass.range_constexpr(8): - neg_beta_dy_regs0[e] = -beta_dy_regs0[e] - neg_beta_dy_regs1[e] = -beta_dy_regs1[e] - neg_beta_dy_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) - neg_beta_dy_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + neg_beta_dy_regs_lo[e] = -beta_dy_regs_lo[e] + neg_beta_dy_regs_hi[e] = -beta_dy_regs_hi[e] + neg_beta_dy_pack_lo = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + neg_beta_dy_pack_hi = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) for reg_idx in cutlass.range_constexpr(4): frag_pair = reg_idx * 2 - neg_beta_dy_pack0[reg_idx] = fp32_to_fp16(neg_beta_dy_regs0[frag_pair], neg_beta_dy_regs0[frag_pair + 1], dtype=cfg.io_dtype) - neg_beta_dy_pack1[reg_idx] = fp32_to_fp16(neg_beta_dy_regs1[frag_pair], neg_beta_dy_regs1[frag_pair + 1], dtype=cfg.io_dtype) - nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_addr_lo + (tmem_col + cfg.tmem_neg_beta_dy_inp_offset), cutlass.Int8), neg_beta_dy_pack0[0:4]) - nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_addr_hi + (tmem_col + cfg.tmem_neg_beta_dy_inp_offset), cutlass.Int8), neg_beta_dy_pack1[0:4]) + neg_beta_dy_pack_lo[reg_idx] = fp32_to_fp16(neg_beta_dy_regs_lo[frag_pair], neg_beta_dy_regs_lo[frag_pair + 1], dtype=cfg.io_dtype) + neg_beta_dy_pack_hi[reg_idx] = fp32_to_fp16(neg_beta_dy_regs_hi[frag_pair], neg_beta_dy_regs_hi[frag_pair + 1], dtype=cfg.io_dtype) + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_lo_addr + (tmem_col + cfg.tmem_neg_beta_dy_input_offset), cutlass.Int8), neg_beta_dy_pack_lo[0:4]) + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_hi_addr + (tmem_col + cfg.tmem_neg_beta_dy_input_offset), cutlass.Int8), neg_beta_dy_pack_hi[0:4]) nvvm.tcgen05_wait("store") - bars.mb_neg_beta_dy_inp_ready.arrive() + bars.mb_neg_beta_dy_input_ready.arrive() - # ---- dBeta v-term parts: dY.(V - state_k), diff_w from Y staging ------- - tok4 = cutlass.Array(cutlass.Float32, 4, alignment=16) + # ---- dBeta v-term part: dY.(V - k state), diff w from Y stage ------------ + token4 = cutlass.Array(cutlass.Float32, 4, alignment=16) for s in cutlass.range_constexpr(4): - tok4[s] = cutlass.Float32(0.0) + token4[s] = cutlass.Float32(0.0) for j in cutlass.range_constexpr(4): - d0_lo, d0_hi = f16x2_to_f32(diff_w0[j], dtype=cfg.io_dtype) - d1_lo, d1_hi = f16x2_to_f32(diff_w1[j], dtype=cfg.io_dtype) - s_lo = cutlass.const_expr(2 * (j // 2)) - s_hi = cutlass.const_expr(2 * (j // 2) + 1) - t_lo, t_hi = ffma2(dy_vec0[2 * j], dy_vec0[2 * j + 1], d0_lo, d0_hi, tok4[s_lo], tok4[s_hi]) - t_lo, t_hi = ffma2(dy_vec1[2 * j], dy_vec1[2 * j + 1], d1_lo, d1_hi, t_lo, t_hi) - tok4[s_lo] = t_lo - tok4[s_hi] = t_hi - - # ---- beta.dY -> sdV (epilogue TMA + dK_decay MMA operand) ------------ - beta_dy_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) - beta_dy_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + d0_lo, d0_hi = f16x2_to_f32(diff_w_lo[j], dtype=cfg.io_dtype) + d1_lo, d1_hi = f16x2_to_f32(diff_w_hi[j], dtype=cfg.io_dtype) + slot = cutlass.const_expr(2 * (j // 2)) + lo, hi = ffma2(dy_vec_lo[2 * j], dy_vec_lo[2 * j + 1], d0_lo, d0_hi, token4[slot], token4[slot + 1]) + lo, hi = ffma2(dy_vec_hi[2 * j], dy_vec_hi[2 * j + 1], d1_lo, d1_hi, lo, hi) + token4[slot] = lo + token4[slot + 1] = hi + + # ---- beta.dY -> sdV (epilogue TMA + dK decay MMA operand) ---------------- + beta_dy_pack_lo = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + beta_dy_pack_hi = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) for reg_idx in cutlass.range_constexpr(4): - beta_dy_pack0[reg_idx] = fp32_to_fp16(beta_dy_regs0[2 * reg_idx], beta_dy_regs0[2 * reg_idx + 1], dtype=cfg.io_dtype) - beta_dy_pack1[reg_idx] = fp32_to_fp16(beta_dy_regs1[2 * reg_idx], beta_dy_regs1[2 * reg_idx + 1], dtype=cfg.io_dtype) + beta_dy_pack_lo[reg_idx] = fp32_to_fp16(beta_dy_regs_lo[2 * reg_idx], beta_dy_regs_lo[2 * reg_idx + 1], dtype=cfg.io_dtype) + beta_dy_pack_hi[reg_idx] = fp32_to_fp16(beta_dy_regs_hi[2 * reg_idx], beta_dy_regs_hi[2 * reg_idx + 1], dtype=cfg.io_dtype) dv_stage = chunk_serial % cfg.smem_dv_stages sdv_stage_base = dv_stage * (cfg.b_t * cfg.d_v) bars.mb_dv_tmastg_done[dv_stage].wait(dv_done_index.phase) dv_done_index = advance(dv_done_index, cfg.smem_dv_stages) nvvm.stmatrix( - sDv_raw.data_ptr() + sdv_stage_base + addr_lo0, - beta_dy_pack0.data_ptr().load(count=4, alignment=4), + sDv_raw.data_ptr() + sdv_stage_base + dy_addr_lo, + beta_dy_pack_lo.data_ptr().load(count=4, alignment=4), nvvm.MMALayout.COL, shape=nvvm.StoreShape.M8N8, ) nvvm.stmatrix( - sDv_raw.data_ptr() + sdv_stage_base + addr_lo1, - beta_dy_pack1.data_ptr().load(count=4, alignment=4), + sDv_raw.data_ptr() + sdv_stage_base + dy_addr_hi, + beta_dy_pack_hi.data_ptr().load(count=4, alignment=4), nvvm.MMALayout.COL, shape=nvvm.StoreShape.M8N8, ) nvvm.fence_proxy("async.shared", space="cta") bars.mb_dv_tmastg_ready[dv_stage].arrive() - # ---- dBeta = sum_v dY.(V - state_k) + M-term ------------------------- + # ---- dBeta = sum over v of dY.(V - k state) + M-term --------------------- nvvm.barrier_cta_sync(cfg.cg1_sync_barrier_id, thread_count=cfg.cg1_threads) bars.mb_dbeta_m_ready.wait(dbeta_m_index.phase) dbeta_m_index = advance(dbeta_m_index, 1) for off in cutlass.range_constexpr(3): step = cutlass.const_expr(4 << off) for s in cutlass.range_constexpr(4): - tok4[s] = tok4[s] + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, tok4[s], step, 31, kind=nvvm.Shfl.BFLY)) - if lane < 4: + token4[s] = token4[s] + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, token4[s], step, 31, kind=nvvm.Shfl.BFLY)) + if lane_idx < 4: for s in cutlass.range_constexpr(4): - sRed_raw[(warp_idx % 4) * cfg.b_t + (lane % 4) * 2 + (s % 2) + 8 * (s // 2)] = tok4[s] + sRed_raw[(warp_idx % 4) * cfg.b_t + (lane_idx % 4) * 2 + (s % 2) + 8 * (s // 2)] = token4[s] nvvm.barrier_cta_sync(cfg.cg1_sync_barrier_id, thread_count=cfg.cg1_threads) if cg1_tidx < cfg.b_t: acc = cutlass.Float32(0.0) @@ -2468,40 +2398,40 @@ def compute1_warp_group( if cutlass.const_expr(cfg.beta_sigmoid): db_val = db_val * (beta_self - beta_self * beta_self) token_idx = chunk_idx * cutlass.Int32(cfg.b_t) + cg1_tidx - if token_idx < seqlen_b and chunk_idx < wend: + if token_idx < batch_seqlen and chunk_idx < write_end: mDbeta[batch_start + token_idx, head_idx] = db_val.to(mDbeta.element_type) nvvm.barrier_cta_sync(cfg.cg1_sync_barrier_id, thread_count=cfg.cg1_threads) - # ---- dH capture for the next ----------------------------------------- + # ---- dH capture for the next --------------------------------------------- bars.mb_dstate_acc_ready.wait(dstate_ready_index.phase) dstate_ready_index = advance(dstate_ready_index, 1) if rev_idx + cutlass.Int32(1) < num_compute_chunks: bars.mb_dstate_smem_done.wait(dstate_smem_done_index.phase) bars.mb_dstate_smem_cg2_done.wait(dstate_smem_done_index.phase) dstate_smem_done_index = advance(dstate_smem_done_index, 1) - row_addr = (tmem_row + tmem_subpartition * cfg.threads_per_warp) << 16 - for sub in cutlass.range_constexpr(cfg.d_k // 32): + row_lo_addr = tmem_row << 16 + for i in cutlass.range_constexpr(cfg.d_k // 32): dstate_vec = nvvm.tcgen05_ld( - "32x32b", nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_dstate_acc_offset + sub * 32), cutlass.Float32), num=32 + "32x32b", nvvm.make_tmem_ptr(row_lo_addr + (tmem_col + cfg.tmem_dstate_acc_offset + i * 32), cutlass.Float32), num=32 ) dstate_pack = cutlass.Array(cutlass.Int32, 16, alignment=16) for pc in cutlass.range_constexpr(16): dstate_pack[pc] = fp32_to_fp16(dstate_vec[2 * pc], dstate_vec[2 * pc + 1], dtype=cfg.io_dtype) nvvm.tcgen05_st( "32x32b", - nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_dstate_inp_offset + sub * 16), cutlass.Int8), + nvvm.make_tmem_ptr(row_lo_addr + (tmem_col + cfg.tmem_dstate_input_offset + i * 16), cutlass.Int8), dstate_pack[0:16], ) nvvm.tcgen05_wait("store") - bars.mb_dstate_inp_ready.arrive() + bars.mb_dstate_input_ready.arrive() - # ---- dh_inp -> sdH: re-read after the TMEM publish --------------- - for sub in cutlass.range_constexpr(cfg.d_k // 32): + # ---- dstate input -> sdH: re-read after the TMEM publish ------------- + for i in cutlass.range_constexpr(cfg.d_k // 32): dstate_words = nvvm.tcgen05_ld( - "32x32b", nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_dstate_inp_offset + sub * 16), cutlass.Float32), num=16 + "32x32b", nvvm.make_tmem_ptr(row_lo_addr + (tmem_col + cfg.tmem_dstate_input_offset + i * 16), cutlass.Float32), num=16 ) for half in cutlass.range_constexpr(4): - d_base = sub * 32 + half * 8 + d_base = i * 32 + half * 8 h_vec = cutlass.Vector.from_elements( (dstate_words[half * 4], dstate_words[half * 4 + 1], dstate_words[half * 4 + 2], dstate_words[half * 4 + 3]), cutlass.Float32, @@ -2512,18 +2442,18 @@ def compute1_warp_group( bars.mb_dstate_smem_ready.arrive() raw_index = advance(raw_index, cfg.smem_raw_stages) - # ---- tile end: dS0 drain / zero-length pass-through ---------------------- + # ---- tile end: dstate0 store / zero-length pass-through ---------------------- if cutlass.const_expr(mDstate0 is not None): if num_compute_chunks > 0: - if wstart == 0: - row_addr = (tmem_row + tmem_subpartition * cfg.threads_per_warp) << 16 + if write_start == 0: + row_lo_addr = tmem_row << 16 dstate0_dst = (mDstate0.iterator + mDstate0.layout((batch_idx, head_idx, value_dim, 0))).raw_ptr() - for sub in cutlass.range_constexpr(cfg.d_k // 32): + for i in cutlass.range_constexpr(cfg.d_k // 32): dstate0_vec = nvvm.tcgen05_ld( - "32x32b", nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_dstate_acc_offset + sub * 32), cutlass.Float32), num=32 + "32x32b", nvvm.make_tmem_ptr(row_lo_addr + (tmem_col + cfg.tmem_dstate_acc_offset + i * 32), cutlass.Float32), num=32 ) for g in cutlass.range_constexpr(8): - (dstate0_dst + sub * 32 + g * 4).store( + (dstate0_dst + i * 32 + g * 4).store( cutlass.Vector.from_elements(tuple(dstate0_vec[g * 4 + t] for t in range(4)), cutlass.Float32), alignment=16, ) @@ -2535,9 +2465,10 @@ def compute1_warp_group( mDstate0[batch_idx, head_idx, value_dim, kd] = mDstate_in[batch_idx, head_idx, value_dim, kd] else: mDstate0[batch_idx, head_idx, value_dim, kd] = cutlass.Float32(0.0) - bars.mb_dstate0_acc_stored.arrive() + if num_compute_chunks > 0: + bars.mb_dstate0_acc_stored.arrive() chunk_serial_base += num_compute_chunks - tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one) bars.mb_tmem_done[0].arrive() @@ -2550,8 +2481,8 @@ def compute2_warp_group( num_ctas, cu_seqlens, mWorkItems, - sSched, - lane, + sScheduler, + lane_idx, tmem_base_holder, warp_idx, sGate_raw, @@ -2572,13 +2503,14 @@ def compute2_warp_group( g_last terms, reverse-cumsums it in registers, stages dGate for the epilogue's TMA store, and stages dQ/dK for the epilogue's TMA stores.""" nvvm.setmaxregister(cfg.num_regs_compute_group_2, nvvm.SetMaxRegisterAction.INCREASE) + elect_one = nvvm.elect_sync() nvvm.barrier_cta_sync(cfg.tmem_lifecycle_barrier_id, thread_count=cfg.tmem_user_threads) tmem_base = tmem_base_holder.load() tmem_col = tmem_base & 0xFFFF tmem_row = tmem_base >> 16 tmem_subpartition = warp_idx % 4 - channel = tmem_subpartition * cfg.threads_per_warp + lane - row_addr = (tmem_row + tmem_subpartition * cfg.threads_per_warp) << 16 + channel = tmem_subpartition * cfg.threads_per_warp + lane_idx + row_lo_addr = tmem_row << 16 cg2_tidx = channel raw_index = PipelineState.start(phase=0) @@ -2588,14 +2520,16 @@ def compute2_warp_group( dk_restore_part_index = PipelineState.start(phase=0) dgate_last_dstate_smem_index = PipelineState.start(phase=0) chunk_serial_base = cutlass.Int32(0) - sched_state = PipelineState.start(phase=0) + scheduler_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) FIRST_STATE_CHUNK = 0 if cfg.use_initial_state else 1 while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) - num_compute_chunks = cend - wstart + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) + num_compute_chunks = compute_end - write_start for rev_idx in cutlass.range(num_compute_chunks, unroll=1): - chunk_idx = cend - cutlass.Int32(1) - rev_idx + chunk_idx = compute_end - cutlass.Int32(1) - rev_idx chunk_serial = chunk_serial_base + rev_idx chunk_start = chunk_idx * cfg.b_t raw_stage = chunk_serial % cfg.smem_raw_stages @@ -2604,13 +2538,13 @@ def compute2_warp_group( if cutlass.const_expr(cfg.use_dstate_in): has_dstate = cutlass.Boolean(True) sGate_ptr = sGate_raw.data_ptr() + raw_stage * (cfg.d_k * cfg.b_t) - writes = chunk_idx < wend + writes = chunk_idx < write_end # ---- gate landed: CG0 publishes the decay ring only after consuming - # the gate TMA, so this wait is CG2's visibility guard for sGate ------ + # the gate TMA, so this wait is CG2's visibility guard for sGate ----------- bars.mb_k_decay_inv_ready[decay_stage].wait((chunk_serial // cfg.smem_decay_stages) % 2) - # ---- per-channel gate factors ---------------------------------------- + # ---- per-channel gate factors -------------------------------------------- f32_seg = channel // 32 f32_dim = channel - f32_seg * 32 f16_seg = channel // 64 @@ -2622,16 +2556,16 @@ def compute2_warp_group( nvvm.fence_proxy("async.shared", space="cta") bars.mb_gate_done[raw_stage].arrive() - # ---- staged raw Q/K: TMEM ring cols for this chunk ------------------- + # ---- staged raw Q/K: TMEM ring cols for this chunk ----------------------- qk_raw_stage = chunk_serial % cfg.tmem_qk_raw_stages - qraw_col = tmem_col + cfg.tmem_qraw_inp_offset + qk_raw_stage * (cfg.b_t // 2) - kraw_col = tmem_col + cfg.tmem_kraw_inp_offset + qk_raw_stage * (cfg.b_t // 2) + qraw_col = tmem_col + cfg.tmem_qraw_input_offset + qk_raw_stage * (cfg.b_t // 2) + kraw_col = tmem_col + cfg.tmem_kraw_input_offset + qk_raw_stage * (cfg.b_t // 2) norm_base = qk_raw_stage * (2 * cfg.b_t) bars.mb_qk_raw_ready[qk_raw_stage].wait((chunk_serial // cfg.tmem_qk_raw_stages) % 2) - # ---- dGate_last hdot: sum_v sdH[v, c] * S0[c, v] --------------------- + # ---- dGate last hdot: sum over v of sdH[v, c] * S0[c, v] ----------------- dgate_last_val = cutlass.Float32(0.0) - bars.mb_state_inp_ready[chunk_serial % 2].wait((chunk_serial // 2) % 2) + bars.mb_state_input_ready[chunk_serial % 2].wait((chunk_serial // 2) % 2) if has_dstate: bars.mb_dstate_smem_ready.wait(dgate_last_dstate_smem_index.phase) dgate_last_dstate_smem_index = advance(dgate_last_dstate_smem_index, 1) @@ -2640,7 +2574,7 @@ def compute2_warp_group( state_vec = nvvm.tcgen05_ld( "32x32b", nvvm.make_tmem_ptr( - row_addr + (tmem_col + cfg.tmem_state_inp_offset + (chunk_serial % 2) * (cfg.d_v // 2) + pl * 32 + row_half * 16), + row_lo_addr + (tmem_col + cfg.tmem_state_input_offset + (chunk_serial % 2) * (cfg.d_v // 2) + pl * 32 + row_half * 16), cutlass.Float32, ), num=16, @@ -2662,9 +2596,9 @@ def compute2_warp_group( part_a, part_b = fadd2(pa0, pb0, pa1, pb1) dgate_last_val = dgate_last_val + (part_a + part_b) bars.mb_dstate_smem_cg2_done.arrive() - bars.mb_state_inp_cg2_done[chunk_serial % 2].arrive() + bars.mb_state_input_cg2_done[chunk_serial % 2].arrive() - # ---- part-drain accumulators ------------------------------------------- + # ---- part-store accumulators --------------------------------------------- dq_n = cutlass.Array(cutlass.Float32, cfg.b_t, alignment=16) dk_n = cutlass.Array(cutlass.Float32, cfg.b_t, alignment=16) dgate_regs = cutlass.Array(cutlass.Float32, cfg.b_t, alignment=16) @@ -2674,14 +2608,14 @@ def compute2_warp_group( for t in cutlass.range_constexpr(cfg.b_t): dk_n[t] = cutlass.Float32(0.0) - # ---- dK_restore part drain: (eGl/eG) scale + dGate_last K-dot ---------- + # ---- dK restore part store: (eGl/eG) scale + dGate last K-dot ------------ if has_dstate: bars.mb_dk_restore_part_acc_ready.wait(dk_restore_part_index.phase) dk_restore_part_index = advance(dk_restore_part_index, 1) dk_restore_part_vec = nvvm.tcgen05_ld( - "32x32b", nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_dk_restore_acc_offset), cutlass.Float32), num=cfg.b_t + "32x32b", nvvm.make_tmem_ptr(row_lo_addr + (tmem_col + cfg.tmem_dk_restore_acc_offset), cutlass.Float32), num=cfg.b_t ) - kr_words = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + kraw_col, cutlass.Float32), num=cfg.b_t // 2) + kr_words = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_lo_addr + kraw_col, cutlass.Float32), num=cfg.b_t // 2) for t in cutlass.range_constexpr(cfg.b_t): dk_hat = egl * cute.math.rcp(eg[t], approx=True, ftz=True) * dk_restore_part_vec[t] dk_n[t] = dk_hat @@ -2691,27 +2625,27 @@ def compute2_warp_group( k_v = k_v * sNorm_raw[norm_base + cfg.b_t + t] dgate_last_acc[t % 4] = dgate_last_acc[t % 4] + k_v * dk_hat - # ---- dQ acc drain: eG.scale --------------------------------------------- + # ---- dQ acc store: eG.scale ---------------------------------------------- bars.mb_dq_acc_ready.wait(dq_acc_index.phase) dq_acc_index = advance(dq_acc_index, 1) - dq_vec = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_dq_acc_offset), cutlass.Float32), num=cfg.b_t) + dq_vec = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_lo_addr + (tmem_col + cfg.tmem_dq_acc_offset), cutlass.Float32), num=cfg.b_t) for t2 in cutlass.range_constexpr(cfg.b_t // 2): t = 2 * t2 - es_lo, es_hi = fmul2(eg[t], eg[t + 1], scale, scale) - dq_n[t], dq_n[t + 1] = fmul2(es_lo, es_hi, dq_vec[t], dq_vec[t + 1]) + lo, hi = fmul2(eg[t], eg[t + 1], scale, scale) + dq_n[t], dq_n[t + 1] = fmul2(lo, hi, dq_vec[t], dq_vec[t + 1]) - # ---- dK_inv part drain: (dA - dM) term, 1/eG scale ---------------------- + # ---- dK inv part store: (dA - dM) term, 1/eG scale ----------------------- bars.mb_dk_inv_part_acc_ready.wait(dk_inv_part_index.phase) dk_inv_part_index = advance(dk_inv_part_index, 1) - dk_inv_part_vec = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_dk_inv_acc_offset), cutlass.Float32), num=cfg.b_t) + dk_inv_part_vec = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_lo_addr + (tmem_col + cfg.tmem_dk_inv_acc_offset), cutlass.Float32), num=cfg.b_t) for t in cutlass.range_constexpr(cfg.b_t): dk_n[t] = dk_n[t] + dk_inv_part_vec[t] * cute.math.rcp(eg[t], approx=True, ftz=True) - # ---- dK_decay part drain: -eG scale, seeds dGate ------------------------ + # ---- dK decay part store: -eG scale, seeds dGate ------------------------- bars.mb_dk_decay_part_acc_ready.wait(dk_decay_part_index.phase) dk_decay_part_index = advance(dk_decay_part_index, 1) dk_decay_part_vec = nvvm.tcgen05_ld( - "32x32b", nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_dk_decay_acc_offset), cutlass.Float32), num=cfg.b_t + "32x32b", nvvm.make_tmem_ptr(row_lo_addr + (tmem_col + cfg.tmem_dk_decay_acc_offset), cutlass.Float32), num=cfg.b_t ) for t in cutlass.range_constexpr(cfg.b_t): dgate_regs[t] = -eg[t] * dk_decay_part_vec[t] @@ -2720,9 +2654,9 @@ def compute2_warp_group( nvvm.tcgen05_wait("load") bars.mb_dqk_acc_done.arrive() - # ---- dGate finalize -------------------------------------------------- - qf_words = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + qraw_col, cutlass.Float32), num=cfg.b_t // 2) - kf_words = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + kraw_col, cutlass.Float32), num=cfg.b_t // 2) + # ---- dGate finalize ------------------------------------------------------ + qf_words = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_lo_addr + qraw_col, cutlass.Float32), num=cfg.b_t // 2) + kf_words = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_lo_addr + kraw_col, cutlass.Float32), num=cfg.b_t // 2) for t in cutlass.range_constexpr(cfg.b_t): q_pair = cutlass.Vector.from_elements((qf_words[t // 2],), cutlass.Float32).bitcast(cfg.io_dtype) k_pair = cutlass.Vector.from_elements((kf_words[t // 2],), cutlass.Float32).bitcast(cfg.io_dtype) @@ -2734,12 +2668,12 @@ def compute2_warp_group( dgate_regs[t] = q_v * dq_n[t] + k_v * (cutlass.Float32(2.0) * dgate_regs[t] - dk_n[t]) dgate_regs[cfg.b_t - 1] = dgate_regs[cfg.b_t - 1] + ((dgate_last_acc[0] + dgate_last_acc[1]) + (dgate_last_acc[2] + dgate_last_acc[3])) - # ---- L2-norm backward row projection --------------------------------- + # ---- L2-norm backward row projection ------------------------------------- if cutlass.const_expr(cfg.l2norm): for grad, qk_col, inv_off in ((dq_n, qraw_col, 0), (dk_n, kraw_col, cfg.b_t)): dots = cutlass.Array(cutlass.Float32, cfg.b_t, alignment=16) for half in cutlass.range_constexpr(2): - p_words = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + qk_col + half * (cfg.b_t // 4), cutlass.Float32), num=cfg.b_t // 4) + p_words = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_lo_addr + qk_col + half * (cfg.b_t // 4), cutlass.Float32), num=cfg.b_t // 4) for tt2 in cutlass.range_constexpr(cfg.b_t // 4): t = half * (cfg.b_t // 2) + 2 * tt2 p_pair = cutlass.Vector.from_elements((p_words[tt2],), cutlass.Float32).bitcast(cfg.io_dtype) @@ -2752,12 +2686,12 @@ def compute2_warp_group( sh_lo = cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, dots[t], step, 31, kind=nvvm.Shfl.BFLY)) sh_hi = cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, dots[t + 1], step, 31, kind=nvvm.Shfl.BFLY)) dots[t], dots[t + 1] = fadd2(dots[t], dots[t + 1], sh_lo, sh_hi) - if lane == 0: + if lane_idx == 0: for t in cutlass.range_constexpr(cfg.b_t): sRed1_raw[tmem_subpartition * cfg.b_t + t] = dots[t] nvvm.barrier_cta_sync(cfg.cg2_sync_barrier_id, thread_count=cfg.cg2_threads) for half in cutlass.range_constexpr(2): - a_words = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + qk_col + half * (cfg.b_t // 4), cutlass.Float32), num=cfg.b_t // 4) + a_words = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_lo_addr + qk_col + half * (cfg.b_t // 4), cutlass.Float32), num=cfg.b_t // 4) for tt2 in cutlass.range_constexpr(cfg.b_t // 4): t = half * (cfg.b_t // 2) + 2 * tt2 a_pair = cutlass.Vector.from_elements((a_words[tt2],), cutlass.Float32).bitcast(cfg.io_dtype) @@ -2774,7 +2708,7 @@ def compute2_warp_group( bars.mb_qk_raw_done[qk_raw_stage].arrive() - # ---- stage dQ/dK for the epilogue TMA stores ------------------------- + # ---- stage dQ/dK for the epilogue TMA stores ----------------------------- dq_stage = chunk_serial % cfg.smem_dq_stages dk_stage = chunk_serial % cfg.smem_dk_stages bars.mb_dq_tmastg_done[dq_stage].wait(((chunk_serial // cfg.smem_dq_stages) + 1) % 2) @@ -2789,19 +2723,19 @@ def compute2_warp_group( bars.mb_dq_tmastg_ready[dq_stage].arrive() bars.mb_dk_tmastg_ready[dk_stage].arrive() - # ---- dGate_last add -------------------------------------------------- + # ---- dGate last add ------------------------------------------------------ if has_dstate: if chunk_idx >= FIRST_STATE_CHUNK: dgate_regs[cfg.b_t - 1] = dgate_regs[cfg.b_t - 1] + egl * dgate_last_val - # ---- dGate reverse cumsum -------------------------------------------- + # ---- dGate reverse cumsum ------------------------------------------------ suffix = cutlass.Float32(0.0) for rt in cutlass.range_constexpr(cfg.b_t): t = cfg.b_t - 1 - rt suffix = suffix + dgate_regs[t] dgate_regs[t] = suffix - # ---- stage dGate for the epilogue TMA store -------------------------- + # ---- stage dGate for the epilogue TMA store ------------------------------ dgate_stage = chunk_serial % cfg.smem_dgate_stages bars.mb_dgate_tmastg_done[dgate_stage].wait(((chunk_serial // cfg.smem_dgate_stages) + 1) % 2) for t in cutlass.range_constexpr(cfg.b_t): @@ -2811,7 +2745,7 @@ def compute2_warp_group( bars.mb_dgate_tmastg_ready[dgate_stage].arrive() raw_index = advance(raw_index, cfg.smem_raw_stages) chunk_serial_base += num_compute_chunks - tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one) bars.mb_tmem_done[0].arrive() @@ -2834,7 +2768,7 @@ def build_descs_body( base_dv, base_dgate, base_checkpoint, - desc_ws: cute.Tensor, + desc_workspace: cute.Tensor, cu_seqlens: cute.Tensor, q: cute.Tensor, k: cute.Tensor, @@ -2863,16 +2797,16 @@ def build_descs_body( prologue kernel after its order pass; warps past the array count fall through the widx guards.""" arr_words = n_batch * cutlass.Int32(TENSOR_MAP_QWORDS) - sub0 = cute.make_tensor(desc_ws.iterator, cute.make_layout((arr_words,), stride=(1,))) - sub1 = cute.make_tensor(desc_ws.iterator + arr_words, cute.make_layout((arr_words,), stride=(1,))) - sub2 = cute.make_tensor(desc_ws.iterator + 2 * arr_words, cute.make_layout((arr_words,), stride=(1,))) - sub3 = cute.make_tensor(desc_ws.iterator + 3 * arr_words, cute.make_layout((arr_words,), stride=(1,))) - sub4 = cute.make_tensor(desc_ws.iterator + 4 * arr_words, cute.make_layout((arr_words,), stride=(1,))) - sub5 = cute.make_tensor(desc_ws.iterator + 5 * arr_words, cute.make_layout((arr_words,), stride=(1,))) - sub6 = cute.make_tensor(desc_ws.iterator + 6 * arr_words, cute.make_layout((arr_words,), stride=(1,))) - sub7 = cute.make_tensor(desc_ws.iterator + 7 * arr_words, cute.make_layout((arr_words,), stride=(1,))) - sub8 = cute.make_tensor(desc_ws.iterator + 8 * arr_words, cute.make_layout((arr_words,), stride=(1,))) - sub9 = cute.make_tensor(desc_ws.iterator + 9 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub0 = cute.make_tensor(desc_workspace.iterator, cute.make_layout((arr_words,), stride=(1,))) + sub1 = cute.make_tensor(desc_workspace.iterator + arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub2 = cute.make_tensor(desc_workspace.iterator + 2 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub3 = cute.make_tensor(desc_workspace.iterator + 3 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub4 = cute.make_tensor(desc_workspace.iterator + 4 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub5 = cute.make_tensor(desc_workspace.iterator + 5 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub6 = cute.make_tensor(desc_workspace.iterator + 6 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub7 = cute.make_tensor(desc_workspace.iterator + 7 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub8 = cute.make_tensor(desc_workspace.iterator + 8 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub9 = cute.make_tensor(desc_workspace.iterator + 9 * arr_words, cute.make_layout((arr_words,), stride=(1,))) if widx == 0: if nvvm.elect_sync(): @@ -2920,7 +2854,7 @@ def build_descs_body( def prologue_kernel( run_order: cutlass.Constexpr[bool], order_gen: cutlass.Constexpr[bool], - has_sched: cutlass.Constexpr[bool], + has_scheduler: cutlass.Constexpr[bool], b_t: cutlass.Constexpr[int], base_q: cutlass.GridConstant[cuda.tensor_map.TensorMap], base_k: cutlass.GridConstant[cuda.tensor_map.TensorMap], @@ -2932,7 +2866,7 @@ def prologue_kernel( base_dv: cutlass.GridConstant[cuda.tensor_map.TensorMap], base_dgate: cutlass.GridConstant[cuda.tensor_map.TensorMap], base_checkpoint: cutlass.GridConstant[cuda.tensor_map.TensorMap], - desc_ws: cute.Tensor, + desc_workspace: cute.Tensor, cu_seqlens: cute.Tensor, q: cute.Tensor, k: cute.Tensor, @@ -2947,7 +2881,7 @@ def prologue_kernel( mStaging: cute.Tensor | None, mCount: cute.Tensor, mWorkItems: cute.Tensor | None, - mSched: cute.Tensor | None, + mScheduler: cute.Tensor | None, n_batch: cutlass.Int32, q_row_stride: cutlass.Int32, k_row_stride: cutlass.Int32, @@ -2963,7 +2897,7 @@ def prologue_kernel( ) -> None: """Single-CTA prologue. Under ``run_order`` this kernel is the first work-item-table consumer, so it LPT-orders the table and zeroes both - consumers' sched rings via :func:`order_body`; it then builds the + consumers' scheduler rings via :func:`order_body`; it then builds the per-batch TMA-descriptor arrays via :func:`build_descs_body`, one warp per array (the extra warps only take part in the order phase).""" tidx, _, _ = cute.arch.thread_idx() @@ -2976,10 +2910,10 @@ def prologue_kernel( n_heads_out = cutlass.Int32(gate.shape[1]) order_body( order_gen, - has_sched, + has_scheduler, b_t, ORDER_THREADS, - ORDER_ELEMS, + ORDER_ELEMENTS, tidx, n_heads_out, n_heads_out * n_batch, @@ -2987,7 +2921,7 @@ def prologue_kernel( mStaging, mCount, mWorkItems, - mSched, + mScheduler, sKey, sIdx, sSpread, @@ -3004,7 +2938,7 @@ def prologue_kernel( base_dv, base_dgate, base_checkpoint, - desc_ws, + desc_workspace, cu_seqlens, q, k, @@ -3037,7 +2971,7 @@ def prologue( b_t: cutlass.Constexpr[int], run_order: cutlass.Constexpr[bool], order_gen: cutlass.Constexpr[bool], - has_sched: cutlass.Constexpr[bool], + has_scheduler: cutlass.Constexpr[bool], q: cute.Tensor, k: cute.Tensor, v: cute.Tensor, @@ -3052,7 +2986,7 @@ def prologue( work_item_staging: cute.Tensor | None, work_count: cute.Tensor, work_items: cute.Tensor | None, - sched_all: cute.Tensor | None, + scheduler_all: cute.Tensor | None, tensormap_workspace: cute.Tensor, stream: cuda_driver.CUstream, ): @@ -3104,7 +3038,7 @@ def prologue( prologue_kernel( run_order, order_gen, - has_sched, + has_scheduler, b_t, base_q, base_k, @@ -3131,7 +3065,7 @@ def prologue( work_item_staging, work_count, work_items, - sched_all, + scheduler_all, cutlass.Int32(batch_size), cutlass.Int32(q.stride[0]), cutlass.Int32(k.stride[0]), @@ -3161,14 +3095,14 @@ def host( d_final_state: cute.Tensor | None, work_items: cute.Tensor | None, work_count: cute.Tensor | None, - sched_ctr: cute.Tensor | None, + scheduler_counter: cute.Tensor | None, tensormap_workspace: cute.Tensor, scale: cutlass.Float32, stream, ) -> None: num_sequences = cu_seqlens.shape[0] - 1 - # ---- launch ------------------------------------------------------------------ + # ---- launch ---------------------------------------------------------------------- n_desc = num_sequences grid_shape = (cfg.max_active_clusters, 1, 1) kernel( @@ -3185,7 +3119,7 @@ def host( d_final_state, work_items, work_count, - sched_ctr, + scheduler_counter, scale, ).launch( grid=grid_shape, @@ -3210,7 +3144,7 @@ def kernel( mDstate_in: cute.Tensor | None, mWorkItems: cute.Tensor, mCount: cute.Tensor, - mSched: cute.Tensor | None, + mScheduler: cute.Tensor | None, scale: cutlass.Float32, ) -> None: """BT=16 KDA backward kernel (persistent, 16 warps).""" @@ -3218,7 +3152,7 @@ def kernel( bidx = cute.arch.block_idx()[0] num_ctas = cute.arch.grid_dim()[0] warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - lane = tidx % cfg.threads_per_warp + lane_idx = tidx % cfg.threads_per_warp total_tiles = mCount[0] beta_expected = cfg.io_dtype if cutlass.const_expr(cfg.beta_sigmoid) else cutlass.Float32 @@ -3242,7 +3176,7 @@ def kernel( SMEM = cutlass.AddressSpace.smem bars = make_kda_bwd_bars(cfg) tmem_base_holder = cutlass.Array(cutlass.Int32, 1, space=SMEM, alignment=4) - sSched = cutlass.Array(cutlass.Int32, cfg.sched_stages, space=SMEM, alignment=16) + sScheduler = cutlass.Array(cutlass.Int32, cfg.scheduler_stages, space=SMEM, alignment=16) bpe = cfg.io_dtype.width // 8 SWZ = 2 LEAD = 16 @@ -3275,7 +3209,7 @@ def kernel( sDk_raw = cutlass.Array(cfg.io_dtype, cfg.dk_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) sDgate_raw = cutlass.Array(cutlass.Float32, cfg.dgate_cosize, space=SMEM, alignment=1024) - sState_alt = SmemTile( + sState_trans = SmemTile( base=sState_raw.data_ptr().toint(), elems_per_stage=((cfg.state_cosize) // (cfg.smem_state_stages)) * bpe, stages=cfg.smem_state_stages, @@ -3283,7 +3217,7 @@ def kernel( stride_byte_offset=STRIDE, layout=SWZ, ) - sState_direct = SmemTile( + sState = SmemTile( base=sState_raw.data_ptr().toint(), elems_per_stage=((cfg.state_cosize) // (cfg.smem_state_stages)) * bpe, stages=cfg.smem_state_stages, @@ -3291,7 +3225,7 @@ def kernel( stride_byte_offset=STRIDE, layout=SWZ, ) - sK_decay_lead16 = SmemTile( + sK_decay = SmemTile( base=sK_decay_raw.data_ptr().toint(), elems_per_stage=((cfg.operand_cosize) // (cfg.smem_decay_stages)) * bpe, stages=cfg.smem_decay_stages, @@ -3299,7 +3233,7 @@ def kernel( stride_byte_offset=STRIDE, layout=SWZ, ) - sK_inv_lead16 = SmemTile( + sK_inv = SmemTile( base=sK_inv_raw.data_ptr().toint(), elems_per_stage=((cfg.operand_cosize) // (cfg.smem_decay_stages)) * bpe, stages=cfg.smem_decay_stages, @@ -3307,7 +3241,7 @@ def kernel( stride_byte_offset=STRIDE, layout=SWZ, ) - sK_restore_lead16 = SmemTile( + sK_restore = SmemTile( base=sK_restore_raw.data_ptr().toint(), elems_per_stage=((cfg.operand_cosize) // (cfg.smem_decay_stages)) * bpe, stages=cfg.smem_decay_stages, @@ -3315,7 +3249,7 @@ def kernel( stride_byte_offset=STRIDE, layout=SWZ, ) - sDo_lead16 = SmemTile( + sDo = SmemTile( base=sDo_raw.data_ptr().toint(), elems_per_stage=((cfg.raw_v_cosize) // (cfg.smem_raw_stages)) * bpe, stages=cfg.smem_raw_stages, @@ -3323,7 +3257,7 @@ def kernel( stride_byte_offset=STRIDE, layout=SWZ, ) - sDo_amaj = SmemTile( + sDo_trans = SmemTile( base=sDo_raw.data_ptr().toint(), elems_per_stage=((cfg.raw_v_cosize) // (cfg.smem_raw_stages)) * bpe, stages=cfg.smem_raw_stages, @@ -3331,7 +3265,7 @@ def kernel( stride_byte_offset=STRIDE, layout=SWZ, ) - sU_lead16 = SmemTile( + sU = SmemTile( base=sU_raw.data_ptr().toint(), elems_per_stage=((cfg.b_t * cfg.d_v) // (1)) * bpe, stages=1, @@ -3339,7 +3273,7 @@ def kernel( stride_byte_offset=STRIDE, layout=SWZ, ) - sDv_lead16 = SmemTile( + sDv = SmemTile( base=sDv_raw.data_ptr().toint(), elems_per_stage=((cfg.dv_cosize) // (cfg.smem_dv_stages)) * bpe, stages=cfg.smem_dv_stages, @@ -3347,7 +3281,7 @@ def kernel( stride_byte_offset=STRIDE, layout=SWZ, ) - sDstate_alt = SmemTile( + sDstate_trans = SmemTile( base=sDstate_raw.data_ptr().toint(), elems_per_stage=((cfg.d_k * cfg.d_v) // (1)) * bpe, stages=1, @@ -3364,7 +3298,7 @@ def kernel( stride_byte_offset=STRIDE, layout=SWZ, ) - sK_inv_amaj = SmemTile( + sK_inv_trans = SmemTile( base=sK_inv_raw.data_ptr().toint(), elems_per_stage=((cfg.operand_cosize) // (cfg.smem_decay_stages)) * bpe, stages=cfg.smem_decay_stages, @@ -3419,22 +3353,22 @@ def kernel( bars.mb_state_done[stage].init() bars.mb_state_cg0_done[stage].init() for stage in cutlass.range_constexpr(2): - bars.mb_state_inp_ready[stage].init() - bars.mb_state_inp_done[stage].init() - bars.mb_state_inp_cg2_done[stage].init() + bars.mb_state_input_ready[stage].init() + bars.mb_state_input_done[stage].init() + bars.mb_state_input_cg2_done[stage].init() elif warp_idx == cfg.tcgen05_mma_warp_id: if elect_one: bars.mb_state_k_acc_ready.init() - bars.mb_y_inp_ready.init() + bars.mb_y_input_ready.init() bars.mb_u_acc_ready.init() bars.mb_u_smem_ready.init() bars.mb_du_acc_ready.init() - bars.mb_du_inp_ready.init() + bars.mb_du_input_ready.init() bars.mb_dy_acc_ready.init() - bars.mb_neg_beta_dy_inp_ready.init() + bars.mb_neg_beta_dy_input_ready.init() bars.mb_dy_smem_ready.init() bars.mb_dstate_acc_ready.init() - bars.mb_dstate_inp_ready.init() + bars.mb_dstate_input_ready.init() bars.mb_dstate_smem_ready.init() bars.mb_dstate_smem_done.init() bars.mb_dstate_smem_cg2_done.init() @@ -3478,9 +3412,9 @@ def kernel( for stage in cutlass.range_constexpr(cfg.smem_dv_stages): bars.mb_dv_tmastg_ready[stage].init() bars.mb_dv_tmastg_done[stage].init() - for stage in cutlass.range_constexpr(cfg.sched_stages): - bars.mb_sched_ready[stage].init() - bars.mb_sched_done[stage].init() + for stage in cutlass.range_constexpr(cfg.scheduler_stages): + bars.mb_scheduler_ready[stage].init() + bars.mb_scheduler_done[stage].init() diag_zero = cfg.io_dtype(0.0) for diag_idx in cutlass.range(tidx, cfg.diag_cosize, cfg.threads_per_cta, unroll=1): sState_scale_diag_raw[diag_idx] = diag_zero @@ -3494,9 +3428,9 @@ def kernel( num_ctas, cu_seqlens, mWorkItems, - mSched, - sSched, - lane, + mScheduler, + sScheduler, + lane_idx, sQ_raw, sK_raw, sV_raw, @@ -3519,8 +3453,8 @@ def kernel( num_ctas, cu_seqlens, mWorkItems, - sSched, - lane, + sScheduler, + lane_idx, sK_decay_raw, sK_inv_raw, sU_raw, @@ -3538,21 +3472,21 @@ def kernel( num_ctas, cu_seqlens, mWorkItems, - sSched, + sScheduler, tmem_base_holder, - sState_alt, - sState_direct, - sK_decay_lead16, - sK_inv_lead16, - sK_inv_amaj, - sK_restore_lead16, - sDo_lead16, - sDo_amaj, + sState_trans, + sState, + sK_decay, + sK_inv, + sK_inv_trans, + sK_restore, + sDo, + sDo_trans, sQ_decay_trans, sK_decay_trans, - sU_lead16, - sDv_lead16, - sDstate_alt, + sU, + sDv, + sDstate_trans, sIntermediate, sState_scale_diag, bars, @@ -3565,8 +3499,8 @@ def kernel( num_ctas, cu_seqlens, mWorkItems, - sSched, - lane, + sScheduler, + lane_idx, sK_inv_raw, sQ_decay_raw, sDo_raw, @@ -3590,8 +3524,8 @@ def kernel( num_ctas, cu_seqlens, mWorkItems, - sSched, - lane, + sScheduler, + lane_idx, tmem_base_holder, warp_idx, scale, @@ -3619,8 +3553,8 @@ def kernel( num_ctas, cu_seqlens, mWorkItems, - sSched, - lane, + sScheduler, + lane_idx, tmem_base_holder, warp_idx, sGate_raw, @@ -3641,8 +3575,8 @@ def kernel( num_ctas, cu_seqlens, mWorkItems, - sSched, - lane, + sScheduler, + lane_idx, tmem_base_holder, warp_idx, mDbeta, @@ -3679,10 +3613,10 @@ class KdaBwdCfg: v_ratio: int n_heads_out: int max_active_clusters: int - dyn_sched: bool = False - sched_stages: int = 8 + dynamic_scheduling: bool = False + scheduler_stages: int = 8 - # ---- fixed constants stamped from CFG by build_cfg --------------------------- + # ---- fixed constants stamped from CFG at build time ------------------------------ compute_group_0_warp_ids: tuple = CFG.COMPUTE_GROUP_0_WARP_IDS compute_group_2_warp_ids: tuple = CFG.COMPUTE_GROUP_2_WARP_IDS compute_group_1_warp_ids: tuple = CFG.COMPUTE_GROUP_1_WARP_IDS @@ -3700,7 +3634,7 @@ class KdaBwdCfg: num_regs_compute_group_2: int = CFG.NUM_REGS_COMPUTE_GROUP_2 num_regs_other: int = CFG.NUM_REGS_OTHER - # ---- named barrier slots (ids 1-4; 0 is the CTA-wide sync) ------------------- + # ---- named barrier slots (ids 1-4; 0 is the CTA-wide sync) ----------------------- cg0_sync_barrier_id: int = 1 cg0_threads: int = 0 cg2_sync_barrier_id: int = 2 @@ -3710,7 +3644,7 @@ class KdaBwdCfg: cg1_sync_barrier_id: int = 4 cg1_threads: int = 0 - # ---- SMEM / TMEM stage counts + TMEM column offsets -------------------------- + # ---- SMEM / TMEM stage counts + TMEM column offsets ------------------------------ smem_raw_stages: int = CFG.SMEM_RAW_STAGES smem_state_stages: int = CFG.SMEM_S_STAGES smem_decay_stages: int = CFG.SMEM_DECAY_STAGES @@ -3722,7 +3656,7 @@ class KdaBwdCfg: smem_beta_stages: int = 4 intermediate_tiles: int = 5 tmem_dstate_acc_offset: int = 0 - tmem_dstate_inp_offset: int = 0 + tmem_dstate_input_offset: int = 0 tmem_state_k_acc_offset: int = 0 tmem_u_acc_offset: int = 0 tmem_du_acc_offset: int = 0 @@ -3732,15 +3666,15 @@ class KdaBwdCfg: tmem_dk_inv_acc_offset: int = 0 tmem_dk_restore_acc_offset: int = 0 tmem_qk_raw_stages: int = 4 - tmem_qraw_inp_offset: int = 0 - tmem_kraw_inp_offset: int = 0 - tmem_y_inp_offset: int = 0 - tmem_du_inp_offset: int = 0 - tmem_neg_beta_dy_inp_offset: int = 0 - tmem_state_inp_offset: int = 0 + tmem_qraw_input_offset: int = 0 + tmem_kraw_input_offset: int = 0 + tmem_y_input_offset: int = 0 + tmem_du_input_offset: int = 0 + tmem_neg_beta_dy_input_offset: int = 0 + tmem_state_input_offset: int = 0 buffer_align_bytes: int = CFG.BUFFER_ALIGN_BYTES - # ---- buffer cosizes / TMA bytes stamped by build_cfg ------------------------- + # ---- buffer cosizes / TMA bytes stamped at build time ---------------------------- raw_qk_cosize: int = 0 raw_v_cosize: int = 0 raw_gate_cosize: int = 0 @@ -3777,7 +3711,7 @@ def build_cfg( v_ratio: int, n_heads_out: int, max_active_clusters: int, - dyn_sched: bool = False, + dynamic_scheduling: bool = False, ) -> KdaBwdCfg: if io_dtype not in (cutlass.Float16, cutlass.BFloat16): raise ValueError(f"io_dtype={io_dtype} not supported; only Float16 and BFloat16 are supported") @@ -3795,7 +3729,7 @@ def build_cfg( v_ratio=v_ratio, n_heads_out=n_heads_out, max_active_clusters=max_active_clusters, - dyn_sched=dyn_sched, + dynamic_scheduling=dynamic_scheduling, ) cfg.threads_per_cta = 16 * cfg.threads_per_warp cfg.cg0_threads = len(cfg.compute_group_0_warp_ids) * cfg.threads_per_warp @@ -3806,9 +3740,9 @@ def build_cfg( ) * cfg.threads_per_warp cfg.tmem_dstate_acc_offset = 0 - cfg.tmem_dstate_inp_offset = cfg.d_k - cfg.tmem_state_inp_offset = cfg.tmem_dstate_inp_offset + cfg.d_k // 2 - cfg.tmem_state_k_acc_offset = cfg.tmem_state_inp_offset + cfg.d_v + cfg.tmem_dstate_input_offset = cfg.d_k + cfg.tmem_state_input_offset = cfg.tmem_dstate_input_offset + cfg.d_k // 2 + cfg.tmem_state_k_acc_offset = cfg.tmem_state_input_offset + cfg.d_v cfg.tmem_u_acc_offset = cfg.tmem_state_k_acc_offset + cfg.b_t cfg.tmem_du_acc_offset = cfg.tmem_u_acc_offset + cfg.b_t cfg.tmem_dy_acc_offset = cfg.tmem_state_k_acc_offset @@ -3816,12 +3750,12 @@ def build_cfg( cfg.tmem_dk_decay_acc_offset = cfg.tmem_dq_acc_offset + cfg.b_t cfg.tmem_dk_inv_acc_offset = cfg.tmem_dk_decay_acc_offset + cfg.b_t cfg.tmem_dk_restore_acc_offset = cfg.tmem_dk_inv_acc_offset + cfg.b_t - cfg.tmem_y_inp_offset = cfg.tmem_dk_restore_acc_offset + cfg.b_t - cfg.tmem_neg_beta_dy_inp_offset = cfg.tmem_y_inp_offset - cfg.tmem_du_inp_offset = cfg.tmem_y_inp_offset + cfg.b_t // 2 - cfg.tmem_qraw_inp_offset = cfg.tmem_du_inp_offset + cfg.b_t // 2 - cfg.tmem_kraw_inp_offset = cfg.tmem_qraw_inp_offset + cfg.tmem_qk_raw_stages * (cfg.b_t // 2) - assert cfg.tmem_kraw_inp_offset + cfg.tmem_qk_raw_stages * (cfg.b_t // 2) <= 512 + cfg.tmem_y_input_offset = cfg.tmem_dk_restore_acc_offset + cfg.b_t + cfg.tmem_neg_beta_dy_input_offset = cfg.tmem_y_input_offset + cfg.tmem_du_input_offset = cfg.tmem_y_input_offset + cfg.b_t // 2 + cfg.tmem_qraw_input_offset = cfg.tmem_du_input_offset + cfg.b_t // 2 + cfg.tmem_kraw_input_offset = cfg.tmem_qraw_input_offset + cfg.tmem_qk_raw_stages * (cfg.b_t // 2) + assert cfg.tmem_kraw_input_offset + cfg.tmem_qk_raw_stages * (cfg.b_t // 2) <= 512 cfg.raw_qk_cosize = cfg.smem_raw_stages * cfg.d_k * cfg.b_t cfg.raw_v_cosize = cfg.smem_raw_stages * cfg.d_v * cfg.b_t @@ -3847,7 +3781,7 @@ def build_cfg( TENSORMAP_STATIC_SLOTS = 0 -# ---- Torch adapter / host-side compilation --------------------------------------- +# ---- Torch adapter / host-side compilation ------------------------------------------- @lru_cache(maxsize=None) @@ -3864,7 +3798,7 @@ def get_compiled_cache( gate_lower_bound: float, beta_sigmoid: bool, use_initial_state: bool, - dyn_sched: bool, + dynamic_scheduling: bool, run_order: bool, order_gen: bool, ): @@ -3898,8 +3832,8 @@ def chunk_kda_bwd_sm100( use_beta_sigmoid: bool = False, work_items=None, work_count=None, - sched_ctr=None, - sched_all=None, + scheduler_counter=None, + scheduler_all=None, work_item_scratch=None, order_in_prologue: bool = False, tensormap_workspace, @@ -3945,9 +3879,9 @@ def chunk_kda_bwd_sm100( use_beta_sigmoid: ``beta`` holds logits; sigmoid in-kernel work_items/work_count: split-K table (``common/split_k.py``, REQUIRED; an uncut table row is the whole (b, h) sequence); each item - computes chunks ``[wstart, cend)`` backward and writes - gradients only for ``[wstart, wend)`` - sched_ctr: ``(2,)`` int32 zeroed scratch enabling the dynamic + computes chunks ``[write_start, compute_end)`` backward and writes + gradients only for ``[write_start, write_end)`` + scheduler_counter: ``(2,)`` int32 zeroed scratch enabling the dynamic (work-stealing) tile scheduler tensormap_workspace: ``tensormap_workspace_bytes(module, B)`` bytes, 128-byte aligned, for the per-(batch, head) TMA-descriptor @@ -3960,12 +3894,14 @@ def chunk_kda_bwd_sm100( use_dstate_in = d_final_state is not None use_dstate0 = d_initial_state is not None if work_items is None or work_count is None: - raise ValueError("work_items/work_count are required (the split-table stage builds them for every launch)") - dyn_sched = sched_ctr is not None + raise ValueError( + "work_items/work_count are required (built by the split-table stage, or by an order-generating prologue when work_item_scratch is None)" + ) + dynamic_scheduling = scheduler_counter is not None run_order = order_in_prologue order_gen = order_in_prologue and work_item_scratch is None - if run_order and sched_all is None: - raise ValueError("order in the prologue requires sched_all (the prologue zeroes both consumers' sched rings)") + if run_order and scheduler_all is None: + raise ValueError("order in the prologue requires scheduler_all (the prologue zeroes both consumers' scheduler rings)") if str(state_checkpoints.dtype).split(".")[-1] != str(q.dtype).split(".")[-1]: raise ValueError(f"state_checkpoints dtype must match the io dtype: got {state_checkpoints.dtype} with io {q.dtype}") for name, hh in (("HQ", HQ), ("HK", HK), ("HV", HV)): @@ -3994,7 +3930,7 @@ def chunk_kda_bwd_sm100( gate_lower_bound, use_beta_sigmoid, use_initial_state, - dyn_sched, + dynamic_scheduling, run_order, order_gen, ) @@ -4015,7 +3951,7 @@ def chunk_kda_bwd_sm100( v_ratio=HO // HV, n_heads_out=HO, max_active_clusters=multiprocessor_count(current_device()), - dyn_sched=dyn_sched, + dynamic_scheduling=dynamic_scheduling, ) dstate0_cute = None @@ -4028,10 +3964,10 @@ def chunk_kda_bwd_sm100( wi_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1), divisibility=1) wc_cute = from_dlpack(work_count, assumed_align=4).mark_layout_dynamic() sc_cute = None - if dyn_sched: - sc_cute = from_dlpack(sched_ctr, assumed_align=4).mark_layout_dynamic() + if dynamic_scheduling: + sc_cute = from_dlpack(scheduler_counter, assumed_align=4).mark_layout_dynamic() - tensormap_ws_cute = from_dlpack(tensormap_workspace, assumed_align=128).mark_layout_dynamic() + tensormap_workspace_cute = from_dlpack(tensormap_workspace, assumed_align=128).mark_layout_dynamic() a_log_cute = from_dlpack(a_log, assumed_align=4) if a_log is not None else None dt_bias_cute = from_dlpack(dt_bias, assumed_align=16) if dt_bias is not None else None @@ -4054,7 +3990,7 @@ def chunk_kda_bwd_sm100( wi_cute, wc_cute, sc_cute, - tensormap_ws_cute, + tensormap_workspace_cute, scale, cu_stream, options="--enable-tvm-ffi --opt-level 2", @@ -4073,7 +4009,7 @@ def chunk_kda_bwd_sm100( dgate_pl = from_dlpack(dgate, assumed_align=16).mark_layout_dynamic(leading_dim=2) state_checkpoints_pl = from_dlpack(state_checkpoints, assumed_align=16).mark_layout_dynamic(leading_dim=3) cu_pl = from_dlpack(cu_seqlens, assumed_align=8).mark_layout_dynamic() - ws_pl = from_dlpack(tensormap_workspace, assumed_align=128).mark_layout_dynamic() + workspace_pl = from_dlpack(tensormap_workspace, assumed_align=128).mark_layout_dynamic() staging_pl = None if run_order and not order_gen: staging_pl = from_dlpack(work_item_scratch, assumed_align=16) @@ -4081,9 +4017,9 @@ def chunk_kda_bwd_sm100( work_items_pl = from_dlpack(work_items, assumed_align=16) work_items_pl.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1), divisibility=1) work_count_pl = from_dlpack(work_count, assumed_align=4).mark_layout_dynamic() - sched_pl = None + scheduler_pl = None if run_order: - sched_pl = from_dlpack(sched_all, assumed_align=4).mark_layout_dynamic() + scheduler_pl = from_dlpack(scheduler_all, assumed_align=4).mark_layout_dynamic() cache["prologue"] = cute.compile( prologue, io_dtype, @@ -4105,8 +4041,8 @@ def chunk_kda_bwd_sm100( staging_pl, work_count_pl, work_items_pl, - sched_pl, - ws_pl, + scheduler_pl, + workspace_pl, cu_stream, options="--enable-tvm-ffi", ) @@ -4125,7 +4061,7 @@ def chunk_kda_bwd_sm100( work_item_scratch if run_order else None, work_count, work_items, - sched_all if run_order else None, + scheduler_all if run_order else None, tensormap_workspace, cu_stream, ) @@ -4141,7 +4077,7 @@ def chunk_kda_bwd_sm100( d_final_state, work_items, work_count, - sched_ctr, + scheduler_counter, tensormap_workspace, scale, cu_stream, @@ -4168,8 +4104,8 @@ def run_bwd( d_final_state, work_items, work_count, - sched_ctr, - sched_all, + scheduler_counter, + scheduler_all, work_item_scratch, tensormap_workspace, scale, @@ -4196,7 +4132,7 @@ def run_bwd( work_item_scratch, work_count, work_items, - sched_all, + scheduler_all, tensormap_workspace, cu_stream, ) @@ -4212,7 +4148,7 @@ def run_bwd( d_final_state, work_items, work_count, - sched_ctr, + scheduler_counter, tensormap_workspace, scale, cu_stream, diff --git a/python/cudnn/linear_attention/frost/kernel/kda_prefill_config.py b/python/cudnn/linear_attention/frost/kernel/kda_prefill_config.py index b77fece84..1cffe8d43 100644 --- a/python/cudnn/linear_attention/frost/kernel/kda_prefill_config.py +++ b/python/cudnn/linear_attention/frost/kernel/kda_prefill_config.py @@ -53,7 +53,7 @@ class Cfg: # --- SMEM / TMEM ring stage counts --- SMEM_RAW_STAGES: int = 8 - SMEM_SCHED_STAGES: int = 8 + SMEM_SCHEDULER_STAGES: int = 8 SMEM_O_STAGES: int = 2 SMEM_DECAY_STAGES: int = 2 SMEM_INTERMEDIATE_STAGES: int = 2 diff --git a/python/cudnn/linear_attention/frost/kernel/kda_prefill_f16.py b/python/cudnn/linear_attention/frost/kernel/kda_prefill_f16.py index 3eb6c3724..149bde260 100644 --- a/python/cudnn/linear_attention/frost/kernel/kda_prefill_f16.py +++ b/python/cudnn/linear_attention/frost/kernel/kda_prefill_f16.py @@ -15,78 +15,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Chunked Kimi Delta Attention (KDA) prefill kernel for Blackwell SM100/SM103 -(Cutlass DSL), BT=16 tiling with a per-key-channel decay. Framework-neutral -entry ``chunk_kda_sm100``. - -Persistent kernel: the grid is the SM count and every warp role -runs a tile-scheduler loop (``decode_work_item``); a tile is one (batch, -head) sequence, or one split-K work item computing chunks ``[cstart, wend)`` -and writing O / checkpoints only for the owned ``[wstart, wend)`` (see -``common/split_k.py``; warmup chunks rebuild the incoming state from -zero). All ring stage/phase bookkeeping runs on cumulative per-CTA chunk -counters so pipelines flow seamlessly across tiles. - -Pipeline (direct CUTLASS primitives, chunk_idx-size 16 KDA schedule): - - load Q/K/V/Gate/Beta - optional in-kernel L2-norm of Q/K (L2NORM specialization) - exp2(G), exp2(-G), stage final-token exp2(G) as exp2(G_last) - super-MMA: KK/A/Neumann inverse (T_inv) + apply Beta - tcgen05-MMA: state*K/state*Q/U solve/state update/O - store O, periodic state checkpoints, final state - -ABI: q `[T, HQ, DK]`, k `[T, HK, DK]`, v `[T, HV, DV]`, gate -`[T, HO, DK]` fp32 (natural-log decay unless SAFE_GATE, which applies the -safe-gate transform from raw gate + a_log/dt_bias), beta `[T, HO]` fp32 -post-sigmoid, cu_seqlens int32, states/checkpoints `[N, HO, DV, DK]` (VK, k -contiguous). -GQA/GVA head broadcast follows repeat_interleave: source head = -head_idx // (HO // H_x). State presence, L2NORM, SAFE_GATE, checkpoints, and the -head ratios are compile-time specializations. - -Warp assignments (16 warps = 512 threads): - warps 0-7 : compute group 0 - Gate prefix scan + decay/restore operands - warps 8-11 : compute group 1 - TMEM value side, O drain, state stores - warp 12 : super-MMA - register-MMA KK^T + Neumann inverse - warp 13 : tcgen05-MMA - the six state GEMMs + the TMEM lifecycle - warp 14 : TMA load - per-chunk input G->S loads - warp 15 : epilogue - register-MMA A + the O TMA store - -SMEM layout (~221 KB total): - Buffer Bytes Stages - Q / K / V raw 32768 8 <-- SW128 TMA ring (io dtype) - gate raw 65536 8 <-- fp32 prefix-scan source - beta 512 8 <-- fp32 per-token scalars - K_inv 8192 2 <-- token-major ldmatrix/tcgen05 B operand - K decay / Q decay 2x 8192 2 <-- tcgen05 SW128 K-box-major A/B operands - K restore 8192 2 <-- tcgen05 B operand for the state update - state-scale diag 12288 3 <-- per-K-atom decay diagonal blocks - intermediate (A / T_inv) 2048 2 <-- SW32 16x16 register-MMA tiles - O staging 8192 2 <-- W128 output drain - -TMEM layout (272 of 512 columns): - Buffer Cols Purpose - state 0-127 state[DK,DV] fp32 recurrent state - state inp 128-191 packed b16 A operand view of the state - q_state_acc 192-223 2-stage state*Q -> O accumulator - state_k_acc 224-239 state*K fp32 accumulator - u_acc 240-255 U fp32 accumulator - y_inp 256-263 packed b16 Y staging: Beta * (V - state*K) - u_inp 264-271 packed b16 U input (b16 U repack) - -GEMM schedule (tcgen05-MMA warp, in issue order per chunk): - state*K -> state_k_acc - state*Q -> q_state_acc (the O acc) - state decay (diag blocks) - U = Y(T) @ T_inv -> u_acc - final_state += U @ K_restore - O += A @ U -> q_state_acc - -Requires a cutlass DSL build providing `cutlass.experimental.*`; not -available in the pip nvidia-cutlass-dsl releases. -""" - from dataclasses import dataclass from functools import lru_cache from typing import Callable, NamedTuple, Optional, Type @@ -98,7 +26,7 @@ import cutlass.cute as cute from cutlass.cute.runtime import from_dlpack -from ..common.split_k import ORDER_CAPACITY, ORDER_ELEMS, ORDER_THREADS, decode_work_item, order_body +from ..common.split_k import ORDER_CAPACITY, ORDER_ELEMENTS, ORDER_THREADS, decode_work_item, order_body from ..common.host import get_dtype from cudnn.frost.buffers import data_ptr from cudnn.frost.device import current_device, multiprocessor_count @@ -113,9 +41,10 @@ ) from cudnn.frost.tile_dsl.handles import GmemTileTma, MmaDesc, SmemTile, tma_slice_runtime_desc from cudnn.frost.tile_dsl.mma import mma_step, mma_ts_step -from cudnn.frost.tile_dsl.swizzle import swizzle_lin_128b, swizzle_lin_S, swizzle_xor_128b +from cudnn.frost.tile_dsl.swizzle import swizzle_xor_128b, swizzle_xor_32b from cudnn.frost.tile_dsl.tma import tma_load_tile, tma_store_commit, tma_store_tile, tma_store_wait, tma_tensormap_acquire from cudnn.frost.tile_dsl.pointwise import ( + sigmoid, opaque_f32_zero, f16x2_to_f32, fadd2, @@ -128,21 +57,12 @@ ) LOG2_E: float = 1.4426950408889634 - - DEFAULT_GATE_LOWER_BOUND: float = -5.0 - - -# Host-side API defaults. - - L2_NORM_EPS: float = 1.0e-12 class KdaBars(NamedTuple): - """Every inter-warp handoff as an ``MBarrier`` over its ring. Consumers - track ``(idx, phase)`` inline; the producer tag selects - the arrive lowering (``TMA_LOAD``/``MMA_COMMIT``/``THREAD``).""" + """Every inter-warp handoff as an ``MBarrier`` over its ring.""" mb_q_ready: MBarrier mb_q_done: MBarrier @@ -161,9 +81,9 @@ class KdaBars(NamedTuple): mb_state_k_acc_ready: MBarrier mb_u_acc_ready: MBarrier - mb_state_inp_ready: MBarrier - mb_y_inp_ready: MBarrier - mb_u_inp_ready: MBarrier + mb_state_input_ready: MBarrier + mb_y_input_ready: MBarrier + mb_u_input_ready: MBarrier mb_t_inv_ready: MBarrier mb_t_inv_done: MBarrier @@ -186,13 +106,12 @@ class KdaBars(NamedTuple): mb_checkpoint_tmastg_ready: MBarrier mb_checkpoint_tmastg_done: MBarrier - mb_sched_ready: MBarrier - mb_sched_done: MBarrier + mb_scheduler_ready: MBarrier + mb_scheduler_done: MBarrier def make_kda_bars(cfg) -> KdaBars: - """Bars factory. MUST be called from inside ``kernel`` (allocates the - mbarrier rings in SMEM ahead of the data buffers).""" + """KdaBars factory.""" def alloc(n): return cutlass.Array(cutlass.Int64, n, space=cutlass.AddressSpace.smem, alignment=8) @@ -216,9 +135,9 @@ def alloc(n): mb_o_acc_done=MBarrier(alloc(cfg.tmem_q_state_acc_stages), stages=cfg.tmem_q_state_acc_stages, init_count=CG1_THREADS, producer=Producer.THREAD), mb_state_k_acc_ready=MBarrier(alloc(1), stages=1, init_count=1, producer=Producer.MMA_COMMIT), mb_u_acc_ready=MBarrier(alloc(1), stages=1, init_count=1, producer=Producer.MMA_COMMIT), - mb_state_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), - mb_y_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), - mb_u_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_state_input_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_y_input_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_u_input_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), mb_t_inv_ready=MBarrier(alloc(cfg.smem_intermediate_stages), stages=cfg.smem_intermediate_stages, init_count=WARP, producer=Producer.THREAD), mb_t_inv_done=MBarrier(alloc(cfg.smem_intermediate_stages), stages=cfg.smem_intermediate_stages, init_count=1, producer=Producer.MMA_COMMIT), mb_a_ready=MBarrier(alloc(cfg.smem_intermediate_stages), stages=cfg.smem_intermediate_stages, init_count=WARP, producer=Producer.THREAD), @@ -248,40 +167,40 @@ def alloc(n): alloc(cfg.smem_checkpoint_stages), stages=cfg.smem_checkpoint_stages, init_count=CG1_THREADS, producer=Producer.THREAD ), mb_checkpoint_tmastg_done=MBarrier(alloc(cfg.smem_checkpoint_stages), stages=cfg.smem_checkpoint_stages, init_count=WARP, producer=Producer.THREAD), - mb_sched_ready=MBarrier(alloc(cfg.sched_stages), stages=cfg.sched_stages, init_count=1, producer=Producer.THREAD), - mb_sched_done=MBarrier(alloc(cfg.sched_stages), stages=cfg.sched_stages, init_count=15, producer=Producer.THREAD), + mb_scheduler_ready=MBarrier(alloc(cfg.scheduler_stages), stages=cfg.scheduler_stages, init_count=1, producer=Producer.THREAD), + mb_scheduler_done=MBarrier(alloc(cfg.scheduler_stages), stages=cfg.scheduler_stages, init_count=15, producer=Producer.THREAD), ) -# ---- Dynamic tile scheduler ------------------------------------------------------ +# ---- Dynamic tile scheduler ---------------------------------------------------------- @cute.jit -def sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ctas): +def scheduler_publish_next(cfg, bars, sScheduler, mScheduler, scheduler_state, tile_idx, num_ctas, elect_one): """TMA-warp side: pull the next tile off the global ticket, publish it.""" - if cutlass.const_expr(cfg.dyn_sched): - bars.mb_sched_done[sched_state.idx].wait(sched_state.phase) - if nvvm.elect_sync(): - fetched = cutlass.Int32(nvvm.atomicrmw("add", mSched.iterator, cutlass.Int32(1), mem_order="relaxed", syncscope="gpu")) - sSched[sched_state.idx] = num_ctas + fetched + if cutlass.const_expr(cfg.dynamic_scheduling): + bars.mb_scheduler_done[scheduler_state.idx].wait(scheduler_state.phase) + if elect_one: + fetched = cutlass.Int32(nvvm.atomicrmw("add", mScheduler.iterator, cutlass.Int32(1), mem_order="relaxed", syncscope="gpu")) + sScheduler[scheduler_state.idx] = num_ctas + fetched nvvm.bar_warp_sync(cute.arch.FULL_MASK) - next_tile = sSched[sched_state.idx] - if nvvm.elect_sync(): - bars.mb_sched_ready[sched_state.idx].arrive() - return next_tile, advance(sched_state, cfg.sched_stages) - return tile_idx + num_ctas, sched_state + next_tile = sScheduler[scheduler_state.idx] + if elect_one: + bars.mb_scheduler_ready[scheduler_state.idx].arrive() + return next_tile, advance(scheduler_state, cfg.scheduler_stages) + return tile_idx + num_ctas, scheduler_state @cute.jit -def sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas): +def scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one): """Consumer side: read the TMA warp's published next tile.""" - if cutlass.const_expr(cfg.dyn_sched): - bars.mb_sched_ready[sched_state.idx].wait(sched_state.phase) - next_tile = sSched[sched_state.idx] - if nvvm.elect_sync(): - bars.mb_sched_done[sched_state.idx].arrive() - return next_tile, advance(sched_state, cfg.sched_stages) - return tile_idx + num_ctas, sched_state + if cutlass.const_expr(cfg.dynamic_scheduling): + bars.mb_scheduler_ready[scheduler_state.idx].wait(scheduler_state.phase) + next_tile = sScheduler[scheduler_state.idx] + if elect_one: + bars.mb_scheduler_done[scheduler_state.idx].arrive() + return next_tile, advance(scheduler_state, cfg.scheduler_stages) + return tile_idx + num_ctas, scheduler_state @cute.jit @@ -292,9 +211,9 @@ def tmaldg_warp( num_ctas, cu_seqlens, mWorkItems, - mSched, - sSched, - lane, + mScheduler, + sScheduler, + lane_idx, sQ_raw, sK_raw, sV_raw, @@ -307,8 +226,13 @@ def tmaldg_warp( ) -> None: """TMA-LDG warp role (warp 14): persistent scheduler loop issuing the per-chunk Q/K/V/Gate G->S loads.""" - elect_one = nvvm.elect_sync() nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) + + raw_index = PipelineState.start(phase=1) + raw_bar_index = PipelineState.start(phase=0) + scheduler_state = PipelineState.start(phase=1) + + elect_one = nvvm.elect_sync() sQ_tma = SmemTile( base=sQ_raw, elems_per_stage=(cfg.d_k * cfg.b_t), @@ -353,12 +277,11 @@ def tmaldg_warp( tma_granu_elems=32, tma_subtile_stride_elems=(cfg.b_t * 32), ) - raw_index = PipelineState.start(phase=1) - raw_bar_index = PipelineState.start(phase=0) - sched_state = PipelineState.start(phase=1) tile_idx = cutlass.Int32(bidx) while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) head_o = head_idx head_q = head_idx if cfg.q_ratio == 1 else head_idx // cutlass.Int32(cfg.q_ratio) head_k = head_idx if cfg.k_ratio == 1 else head_idx // cutlass.Int32(cfg.k_ratio) @@ -373,31 +296,31 @@ def tmaldg_warp( tma_tensormap_acquire(desc_k_slot) tma_tensormap_acquire(desc_v_slot) tma_tensormap_acquire(desc_gate_slot) - for chunk_idx in cutlass.range(cstart, wend, 1, unroll=1): + for chunk_idx in cutlass.range(compute_start, write_end, 1, unroll=1): chunk_start = chunk_idx * cfg.b_t - # ---- Q load ---------------------------------------------------------- + # ---- Q load -------------------------------------------------------------- bars.mb_q_done[raw_index.idx].wait(raw_index.phase) if elect_one: bars.mb_q_ready[raw_bar_index.idx].arrive(n_bytes=cfg.tma_q_bytes) q_slice = tma_slice_runtime_desc(desc_q_slot, cutlass.Int32(0), head_q, chunk_start) tma_load_tile(sQ_tma[raw_index.idx], q_slice, bars.mb_q_ready[raw_bar_index.idx].smem_ptr, acquire=False) - # ---- K load ---------------------------------------------------------- + # ---- K load -------------------------------------------------------------- bars.mb_k_done[raw_index.idx].wait(raw_index.phase) if elect_one: bars.mb_k_ready[raw_bar_index.idx].arrive(n_bytes=cfg.tma_k_bytes) k_slice = tma_slice_runtime_desc(desc_k_slot, cutlass.Int32(0), head_k, chunk_start) tma_load_tile(sK_tma[raw_index.idx], k_slice, bars.mb_k_ready[raw_bar_index.idx].smem_ptr, acquire=False) - # ---- Gate load ------------------------------------------------------- + # ---- Gate load: GMEM -> SMEM --------------------------------------------- bars.mb_gate_done[raw_index.idx].wait(raw_index.phase) if elect_one: bars.mb_gate_ready[raw_bar_index.idx].arrive(n_bytes=cfg.tma_gate_bytes) gate_slice = tma_slice_runtime_desc(desc_gate_slot, cutlass.Int32(0), head_o, chunk_start) tma_load_tile(sGate_tma[raw_index.idx], gate_slice, bars.mb_gate_ready[raw_bar_index.idx].smem_ptr, acquire=False) - # ---- V load ---------------------------------------------------------- + # ---- V load -------------------------------------------------------------- bars.mb_v_done[raw_index.idx].wait(raw_index.phase) if elect_one: bars.mb_v_ready[raw_bar_index.idx].arrive(n_bytes=cfg.tma_v_bytes) @@ -406,7 +329,7 @@ def tmaldg_warp( raw_index = advance(raw_index, cfg.smem_raw_stages) raw_bar_index = advance(raw_bar_index, cfg.smem_raw_bar_stages) - tile_idx, sched_state = sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_publish_next(cfg, bars, sScheduler, mScheduler, scheduler_state, tile_idx, num_ctas, elect_one) @cute.jit @@ -417,8 +340,8 @@ def super_mma_warp( num_ctas, cu_seqlens, mWorkItems, - sSched, - lane, + sScheduler, + lane_idx, sK_inv_raw, sIntermediate_raw, sBeta_raw, @@ -428,29 +351,31 @@ def super_mma_warp( """Super-MMA warp role (warp 12): persistent scheduler loop computing the Neumann-series T_inv via register MMA.""" nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) + elect_one = nvvm.elect_sync() raw_index = PipelineState.start(phase=0) t_inv_free = PipelineState.start(phase=1) k_decay_ready = PipelineState.start(phase=0) - # ---- ldmatrix/stmatrix lane decode ------------------------------------------- - rhs_row_coord = lane % 8 + (cutlass.Int32(8) if (lane // 16) else cutlass.Int32(0)) - rhs_col_offset = cutlass.Int32(8) if ((lane // 8) % 2) else cutlass.Int32(0) - lhs_row_coord = lane % 8 + (cutlass.Int32(8) if ((lane // 8) % 2) else cutlass.Int32(0)) - lhs_col_offset = cutlass.Int32(8) if ((lane // 8) // 2) else cutlass.Int32(0) - decay_key_mask = cutlass.Int32(8) - stsm_row_coord = lane & 7 - stsm_col_coord = cutlass.Int32(0) - if (lane // 8) & 1: - stsm_row_coord = stsm_row_coord + cutlass.Int32(8) - if lane // 8 >= 2: - stsm_col_coord = cutlass.Int32(8) - stsm_idx = swizzle_lin_S(stsm_row_coord * cfg.b_t + (stsm_col_coord ^ (cfg.b_t // 2)), bbits=1, mbase=3, sshift=3) + # ---- ldmatrix/stmatrix lane decode ----------------------------------------------- + k_inv_row_coord = lane_idx % 8 + (cutlass.Int32(8) if (lane_idx // 16) else cutlass.Int32(0)) + k_inv_col_offset = cutlass.Int32(8) if ((lane_idx // 8) % 2) else cutlass.Int32(0) + k_decay_row_coord = lane_idx % 8 + (cutlass.Int32(8) if ((lane_idx // 8) % 2) else cutlass.Int32(0)) + k_decay_col_offset = cutlass.Int32(8) if ((lane_idx // 8) // 2) else cutlass.Int32(0) + t_inv_row_coord = lane_idx & 7 + t_inv_col_coord = cutlass.Int32(0) + if (lane_idx // 8) & 1: + t_inv_row_coord = t_inv_row_coord + cutlass.Int32(8) + if lane_idx // 8 >= 2: + t_inv_col_coord = cutlass.Int32(8) + t_inv_idx = t_inv_row_coord * cfg.b_t + swizzle_xor_32b(t_inv_row_coord, t_inv_col_coord) cum_chunk_base = cutlass.Int32(0) - sched_state = PipelineState.start(phase=0) + scheduler_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) - num_chunks_tile = wend - cstart # processed chunks; ring bookkeeping runs on cum_chunk_base + local_chunk_idx + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) + num_chunks_tile = write_end - compute_start for local_chunk_idx in cutlass.range(num_chunks_tile, unroll=1): cum_chunk = cum_chunk_base + local_chunk_idx decay_stage = k_decay_ready.idx @@ -463,54 +388,52 @@ def super_mma_warp( bars.mb_k_decay_inv_cg0_ready[decay_stage].wait(k_decay_ready.phase) k_decay_ready = advance(k_decay_ready, cfg.smem_decay_stages) - # ---- KK = K_decay @ K_inv^T ------------------------------------------ + # ---- KK = K decay @ K inv^T ---------------------------------------------- kk_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): kk_acc[accum_idx] = cutlass.Float32(0.0) - for k_block in cutlass.range_constexpr((cfg.d_k // 16)): - # Load B operand - k_inv_col = k_block * 16 + rhs_col_offset + for i in cutlass.range_constexpr((cfg.d_k // 16)): + k_inv_col = i * 16 + k_inv_col_offset k_inv_segment = k_inv_col // 64 - rhs_frag = nvvm.ldmatrix( + k_inv_frag = nvvm.ldmatrix( sK_inv_ptr + k_inv_segment * (cfg.b_t * 64) - + rhs_row_coord * 64 - + swizzle_xor_128b(rhs_row_coord, k_inv_col - k_inv_segment * 64, elem_bytes=2), + + k_inv_row_coord * 64 + + swizzle_xor_128b(k_inv_row_coord, k_inv_col - k_inv_segment * 64, elem_bytes=2), 4, nvvm.MMALayout.ROW, ) - # Load A operand - storage_key = (k_block * 16 + lhs_col_offset) ^ decay_key_mask - storage_slice = storage_key // 64 - kk_lhs_frag = nvvm.ldmatrix( + k_decay_col = i * 16 + k_decay_col_offset + k_decay_segment = k_decay_col // 64 + k_decay_frag = nvvm.ldmatrix( sK_decay_ptr - + storage_slice * (cfg.b_t * 64) - + swizzle_xor_128b(lhs_row_coord, lhs_row_coord * 64 + storage_key - storage_slice * 64, elem_bytes=2), + + k_decay_segment * (cfg.b_t * 64) + + swizzle_xor_128b(k_decay_row_coord, k_decay_row_coord * 64 + k_decay_col - k_decay_segment * 64, elem_bytes=2), 4, nvvm.MMALayout.ROW, ) mma_step( kk_acc, - (kk_lhs_frag[0], kk_lhs_frag[1], kk_lhs_frag[2], kk_lhs_frag[3]), - (rhs_frag[0], rhs_frag[1], rhs_frag[2], rhs_frag[3]), + (k_decay_frag[0], k_decay_frag[1], k_decay_frag[2], k_decay_frag[3]), + (k_inv_frag[0], k_inv_frag[1], k_inv_frag[2], k_inv_frag[3]), k_step=0, M=16, N=16, ab_dtype=cfg.io_dtype, ) - # ---- L = Beta * tril(KK, -1) fragment -------------------------------- + # ---- L = Beta * tril(KK, -1) --------------------------------------------- bars.mb_beta_ready[raw_index.idx].wait(raw_index.phase) - row_lo = lane // 4 + row_lo = lane_idx // 4 row_hi = row_lo + cutlass.Int32(8) beta_lo = (sBeta_ptr + row_lo).load().to(cutlass.Float32) beta_hi = (sBeta_ptr + row_hi).load().to(cutlass.Float32) l_regs = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): row_coord = row_hi if cutlass.const_expr(accum_idx % 4 >= 2) else row_lo - col_coord = (accum_idx // 4) * 8 + 2 * (lane % 4) + col_coord = (accum_idx // 4) * 8 + 2 * (lane_idx % 4) if cutlass.const_expr(accum_idx % 2 == 1): col_coord = col_coord + cutlass.Int32(1) l_regs[accum_idx] = kk_acc[accum_idx] if row_coord > col_coord else cutlass.Float32(0.0) @@ -524,13 +447,13 @@ def super_mma_warp( l_a3 = fp32_to_fp16(l_regs[6], l_regs[7], dtype=cfg.io_dtype) l_values = cutlass.Vector.from_elements((l_a0, l_a1, l_a2, l_a3), cutlass.Int32).bitcast(cfg.io_dtype).to(cutlass.Float32) - # ---- T_inv = I - L, then three Neumann doubling rounds --------------- + # ---- T^-1 = I - L, then three Neumann doubling rounds -------------------- tinv_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): row_coord = row_lo if cutlass.const_expr(accum_idx % 4 >= 2): row_coord = row_hi - col_coord = (accum_idx // 4) * 8 + 2 * (lane % 4) + col_coord = (accum_idx // 4) * 8 + 2 * (lane_idx % 4) if cutlass.const_expr(accum_idx % 2 == 1): col_coord = col_coord + cutlass.Int32(1) eye = cutlass.Float32(1.0) if row_coord == col_coord else cutlass.Float32(0.0) @@ -538,8 +461,8 @@ def super_mma_warp( lpow_a0, lpow_a1, lpow_a2, lpow_a3 = l_a0, l_a1, l_a2, l_a3 mov_lpow0, mov_lpow1, mov_lpow2, mov_lpow3 = movmatrix_16b(l_a0), movmatrix_16b(l_a1), movmatrix_16b(l_a2), movmatrix_16b(l_a3) - for _round in cutlass.range_constexpr(3): - # ---- Lpow = Lpow @ Lpow ------------------------------------------ + for neumann_round in cutlass.range_constexpr(3): + # ---- Lpow = Lpow @ Lpow ---------------------------------------------- sq_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): sq_acc[accum_idx] = cutlass.Float32(0.0) @@ -557,7 +480,7 @@ def super_mma_warp( lpow_a2 = fp32_to_fp16(sq_acc[4], sq_acc[5], dtype=cfg.io_dtype) lpow_a3 = fp32_to_fp16(sq_acc[6], sq_acc[7], dtype=cfg.io_dtype) mov_lpow0, mov_lpow1, mov_lpow2, mov_lpow3 = movmatrix_16b(lpow_a0), movmatrix_16b(lpow_a1), movmatrix_16b(lpow_a2), movmatrix_16b(lpow_a3) - # ---- T_inv += T_inv @ Lpow --------------------------------------- + # ---- T^-1 += T^-1 @ Lpow --------------------------------------------- upd_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): upd_acc[accum_idx] = cutlass.Float32(0.0) @@ -586,7 +509,7 @@ def super_mma_warp( bars.mb_t_inv_done[intermediate_stage].wait(t_inv_free.phase) t_inv_free = advance(t_inv_free, cfg.smem_intermediate_stages) nvvm.stmatrix( - sIntermediate_ptr + (cfg.b_t * cfg.b_t) + stsm_idx, + sIntermediate_ptr + (cfg.b_t * cfg.b_t) + t_inv_idx, [ fp32_to_fp16(tinv_acc[0], tinv_acc[1], dtype=cfg.io_dtype), fp32_to_fp16(tinv_acc[2], tinv_acc[3], dtype=cfg.io_dtype), @@ -601,7 +524,7 @@ def super_mma_warp( bars.mb_decay_super_done[decay_stage].arrive() raw_index = advance(raw_index, cfg.smem_raw_bar_stages) cum_chunk_base += num_chunks_tile - tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one) @cute.jit @@ -612,11 +535,11 @@ def tcgen05_mma_warp( num_ctas, cu_seqlens, mWorkItems, - sSched, + sScheduler, tmem_base_slot, sIntermediate, sK_decay, - sK_restore, + sK_restore_trans, sQ_decay, sState_scale_diag, bars, @@ -628,23 +551,23 @@ def tcgen05_mma_warp( nvvm.tcgen05_alloc(tmem_base_slot, cutlass.Int32(512), group=nvvm.CTAGroup.CTA_1) nvvm.barrier_cta_sync(cfg.tmem_lifecycle_barrier_id, thread_count=cfg.tmem_user_threads) tmem_base = tmem_base_slot.load() - state_inp_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_state_inp_offset, cutlass.Int8) + state_input_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_state_input_offset, cutlass.Int8) state_dsts = tuple(nvvm.make_tmem_ptr(tmem_base + cfg.tmem_state_acc_offset + k * 16, cutlass.Float32) for k in range(cfg.d_k // 16)) state_k_acc_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_state_k_acc_offset, cutlass.Float32) u_acc_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_u_acc_offset, cutlass.Float32) - y_inp_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_y_inp_offset, cutlass.Int8) - u_inp_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_u_inp_offset, cutlass.Int8) + y_input_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_y_input_offset, cutlass.Int8) + u_input_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_u_input_offset, cutlass.Int8) state_dst_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_state_acc_offset, cutlass.Float32) - state_inp_index = PipelineState.start(phase=0) + state_input_index = PipelineState.start(phase=0) state_read_index = PipelineState.start(phase=0) - y_inp_index = PipelineState.start(phase=0) - u_inp_index = PipelineState.start(phase=0) + y_input_index = PipelineState.start(phase=0) + u_input_index = PipelineState.start(phase=0) qk_scale_index = PipelineState.start(phase=0) k_decay_ready = PipelineState.start(phase=0) intermediate_ready = PipelineState.start(phase=0) o_acc_free = PipelineState.start(phase=1) - # ---- chunk-invariant GEMM descriptors ---------------------------------------- + # ---- chunk-invariant GEMM descriptors -------------------------------------------- bpe = cfg.io_dtype.width // 8 idesc_acc = nvvm.Tcgen05InstrDesc.build( c_dtype=cutlass.Float32, @@ -670,7 +593,7 @@ def tcgen05_mma_warp( m_dim=cfg.d_v, b_major=1, ) - bmm_state_desc = MmaDesc( + bmm_state_k_decay_desc = MmaDesc( M=cfg.d_v, N=cfg.b_t, K=cfg.d_k, @@ -682,7 +605,8 @@ def tcgen05_mma_warp( idesc=idesc_acc, kind=nvvm.Tcgen05MMAKind.F16, ) - bmm_diag_desc = MmaDesc( + bmm_state_q_decay_desc = bmm_state_k_decay_desc + bmm_state_diag_desc = MmaDesc( M=cfg.d_v, N=16, K=16, @@ -694,7 +618,7 @@ def tcgen05_mma_warp( idesc=idesc_diag, kind=nvvm.Tcgen05MMAKind.F16, ) - bmm_intermediate_desc = MmaDesc( + bmm_u_a_desc = MmaDesc( M=cfg.d_v, N=cfg.b_t, K=cfg.b_t, @@ -706,7 +630,8 @@ def tcgen05_mma_warp( idesc=idesc_acc, kind=nvvm.Tcgen05MMAKind.F16, ) - bmm_final_state_desc = MmaDesc( + bmm_y_t_inv_desc = bmm_u_a_desc + bmm_u_k_restore_desc = MmaDesc( M=cfg.d_v, N=cfg.d_k, K=cfg.b_t, @@ -718,40 +643,47 @@ def tcgen05_mma_warp( idesc=idesc_final_state, kind=nvvm.Tcgen05MMAKind.F16, ) - STATE_A_SEG = bmm_state_desc.sps_B * bmm_state_desc.tmem_advance_A - STATE_B_SEG = bmm_state_desc.smem_subtile_B >> 4 + STATE_A_SEG = bmm_state_k_decay_desc.sps_B * bmm_state_k_decay_desc.tmem_advance_A + STATE_B_SEG = bmm_state_k_decay_desc.smem_subtile_B >> 4 cum_chunk_base = cutlass.Int32(0) - sched_state = PipelineState.start(phase=0) + scheduler_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) - num_chunks_tile = wend - cstart + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) + num_chunks_tile = write_end - compute_start + if cutlass.const_expr(cfg.use_initial_state): + seed_state = compute_start == 0 for local_chunk_idx in cutlass.range(num_chunks_tile, unroll=1): cum_chunk = cum_chunk_base + local_chunk_idx - have_state = cutlass.Boolean(True) if cutlass.const_expr(cfg.use_initial_state) else local_chunk_idx > 0 + if cutlass.const_expr(cfg.use_initial_state): + have_state = local_chunk_idx > 0 or seed_state + else: + have_state = local_chunk_idx > 0 q_state_acc_stage = o_acc_free.idx decay_stage = k_decay_ready.idx state_scale_diag_stage = qk_scale_index.idx intermediate_stage = intermediate_ready.idx sK_decay_stage = sK_decay[decay_stage] sQ_decay_stage = sQ_decay[decay_stage] - sK_restore_stage = sK_restore[decay_stage] + sK_restore_stage = sK_restore_trans[decay_stage] sState_scale_diag_stage = sState_scale_diag[state_scale_diag_stage] sIntermediate_stage = sIntermediate[intermediate_stage] - # ---- state_k = S(T) @ K_decay^T -------------------------------------- + # ---- k state = state(T) @ K decay^T -------------------------------------- bars.mb_k_decay_inv_cg0_ready[decay_stage].wait(k_decay_ready.phase) k_decay_ready = advance(k_decay_ready, cfg.smem_decay_stages) if have_state: - bars.mb_state_inp_ready.wait(state_inp_index.phase) - state_inp_index = advance(state_inp_index, 1) + bars.mb_state_input_ready.wait(state_input_index.phase) + state_input_index = advance(state_input_index, 1) desc_k_decay = sK_decay_stage.desc() - for s in cutlass.range_constexpr(bmm_state_desc.num_subtiles_B): - for k in cutlass.range_constexpr(bmm_state_desc.sps_B): + for s in cutlass.range_constexpr(bmm_state_k_decay_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_state_k_decay_desc.sps_B): mma_ts_step( - bmm_state_desc, - state_inp_ptr.subview(s * STATE_A_SEG), + bmm_state_k_decay_desc, + state_input_ptr.subview(s * STATE_A_SEG), desc_k_decay + s * STATE_B_SEG, state_k_acc_ptr, k, @@ -761,18 +693,18 @@ def tcgen05_mma_warp( if elect_one: bars.mb_state_k_acc_ready.arrive(cta_group=1) - # ---- q_state = state(T) @ Q_decay^T --------------------------------- + # ---- q state = state(T) @ Q decay^T -------------------------------------- bars.mb_qk_scale_ready[qk_scale_index.idx].wait(qk_scale_index.phase) bars.mb_o_acc_done[q_state_acc_stage].wait(o_acc_free.phase) o_acc_free = advance(o_acc_free, cfg.tmem_q_state_acc_stages) q_state_acc_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_q_state_acc_offset + q_state_acc_stage * cfg.b_t, cutlass.Float32) if have_state: desc_q_decay = sQ_decay_stage.desc() - for s in cutlass.range_constexpr(bmm_state_desc.num_subtiles_B): - for k in cutlass.range_constexpr(bmm_state_desc.sps_B): + for s in cutlass.range_constexpr(bmm_state_q_decay_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_state_q_decay_desc.sps_B): mma_ts_step( - bmm_state_desc, - state_inp_ptr.subview(s * STATE_A_SEG), + bmm_state_q_decay_desc, + state_input_ptr.subview(s * STATE_A_SEG), desc_q_decay + s * STATE_B_SEG, q_state_acc_ptr, k, @@ -782,19 +714,19 @@ def tcgen05_mma_warp( if elect_one: bars.mb_decay_tcgen05_done[decay_stage].arrive(cta_group=1) - # ---- S decay = S(T) @ diag(exp2(G_last)) --------- + # ---- state decay = state(T) @ diag(exp2(g last)) (per-k-atom blocks) ----- if cutlass.const_expr(cfg.enable_checkpoints): if have_state: bars.mb_state_acc_read_done.wait(state_read_index.phase) state_read_index = advance(state_read_index, 1) if have_state: desc_diag = sState_scale_diag_stage.desc() - for k_block in cutlass.range_constexpr(cfg.d_k // 16): + for i in cutlass.range_constexpr(cfg.d_k // 16): mma_ts_step( - bmm_diag_desc, - state_inp_ptr.subview(k_block * bmm_diag_desc.tmem_advance_A), - desc_diag.advance_start_address(k_block * 256 * 2), - state_dsts[k_block], + bmm_state_diag_desc, + state_input_ptr.subview(i * bmm_state_diag_desc.tmem_advance_A), + desc_diag.advance_start_address(i * 256 * 2), + state_dsts[i], 0, cutlass.Boolean(False), ) @@ -802,33 +734,33 @@ def tcgen05_mma_warp( if elect_one: bars.mb_state_scale_diag_done[state_scale_diag_stage].arrive(cta_group=1) - # ---- U = Y(T) @ T_inv ------------------------------------------------ + # ---- U = Y(T) @ T^-1 ----------------------------------------------------- bars.mb_t_inv_ready[intermediate_stage].wait(intermediate_ready.phase) - bars.mb_y_inp_ready.wait(y_inp_index.phase) - y_inp_index = advance(y_inp_index, 1) + bars.mb_y_input_ready.wait(y_input_index.phase) + y_input_index = advance(y_input_index, 1) d_int = sIntermediate_stage.shifted((cfg.b_t * cfg.b_t)).desc() - mma_ts_step(bmm_intermediate_desc, y_inp_ptr, d_int, u_acc_ptr, 0, cutlass.Boolean(False)) + mma_ts_step(bmm_y_t_inv_desc, y_input_ptr, d_int, u_acc_ptr, 0, cutlass.Boolean(False)) if elect_one: bars.mb_t_inv_done[intermediate_stage].arrive(cta_group=1) bars.mb_u_acc_ready.arrive(cta_group=1) - # ---- final_state += U(T) @ K_restore --------------------------------- - bars.mb_u_inp_ready.wait(u_inp_index.phase) - u_inp_index = advance(u_inp_index, 1) + # ---- final state += U(T) @ K restore ------------------------------------- + bars.mb_u_input_ready.wait(u_input_index.phase) + u_input_index = advance(u_input_index, 1) desc_k_restore = sK_restore_stage.desc() - mma_ts_step(bmm_final_state_desc, u_inp_ptr, desc_k_restore, state_dst_ptr, 0, have_state) + mma_ts_step(bmm_u_k_restore_desc, u_input_ptr, desc_k_restore, state_dst_ptr, 0, have_state) if elect_one: bars.mb_k_restore_done[decay_stage].arrive(cta_group=1) bars.mb_state_acc_done.arrive(cta_group=1) - # ---- O += U(T) @ A --------------------------------------------------- + # ---- O += U(T) @ A ------------------------------------------------------- bars.mb_a_ready[intermediate_stage].wait(intermediate_ready.phase) intermediate_ready = advance(intermediate_ready, cfg.smem_intermediate_stages) d_int = sIntermediate_stage.desc() mma_ts_step( - bmm_intermediate_desc, - u_inp_ptr, + bmm_u_a_desc, + u_input_ptr, d_int, nvvm.make_tmem_ptr(tmem_base + cfg.tmem_q_state_acc_offset + q_state_acc_stage * cfg.b_t, cutlass.Float32), 0, @@ -840,7 +772,7 @@ def tcgen05_mma_warp( qk_scale_index = advance(qk_scale_index, cfg.smem_state_scale_diag_stages) cum_chunk_base += num_chunks_tile - tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one) bars.mb_tmem_done[0].wait(0) nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_1) nvvm.tcgen05_dealloc( @@ -858,8 +790,8 @@ def epilogue_warp( num_ctas, cu_seqlens, mWorkItems, - sSched, - lane, + sScheduler, + lane_idx, mO, sK_inv_raw, sO_raw, @@ -902,39 +834,40 @@ def epilogue_warp( qk_scale_index = PipelineState.start(phase=0) a_free = PipelineState.start(phase=1) - # ---- ldmatrix/stmatrix lane decode ------------------------------------------- - rhs_row_coord = lane % 8 + (cutlass.Int32(8) if (lane // 16) else cutlass.Int32(0)) - rhs_col_offset = cutlass.Int32(8) if ((lane // 8) % 2) else cutlass.Int32(0) - lhs_row_coord = lane % 8 + (cutlass.Int32(8) if ((lane // 8) % 2) else cutlass.Int32(0)) - lhs_col_offset = cutlass.Int32(8) if ((lane // 8) // 2) else cutlass.Int32(0) - decay_key_mask = cutlass.Int32(8) - stsm_row_coord = lane & 7 - stsm_col_coord = cutlass.Int32(0) - if (lane // 8) & 1: - stsm_row_coord = stsm_row_coord + cutlass.Int32(8) - if lane // 8 >= 2: - stsm_col_coord = cutlass.Int32(8) - stsm_idx = swizzle_lin_S(stsm_row_coord * cfg.b_t + (stsm_col_coord ^ (cfg.b_t // 2)), bbits=1, mbase=3, sshift=3) + # ---- ldmatrix/stmatrix lane decode ----------------------------------------------- + k_inv_row_coord = lane_idx % 8 + (cutlass.Int32(8) if (lane_idx // 16) else cutlass.Int32(0)) + k_inv_col_offset = cutlass.Int32(8) if ((lane_idx // 8) % 2) else cutlass.Int32(0) + q_decay_row_coord = lane_idx % 8 + (cutlass.Int32(8) if ((lane_idx // 8) % 2) else cutlass.Int32(0)) + q_decay_col_offset = cutlass.Int32(8) if ((lane_idx // 8) // 2) else cutlass.Int32(0) + a_row_coord = lane_idx & 7 + a_col_coord = cutlass.Int32(0) + if (lane_idx // 8) & 1: + a_row_coord = a_row_coord + cutlass.Int32(8) + if lane_idx // 8 >= 2: + a_col_coord = cutlass.Int32(8) + a_idx = a_row_coord * cfg.b_t + swizzle_xor_32b(a_row_coord, a_col_coord) cum_chunk_base = cutlass.Int32(0) - sched_state = PipelineState.start(phase=0) + scheduler_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) head_o = head_idx o_slot = batch_idx * cutlass.Int32(TENSOR_MAP_QWORDS) desc_o_slot = (desc_o_base + o_slot).tospace(cutlass.AddressSpace.generic) if cutlass.const_expr(cfg.enable_checkpoints): desc_checkpoint_slot = (desc_checkpoint_base + o_slot).tospace(cutlass.AddressSpace.generic) checkpoint_chunks = checkpoint_every_n_tokens // cutlass.Int32(cfg.b_t) - checkpoint_quot = (cstart + cutlass.Int32(1)) // checkpoint_chunks - checkpoint_mod = (cstart + cutlass.Int32(1)) % checkpoint_chunks + checkpoint_quot = (compute_start + cutlass.Int32(1)) // checkpoint_chunks + checkpoint_mod = (compute_start + cutlass.Int32(1)) % checkpoint_chunks if elect_one: tma_tensormap_acquire(desc_checkpoint_slot) if elect_one: tma_tensormap_acquire(desc_o_slot) - num_chunks_tile = wend - cstart + num_chunks_tile = write_end - compute_start if cutlass.const_expr(cfg.enable_checkpoints): - if num_chunks_tile > 0 and wstart == 0: + if num_chunks_tile > 0 and write_start == 0: checkpoint_stage = checkpoint_ready_index.idx bars.mb_checkpoint_tmastg_ready[checkpoint_stage].wait(checkpoint_ready_index.phase) checkpoint_ready_index = advance(checkpoint_ready_index, cfg.smem_checkpoint_stages) @@ -944,7 +877,7 @@ def epilogue_warp( tma_store_wait(0) bars.mb_checkpoint_tmastg_done[checkpoint_stage].arrive() for local_chunk_idx in cutlass.range(num_chunks_tile, unroll=1): - chunk_idx = cstart + local_chunk_idx + chunk_idx = compute_start + local_chunk_idx cum_chunk = cum_chunk_base + local_chunk_idx decay_stage = cum_chunk % cfg.smem_decay_stages intermediate_stage = a_free.idx @@ -955,38 +888,36 @@ def epilogue_warp( bars.mb_qk_scale_ready[qk_scale_index.idx].wait(qk_scale_index.phase) - # ---- A = Q_decay @ K_inv^T ------------------------------------------ + # ---- A = Q decay @ K inv^T ----------------------------------------------- a_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): a_acc[accum_idx] = cutlass.Float32(0.0) - for k_block in cutlass.range_constexpr((cfg.d_k // 16)): - # Load B operand - k_inv_col = k_block * 16 + rhs_col_offset + for i in cutlass.range_constexpr((cfg.d_k // 16)): + k_inv_col = i * 16 + k_inv_col_offset k_inv_segment = k_inv_col // 64 - rhs_frag = nvvm.ldmatrix( + k_inv_frag = nvvm.ldmatrix( sK_inv_ptr + k_inv_segment * (cfg.b_t * 64) - + rhs_row_coord * 64 - + swizzle_xor_128b(rhs_row_coord, k_inv_col - k_inv_segment * 64, elem_bytes=2), + + k_inv_row_coord * 64 + + swizzle_xor_128b(k_inv_row_coord, k_inv_col - k_inv_segment * 64, elem_bytes=2), 4, nvvm.MMALayout.ROW, ) - # Load A operand - storage_key = (k_block * 16 + lhs_col_offset) ^ decay_key_mask - storage_slice = storage_key // 64 - a_lhs_frag = nvvm.ldmatrix( + q_decay_col = i * 16 + q_decay_col_offset + q_decay_segment = q_decay_col // 64 + q_decay_frag = nvvm.ldmatrix( sQ_decay_ptr - + storage_slice * (cfg.b_t * 64) - + swizzle_xor_128b(lhs_row_coord, lhs_row_coord * 64 + storage_key - storage_slice * 64, elem_bytes=2), + + q_decay_segment * (cfg.b_t * 64) + + swizzle_xor_128b(q_decay_row_coord, q_decay_row_coord * 64 + q_decay_col - q_decay_segment * 64, elem_bytes=2), 4, nvvm.MMALayout.ROW, ) mma_step( a_acc, - (a_lhs_frag[0], a_lhs_frag[1], a_lhs_frag[2], a_lhs_frag[3]), - (rhs_frag[0], rhs_frag[1], rhs_frag[2], rhs_frag[3]), + (q_decay_frag[0], q_decay_frag[1], q_decay_frag[2], q_decay_frag[3]), + (k_inv_frag[0], k_inv_frag[1], k_inv_frag[2], k_inv_frag[3]), k_step=0, M=16, N=16, @@ -994,10 +925,10 @@ def epilogue_warp( ) for accum_idx in cutlass.range_constexpr(8): - row_coord = lane // 4 + row_coord = lane_idx // 4 if cutlass.const_expr(accum_idx % 4 >= 2): row_coord = row_coord + cutlass.Int32(8) - col_coord = (accum_idx // 4) * 8 + 2 * (lane % 4) + col_coord = (accum_idx // 4) * 8 + 2 * (lane_idx % 4) if cutlass.const_expr(accum_idx % 2 == 1): col_coord = col_coord + cutlass.Int32(1) a_acc[accum_idx] = a_acc[accum_idx] if row_coord >= col_coord else cutlass.Float32(0.0) @@ -1005,7 +936,7 @@ def epilogue_warp( bars.mb_a_done[intermediate_stage].wait(a_free.phase) a_free = advance(a_free, cfg.smem_intermediate_stages) nvvm.stmatrix( - sIntermediate_ptr + stsm_idx, + sIntermediate_ptr + a_idx, [ fp32_to_fp16(a_acc[0], a_acc[1], dtype=cfg.io_dtype), fp32_to_fp16(a_acc[2], a_acc[3], dtype=cfg.io_dtype), @@ -1020,7 +951,7 @@ def epilogue_warp( bars.mb_decay_super_done[decay_stage].arrive() qk_scale_index = advance(qk_scale_index, cfg.qk_scale_ready_stages) - # ---- checkpoint + O drain: checkpoint stores first (CG1 stages the checkpoint before O) ------------- + # ---- checkpoint + O store: checkpoint first ------------------------------ if local_chunk_idx > 0: output_chunk = chunk_idx - cutlass.Int32(1) output_chunk_start = output_chunk * cfg.b_t @@ -1028,9 +959,9 @@ def epilogue_warp( did_checkpoint = cutlass.Int32(0) checkpoint_stage = cutlass.Int32(0) if cutlass.const_expr(cfg.enable_checkpoints): - # ---- checkpoint store ---------------------------------------- + # ---- checkpoint store -------------------------------------------- do_checkpoint = checkpoint_mod == 0 - do_checkpoint = do_checkpoint and chunk_idx >= wstart + do_checkpoint = do_checkpoint and chunk_idx >= write_start checkpoint_stage = checkpoint_ready_index.idx if do_checkpoint: bars.mb_checkpoint_tmastg_ready[checkpoint_ready_index.idx].wait(checkpoint_ready_index.phase) @@ -1047,7 +978,7 @@ def epilogue_warp( bars.mb_o_tmastg_ready[o_stage].wait(((cum_chunk - cutlass.Int32(1)) // cfg.smem_o_stages) % 2) o_slice = tma_slice_runtime_desc(desc_o_slot, cutlass.Int32(0), head_o, output_chunk_start) did_o = cutlass.Int32(0) - if output_chunk >= wstart: + if output_chunk >= write_start: tma_store_tile(sO_tma[o_stage], o_slice, acquire=False) tma_store_commit() did_o = cutlass.Int32(1) @@ -1069,9 +1000,9 @@ def epilogue_warp( tma_store_wait(0) bars.mb_o_tmastg_done[o_stage].arrive() - # ---- last computed chunk drain (always owned: it is wend - 1) ------------ + # ---- last computed chunk store (always owned: it is wend - 1) ---------------- if num_chunks_tile > 0: - output_chunk = wend - cutlass.Int32(1) + output_chunk = write_end - cutlass.Int32(1) last_cum_chunk = cum_chunk_base + num_chunks_tile - cutlass.Int32(1) output_chunk_start = output_chunk * cfg.b_t o_stage = last_cum_chunk % cfg.smem_o_stages @@ -1082,7 +1013,7 @@ def epilogue_warp( tma_store_wait(0) bars.mb_o_tmastg_done[o_stage].arrive() cum_chunk_base += num_chunks_tile - tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one) @cute.jit @@ -1090,9 +1021,7 @@ def gate_scale(cfg, raw_gate: cutlass.Float32) -> cutlass.Float32: """Map raw gate to the log2-domain decay increment used by KDA.""" if cutlass.const_expr(cfg.safe_gate): - half = cutlass.Float32(0.5) - sigmoid = cute.math.tanh(raw_gate * half, approx=True) * half + half - return cfg.gate_scale_log2 * sigmoid + return cfg.gate_scale_log2 * sigmoid(raw_gate) return raw_gate * cutlass.Float32(LOG2_E) @@ -1104,8 +1033,8 @@ def compute0_warp_group( num_ctas, cu_seqlens, mWorkItems, - sSched, - lane, + sScheduler, + lane_idx, warp_idx, mQ, mA_log, @@ -1125,31 +1054,36 @@ def compute0_warp_group( """CG0 warp role (warps 0-7, two ping-pong groups): persistent scheduler loop for the Gate prefix scan and decay/restore operand materialization.""" nvvm.setmaxregister(cfg.num_regs_compute_group_0, nvvm.SetMaxRegisterAction.INCREASE) + elect_one = nvvm.elect_sync() + + scheduler_state = PipelineState.start(phase=0) + cg0_warp = warp_idx - cfg.compute_group_0_warp_ids[0] - cg0_group_id = cg0_warp // cfg.cg0_warps_per_group cg0_local_warp = cg0_warp % cfg.cg0_warps_per_group - prefix_dim = cg0_local_warp * cfg.threads_per_warp + lane + + cg0_group_id = cg0_warp // cfg.cg0_warps_per_group + channel_dim = cg0_local_warp * cfg.threads_per_warp + lane_idx cg0_a_log_exp = cutlass.Float32(1.0) cg0_dt_bias_value = cutlass.Float32(0.0) cum_chunk_base = cutlass.Int32(0) - sched_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) opaque_one = opaque_f32_zero() + cutlass.Float32(1.0) while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) head_o = head_idx - num_chunks_tile = wend - cstart + num_chunks_tile = write_end - compute_start if cutlass.const_expr(cfg.safe_gate): if num_chunks_tile > 0: cg0_a_log_exp = cute.math.exp2(mA_log[head_o].to(cutlass.Float32) * LOG2_E, fastmath=True) - cg0_dt_bias_value = mDt_bias[head_o, prefix_dim].to(cutlass.Float32) - # tile entry: both ping-pong groups inherit each other's delivery proofs (parity-swap guard) + cg0_dt_bias_value = mDt_bias[head_o, channel_dim].to(cutlass.Float32) nvvm.barrier_cta_sync(cfg.cg0_tile_entry_barrier_id, thread_count=cfg.cg0_group_count * cfg.cg0_threads_per_group) group_cum_chunk_start = cum_chunk_base + cutlass.Int32(cg0_group_id) diag_ring_idx = group_cum_chunk_start % cutlass.Int32(cfg.smem_state_scale_diag_stages) diag_ring_phase = (group_cum_chunk_start // cutlass.Int32(cfg.smem_state_scale_diag_stages)) % cutlass.Int32(2) for local_chunk_idx in cutlass.range(cg0_group_id, num_chunks_tile, cfg.cg0_group_count, unroll=1): - chunk_idx = cstart + local_chunk_idx + chunk_idx = compute_start + local_chunk_idx cum_chunk = cum_chunk_base + local_chunk_idx chunk_start = chunk_idx * cfg.b_t decay_stage = cum_chunk % cfg.smem_decay_stages @@ -1166,63 +1100,53 @@ def compute0_warp_group( sK_restore_ptr = sK_restore_raw.data_ptr() + decay_stage * (cfg.d_k * cfg.b_t) sState_scale_diag_ptr = sState_scale_diag_raw.data_ptr() + state_scale_diag_stage * ((cfg.d_k // 16) * 256) - # ---- Beta scalars --------------------------------------------------- + # ---- Beta scalars -------------------------------------------------------- if cg0_local_warp == 0: bars.mb_beta_done[raw_bar_stage].wait(((cum_chunk // cfg.smem_raw_bar_stages) + 1) % 2) - if lane < cfg.b_t: - token_idx = chunk_idx * cfg.b_t + lane + if lane_idx < cfg.b_t: + token_idx = chunk_idx * cfg.b_t + lane_idx beta_value = cutlass.Float32(0.0) - if token_idx < seqlen_b: + if token_idx < batch_seqlen: beta_value = mBeta[batch_start + token_idx, head_o].to(cutlass.Float32) if cutlass.const_expr(cfg.beta_sigmoid): - half = cutlass.Float32(0.5) - beta_value = (cute.math.tanh(beta_value * half, approx=True) * half + half).to(mBeta.element_type).to(cutlass.Float32) - sBeta_raw[raw_bar_stage * cfg.b_t + lane] = beta_value + beta_value = sigmoid(beta_value).to(mBeta.element_type).to(cutlass.Float32) + sBeta_raw[raw_bar_stage * cfg.b_t + lane_idx] = beta_value bars.mb_beta_ready[raw_bar_stage].arrive() bars.mb_gate_ready[raw_bar_stage].wait((cum_chunk // cfg.smem_raw_bar_stages) % 2) row_group_start = cg0_local_warp * (cfg.b_t // cfg.cg0_warps_per_group) - lane_row_group = lane // 8 - lane_in_row_group = lane - lane_row_group * 8 + lane_row_group = lane_idx // 8 + lane_in_row_group = lane_idx - lane_row_group * 8 decay_row = row_group_start + lane_row_group - decay_key_mask = cutlass.Int32(8) - prefix_dim = cg0_local_warp * cfg.threads_per_warp + lane + channel_dim = cg0_local_warp * cfg.threads_per_warp + lane_idx - # ---- Gate prefix scan ----------------------------------------------- - f32_segment = prefix_dim // 32 + # ---- Gate prefix scan ---------------------------------------------------- + f32_segment = channel_dim // 32 prefix_seg_base = f32_segment * (cfg.b_t * 32) - prefix_col = prefix_dim - f32_segment * 32 + prefix_col = channel_dim - f32_segment * 32 gate_raw = cutlass.Array(cutlass.Float32, cfg.b_t, alignment=16) for row in cutlass.range_constexpr(cfg.b_t): prefix_idx = prefix_seg_base + swizzle_xor_128b(row, row * 32 + prefix_col, elem_bytes=4) gate_raw[row] = (sGate_ptr + prefix_idx).load() g_prefix_regs = cutlass.Array(cutlass.Float32, cfg.b_t, alignment=16) if cutlass.const_expr(cfg.safe_gate): - valid_rows = seqlen_b - chunk_idx * cutlass.Int32(cfg.b_t) - valid_mask = cutlass.vector.create_mask([cfg.b_t], [valid_rows]) - for row_pair in cutlass.range_constexpr(cfg.b_t // 2): - row0 = row_pair * 2 - row1 = row0 + 1 - gate0 = cg0_a_log_exp * (gate_raw[row0] + cg0_dt_bias_value) - gate1 = cg0_a_log_exp * (gate_raw[row1] + cg0_dt_bias_value) - gate0 = gate_scale( - cfg, - gate0, - ) - gate1 = gate_scale( - cfg, - gate1, - ) - gate_pair = cutlass.Vector.from_elements((gate0, gate1), cutlass.Float32) - gate_pair = cutlass.vector.where(valid_mask[row0 : row1 + 1], gate_pair, 0.0) - g_prefix_regs[row0] = gate_pair[0] - g_prefix_regs[row1] = gate_pair[1] + for row in cutlass.range_constexpr(cfg.b_t): + gate = gate_raw[row] + token_idx = chunk_idx * cutlass.Int32(cfg.b_t) + cutlass.Int32(row) + if token_idx < batch_seqlen: + gate = gate_scale( + cfg, + cg0_a_log_exp * (gate + cg0_dt_bias_value), + ) + else: + gate = cutlass.Float32(0.0) + g_prefix_regs[row] = gate else: for row in cutlass.range_constexpr(cfg.b_t): gate = gate_raw[row] token_idx = chunk_idx * cutlass.Int32(cfg.b_t) + cutlass.Int32(row) - if token_idx < seqlen_b: + if token_idx < batch_seqlen: gate = gate_scale( cfg, gate, @@ -1243,7 +1167,7 @@ def compute0_warp_group( g_prefix_regs[row1] = prefix1 prefix_acc = prefix1 - # ---- exp2(G): stage prefixes + final-token decay --------------------- + # ---- exp2(g): stage prefixes + final-token decay ------------------------- for row in cutlass.range_constexpr(cfg.b_t): g_prefix_regs[row] = cute.math.exp2(g_prefix_regs[row], fastmath=True) @@ -1252,13 +1176,11 @@ def compute0_warp_group( prefix_idx = prefix_seg_base + swizzle_xor_128b(row, row * 32 + prefix_col, elem_bytes=4) (sGate_ptr + prefix_idx).store(g_prefix_regs[row]) - # ---- state-scale diag: stage exp2(G_last) decay blocks --------------- + # ---- state-scale diag: stage exp2(g last) decay blocks ------------------- bars.mb_state_scale_diag_done[state_scale_diag_stage].wait(diag_ring_phase ^ cutlass.Int32(1)) - block = prefix_dim // cutlass.Int32(16) - coord = prefix_dim - block * cutlass.Int32(16) - storage_col = coord ^ cutlass.Int32((cfg.b_t // 2)) - linear_idx = block * cutlass.Int32(256) + coord * cutlass.Int32(16) + storage_col - diag_idx = swizzle_lin_S(linear_idx, bbits=1, mbase=3, sshift=3) + block = channel_dim // cutlass.Int32(16) + coord = channel_dim - block * cutlass.Int32(16) + diag_idx = block * cutlass.Int32(256) + coord * cutlass.Int32(16) + swizzle_xor_32b(channel_dim, coord) sState_scale_diag_ptr[diag_idx] = exp_g_last.to(cfg.io_dtype) nvvm.barrier_cta_sync(cfg.cg0_group_sync_barrier_base_id + cg0_group_id, thread_count=cfg.cg0_threads_per_group) @@ -1269,7 +1191,7 @@ def compute0_warp_group( raw_q_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) raw_k_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) - # ---- optional Q/K L2-norm ------------------------------------------- + # ---- optional Q/K L2-norm ------------------------------------------------ if cutlass.const_expr(cfg.l2norm): q_sq_even = opaque_f32_zero() k_sq_even = opaque_f32_zero() @@ -1311,7 +1233,7 @@ def compute0_warp_group( q_inv_norm = cute.math.rsqrt(cute.math.max(q_sum_sq, norm_floor_sq), fastmath=True) k_inv_norm = cute.math.rsqrt(cute.math.max(k_sum_sq, norm_floor_sq), fastmath=True) - # ---- decay/restore operands: exp2(+-G) ------------------------------ + # ---- decay/restore operands: exp2(+-g) applied per key channel ----------- exp_g_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) exp_g_last_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) for dim_half in cutlass.range_constexpr(2): @@ -1334,7 +1256,7 @@ def compute0_warp_group( dim_base = dim_half * (cfg.d_k // 2) + lane_in_row_group * 8 reg_base = dim_half * 8 - # ---- K decay + K_inv operands: K * exp2(+G) and K * exp2(-G) ----- + # ---- K decay + K inv operands: K * exp2(+g) and K * exp2(-g) --------- k_decay_pack = cutlass.Array(cutlass.Int32, 4, alignment=16) for pair_idx in cutlass.range_constexpr(4): dim0 = pair_idx * 2 @@ -1376,11 +1298,9 @@ def compute0_warp_group( f16_segment_dim = dim_base - f16_segment * 64 k_inv_swizzled_idx = f16_segment * (cfg.b_t * 64) + decay_row * 64 + swizzle_xor_128b(decay_row, f16_segment_dim, elem_bytes=2) (sK_inv_ptr + k_inv_swizzled_idx).store(k_inv_vec, alignment=16) - storage_key = dim_base ^ decay_key_mask - storage_slice = storage_key // 64 - decay_swizzled_idx = storage_slice * (cfg.b_t * 64) + swizzle_xor_128b( - decay_row, decay_row * 64 + storage_key - storage_slice * 64, elem_bytes=2 - ) + decay_col = dim_base + decay_segment = decay_col // 64 + decay_swizzled_idx = decay_segment * (cfg.b_t * 64) + swizzle_xor_128b(decay_row, decay_row * 64 + decay_col - decay_segment * 64, elem_bytes=2) (sK_decay_ptr + decay_swizzled_idx).store(k_decay_vec, alignment=16) nvvm.fence_proxy("async.shared", space="cta") bars.mb_k_decay_inv_cg0_ready[decay_stage].arrive() @@ -1388,7 +1308,7 @@ def compute0_warp_group( bars.mb_k_done[raw_stage].arrive() bars.mb_gate_done[raw_stage].arrive() - # ---- Q_decay operand: Q * q_inv_norm -------------------------------- + # ---- Q decay operand: Q * q inv norm ------------------------------------- for dim_half in cutlass.range_constexpr(2): dim_base = dim_half * (cfg.d_k // 2) + lane_in_row_group * 8 reg_base = dim_half * 8 @@ -1412,14 +1332,12 @@ def compute0_warp_group( ), cutlass.Int32, ).bitcast(cfg.io_dtype) - storage_key = dim_base ^ decay_key_mask - storage_slice = storage_key // 64 - decay_swizzled_idx = storage_slice * (cfg.b_t * 64) + swizzle_xor_128b( - decay_row, decay_row * 64 + storage_key - storage_slice * 64, elem_bytes=2 - ) + decay_col = dim_base + decay_segment = decay_col // 64 + decay_swizzled_idx = decay_segment * (cfg.b_t * 64) + swizzle_xor_128b(decay_row, decay_row * 64 + decay_col - decay_segment * 64, elem_bytes=2) (sQ_decay_ptr + decay_swizzled_idx).store(q_decay_vec, alignment=16) - # ---- K_restore operand: K_inv * exp_g_last -------------------------- + # ---- K restore operand: K inv * exp2(g last) ----------------------------- bars.mb_k_restore_done[decay_stage].wait(((cum_chunk // cfg.smem_decay_stages + 1) % 2)) for dim_half in cutlass.range_constexpr(2): dim_base = dim_half * (cfg.d_k // 2) + lane_in_row_group * 8 @@ -1430,10 +1348,9 @@ def compute0_warp_group( dim1 = dim0 + 1 exp_g_last_pair = fp32_to_fp16(exp_g_last_regs[reg_base + dim0], exp_g_last_regs[reg_base + dim1], dtype=cfg.io_dtype) k_restore_pack[pair_idx] = mul_f16x2(k_inv_pack[dim_half * 4 + pair_idx], exp_g_last_pair, cfg.io_dtype) - storage_row = decay_row ^ (cfg.b_t // 2) f16_segment = dim_base // 64 f16_segment_dim = dim_base - f16_segment * 64 - k_restore_idx = f16_segment * (cfg.b_t * 64) + storage_row * 64 + swizzle_xor_128b(storage_row, f16_segment_dim, elem_bytes=2) + k_restore_idx = f16_segment * (cfg.b_t * 64) + decay_row * 64 + swizzle_xor_128b(decay_row, f16_segment_dim, elem_bytes=2) k_restore_vec = cutlass.Vector.from_elements( ( k_restore_pack[0], @@ -1451,7 +1368,7 @@ def compute0_warp_group( diag_ring_idx = diag_ring_idx - cutlass.Int32(cfg.smem_state_scale_diag_stages) if wrapped else diag_ring_idx diag_ring_phase = diag_ring_phase ^ (cutlass.Int32(1) if wrapped else cutlass.Int32(0)) cum_chunk_base += num_chunks_tile - tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one) @cute.jit @@ -1462,8 +1379,8 @@ def compute1_warp_group( num_ctas, cu_seqlens, mWorkItems, - sSched, - lane, + sScheduler, + lane_idx, tmem_base_slot, warp_idx, mState_out, @@ -1480,144 +1397,126 @@ def compute1_warp_group( """CG1 warp role (warps 8-11): persistent scheduler loop for the value-side TMEM staging, output drain, and state stores.""" nvvm.setmaxregister(cfg.num_regs_compute_group_1, nvvm.SetMaxRegisterAction.INCREASE) + elect_one = nvvm.elect_sync() + + checkpoint_done_index = PipelineState.start(phase=1) + sO_ptr = sO_raw.data_ptr() sCheckpoint_ptr = sCheckpoint_raw.data_ptr() if cutlass.const_expr(cfg.enable_checkpoints) else sO_raw.data_ptr() - checkpoint_done_index = PipelineState.start(phase=1) nvvm.barrier_cta_sync(cfg.tmem_lifecycle_barrier_id, thread_count=cfg.tmem_user_threads) tmem_base = tmem_base_slot.load() tmem_col = tmem_base & 0xFFFF tmem_row = tmem_base >> 16 tmem_subpartition = warp_idx % (cfg.d_v // cfg.threads_per_warp) - ov_token_coord = (lane // 16) * 8 + (lane & 7) - ov_col_coord = ((lane // 8) & 1) * 8 - row_id = tmem_row + tmem_subpartition * cfg.threads_per_warp - value_dim = tmem_subpartition * cfg.threads_per_warp + lane + ov_row_coord = (lane_idx // 16) * 8 + (lane_idx & 7) + ov_col_offset = ((lane_idx // 8) & 1) * 8 + value_dim = tmem_subpartition * cfg.threads_per_warp + lane_idx value_dim_base = tmem_subpartition * cfg.threads_per_warp - row_addr = row_id << 16 - row16_addr = (row_id + 16) << 16 - st_row_addr = tmem_row << 16 - st_row16_addr = (tmem_row + 16) << 16 + row_lo_addr = tmem_row << 16 + row_hi_addr = (tmem_row + 16) << 16 state_col_id = tmem_col + cfg.tmem_state_acc_offset - packed_col_id = tmem_col + cfg.tmem_state_inp_offset + packed_col_id = tmem_col + cfg.tmem_state_input_offset state_k_col_id = tmem_col + cfg.tmem_state_k_acc_offset - y_inp_col_id = tmem_col + cfg.tmem_y_inp_offset - u_acc_addr = row_addr + tmem_col + cfg.tmem_u_acc_offset - u_inp_addr = st_row_addr + tmem_col + cfg.tmem_u_inp_offset + y_input_col_id = tmem_col + cfg.tmem_y_input_offset + u_acc_addr = row_lo_addr + tmem_col + cfg.tmem_u_acc_offset + u_input_addr = row_lo_addr + tmem_col + cfg.tmem_u_input_offset q_state_col_base = tmem_col + cfg.tmem_q_state_acc_offset - ov_swz_off0 = ( - (value_dim_base + ov_col_coord) // 64 * (cfg.b_t * 64) - + ov_token_coord * 64 - + swizzle_xor_128b(ov_token_coord, (value_dim_base + ov_col_coord) % 64, elem_bytes=2) + ov_swz_off_lo = ( + (value_dim_base + ov_col_offset) // 64 * (cfg.b_t * 64) + + ov_row_coord * 64 + + swizzle_xor_128b(ov_row_coord, (value_dim_base + ov_col_offset) % 64, elem_bytes=2) ) - ov_swz_off = ( - (value_dim_base + 16 + ov_col_coord) // 64 * (cfg.b_t * 64) - + ov_token_coord * 64 - + swizzle_xor_128b(ov_token_coord, (value_dim_base + 16 + ov_col_coord) % 64, elem_bytes=2) + ov_swz_off_hi = ( + (value_dim_base + 16 + ov_col_offset) // 64 * (cfg.b_t * 64) + + ov_row_coord * 64 + + swizzle_xor_128b(ov_row_coord, (value_dim_base + 16 + ov_col_offset) % 64, elem_bytes=2) ) - checkpoint_swz_off0 = (value_dim_base + ov_col_coord) // 64 * (cfg.d_k * 64) - checkpoint_swz_col0 = (value_dim_base + ov_col_coord) % 64 - checkpoint_swz_off = (value_dim_base + 16 + ov_col_coord) // 64 * (cfg.d_k * 64) - checkpoint_swz_col = (value_dim_base + 16 + ov_col_coord) % 64 + checkpoint_swz_off_lo = (value_dim_base + ov_col_offset) // 64 * (cfg.d_k * 64) + checkpoint_swz_col_lo = (value_dim_base + ov_col_offset) % 64 + checkpoint_swz_off_hi = (value_dim_base + 16 + ov_col_offset) // 64 * (cfg.d_k * 64) + checkpoint_swz_col_hi = (value_dim_base + 16 + ov_col_offset) % 64 state_k_acc_index = PipelineState.start(phase=0) u_acc_index = PipelineState.start(phase=0) o_acc_index = PipelineState.start(phase=0) state_upd_index = PipelineState.start(phase=0) raw_index = PipelineState.start(phase=0) - raw_bar_index = PipelineState.start(phase=0) # even-depth ready/beta-ring slot (decoupled from the data ring) + raw_bar_index = PipelineState.start(phase=0) # even-depth ready/beta-ring slot cum_chunk_base = cutlass.Int32(0) - sched_state = PipelineState.start(phase=0) + scheduler_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) head_o = head_idx - num_chunks_tile = wend - cstart + num_chunks_tile = write_end - compute_start if num_chunks_tile > 0: - # ---- first chunk: seed state TMEM from mState_init ---------- - seed_from_initial_state = cstart == 0 + # ---- first chunk: seed state TMEM from mState init ----------------------- + seed_from_initial_state = compute_start == 0 + sV_ptr = sV_raw.data_ptr() + raw_index.idx * (cfg.d_v * cfg.b_t) + sBeta_ptr = sBeta_raw.data_ptr() + raw_bar_index.idx * cfg.b_t + + # ---- state seed: initial state GMEM -> state TMEM ------------------------ if cutlass.const_expr(mState_init is not None): if seed_from_initial_state: seed_vw = 16 // (mState_init.element_type.width // 8) seed_src = (mState_init.iterator + mState_init.layout((batch_idx, head_o, value_dim, 0))).raw_ptr() - for key_block_start in cutlass.range_constexpr(0, cfg.d_k, 32): - state_block = cutlass.Array(cutlass.Float32, 32, alignment=16) - for g in cutlass.range_constexpr(32 // seed_vw): - seed_chunk = (seed_src + key_block_start + g * seed_vw).load(count=seed_vw, alignment=16) + state_vecs = [] + for i in cutlass.range_constexpr(cfg.d_k // 16): + state_block = [] + for g in cutlass.range_constexpr(16 // seed_vw): + seed_chunk = (seed_src + i * 16 + g * seed_vw).load(count=seed_vw, alignment=16) for t in cutlass.range_constexpr(seed_vw): - state_block[g * seed_vw + t] = seed_chunk[t].to(cutlass.Float32) - - nvvm.tcgen05_st( - "32x32b", - nvvm.make_tmem_ptr((row_id << 16) + (tmem_col + cfg.tmem_state_acc_offset + key_block_start), cutlass.Float32), - state_block[0:32], - ) - else: - for key_block_start in cutlass.range_constexpr(0, cfg.d_k, 32): - state_block = cutlass.Array(cutlass.Float32, 32, alignment=16) - for col in cutlass.range_constexpr(32): - state_block[col] = cutlass.Float32(0.0) + state_block.append(seed_chunk[t].to(cutlass.Float32)) + state_vecs.append(state_block) + for i in cutlass.range_constexpr(cfg.d_k // 16): + packed_state = cutlass.Array(cutlass.Int32, 8, alignment=16) + for packed_col in cutlass.range_constexpr(8): + packed_state[packed_col] = fp32_to_fp16(state_vecs[i][2 * packed_col], state_vecs[i][2 * packed_col + 1], dtype=cfg.io_dtype) nvvm.tcgen05_st( "32x32b", - nvvm.make_tmem_ptr((row_id << 16) + (tmem_col + cfg.tmem_state_acc_offset + key_block_start), cutlass.Float32), - state_block[0:32], + nvvm.make_tmem_ptr(row_lo_addr + packed_col_id + i * 8, cutlass.Int8), + packed_state[0:8], ) - if cutlass.const_expr(mState_init is not None): - nvvm.tcgen05_wait("store") - sV_ptr = sV_raw.data_ptr() + raw_index.idx * (cfg.d_v * cfg.b_t) - sBeta_ptr = sBeta_raw.data_ptr() + raw_bar_index.idx * cfg.b_t - - # ---- state repack: acc TMEM -> packed b16 TMEM ---------------------- - if cutlass.const_expr(mState_init is not None): - state_vecs = [] - for sub in cutlass.range_constexpr(cfg.d_k // 16): - state_vecs.append(nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + state_col_id + sub * 16, cutlass.Float32), num=16)) - - for sub in cutlass.range_constexpr(cfg.d_k // 16): - packed_state = cutlass.Array(cutlass.Int32, 8, alignment=16) - for packed_col in cutlass.range_constexpr(8): - source_pair = packed_col ^ 4 - packed_state[packed_col] = fp32_to_fp16(state_vecs[sub][2 * source_pair], state_vecs[sub][2 * source_pair + 1], dtype=cfg.io_dtype) - nvvm.tcgen05_st( - "32x32b", - nvvm.make_tmem_ptr((tmem_row << 16) + packed_col_id + sub * 8, cutlass.Int8), - packed_state[0:8], - ) - nvvm.tcgen05_wait("store") - bars.mb_state_inp_ready.arrive() - if cutlass.const_expr(cfg.enable_checkpoints): - if wstart == 0: - checkpoint_stage = checkpoint_done_index.idx - bars.mb_checkpoint_tmastg_done[checkpoint_stage].wait(checkpoint_done_index.phase) - checkpoint_done_index = advance(checkpoint_done_index, cfg.smem_checkpoint_stages) - checkpoint_stage_base = checkpoint_stage * (cfg.d_k * cfg.d_v) - for sub in cutlass.range_constexpr(cfg.d_k // 16): - for g in cutlass.range_constexpr(2): - packs = tuple( - fp32_to_fp16(state_vecs[sub][g * 8 + 2 * t], state_vecs[sub][g * 8 + 2 * t + 1], dtype=cfg.io_dtype) for t in range(4) - ) - dk = sub * 16 + g * 8 - checkpoint_addr = ( - checkpoint_stage_base + (dk // 64) * (cfg.d_v * 64) + value_dim * 64 + swizzle_xor_128b(value_dim, dk % 64, elem_bytes=2) - ) - (sCheckpoint_ptr + checkpoint_addr).store( - cutlass.Vector.from_elements(packs, cutlass.Int32).bitcast(cfg.io_dtype), alignment=16 - ) - nvvm.tcgen05_wait("load") - nvvm.fence_proxy("async.shared", space="cta") - bars.mb_checkpoint_tmastg_ready[checkpoint_stage].arrive() - bars.mb_state_acc_read_done.arrive() + nvvm.tcgen05_wait("store") + bars.mb_state_input_ready.arrive() + if cutlass.const_expr(cfg.enable_checkpoints): + if write_start == 0: + checkpoint_stage = checkpoint_done_index.idx + bars.mb_checkpoint_tmastg_done[checkpoint_stage].wait(checkpoint_done_index.phase) + checkpoint_done_index = advance(checkpoint_done_index, cfg.smem_checkpoint_stages) + checkpoint_stage_base = checkpoint_stage * (cfg.d_k * cfg.d_v) + for i in cutlass.range_constexpr(cfg.d_k // 16): + for g in cutlass.range_constexpr(2): + packs = tuple( + fp32_to_fp16(state_vecs[i][g * 8 + 2 * t], state_vecs[i][g * 8 + 2 * t + 1], dtype=cfg.io_dtype) for t in range(4) + ) + dk = i * 16 + g * 8 + checkpoint_addr = ( + checkpoint_stage_base + + (dk // 64) * (cfg.d_v * 64) + + value_dim * 64 + + swizzle_xor_128b(value_dim, dk % 64, elem_bytes=2) + ) + (sCheckpoint_ptr + checkpoint_addr).store( + cutlass.Vector.from_elements(packs, cutlass.Int32).bitcast(cfg.io_dtype), alignment=16 + ) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_checkpoint_tmastg_ready[checkpoint_stage].arrive() + bars.mb_state_acc_read_done.arrive() if cutlass.const_expr(cfg.enable_checkpoints and mState_init is None): - if wstart == 0: + if write_start == 0: checkpoint_stage = checkpoint_done_index.idx bars.mb_checkpoint_tmastg_done[checkpoint_stage].wait(checkpoint_done_index.phase) checkpoint_done_index = advance(checkpoint_done_index, cfg.smem_checkpoint_stages) checkpoint_stage_base = checkpoint_stage * (cfg.d_k * cfg.d_v) zero_packs = tuple(cutlass.Int32(0) for _ in range(4)) - for sub in cutlass.range_constexpr(cfg.d_k // 16): + for i in cutlass.range_constexpr(cfg.d_k // 16): for g in cutlass.range_constexpr(2): - dk = sub * 16 + g * 8 + dk = i * 16 + g * 8 checkpoint_addr = ( checkpoint_stage_base + (dk // 64) * (cfg.d_v * 64) + value_dim * 64 + swizzle_xor_128b(value_dim, dk % 64, elem_bytes=2) ) @@ -1627,81 +1526,62 @@ def compute1_warp_group( nvvm.fence_proxy("async.shared", space="cta") bars.mb_checkpoint_tmastg_ready[checkpoint_stage].arrive() - # ---- Y staging: Y = Beta * (V - state*K) ----------------------------- + # ---- Y stage: Y = Beta * (V - k state) ----------------------------------- bars.mb_v_ready[raw_bar_index.idx].wait(raw_bar_index.phase) - raw_v_frag0 = nvvm.ldmatrix( - sV_ptr + ov_swz_off0, + raw_v_frag_lo = nvvm.ldmatrix( + sV_ptr + ov_swz_off_lo, 4, nvvm.MMALayout.COL, ) - raw_v_frag1 = nvvm.ldmatrix( - sV_ptr + ov_swz_off, + raw_v_frag_hi = nvvm.ldmatrix( + sV_ptr + ov_swz_off_hi, 4, nvvm.MMALayout.COL, ) bars.mb_beta_ready[raw_bar_index.idx].wait(raw_bar_index.phase) - if cutlass.const_expr(mState_init is not None): - bars.mb_state_k_acc_ready.wait(state_k_acc_index.phase) - - state_k_vec0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_addr + state_k_col_id, cutlass.Float32), num=2) - state_k_vec1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row16_addr + state_k_col_id, cutlass.Float32), num=2) - beta_pack = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) for reg_idx in cutlass.range_constexpr(4): - token0 = (((reg_idx // 2) * 4 + (lane & 3)) ^ 4) * 2 + token0 = ((reg_idx // 2) * 4 + (lane_idx & 3)) * 2 beta0 = (sBeta_ptr + token0).load().to(cutlass.Float32) beta1 = (sBeta_ptr + token0 + 1).load().to(cutlass.Float32) beta_pack[reg_idx] = fp32_to_fp16(beta0, beta1, dtype=cfg.io_dtype) - y_inp_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) - for reg_idx in cutlass.range_constexpr(4): - raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1) - frag_pair = (reg_idx ^ 2) * 2 - if cutlass.const_expr(mState_init is not None): - state_k_val0, state_k_val1 = state_k_vec0[frag_pair], state_k_vec0[frag_pair + 1] - state_k_pair = fp32_to_fp16(state_k_val0, state_k_val1, dtype=cfg.io_dtype) - diff_pair = sub_f16x2( - raw_v_frag0[raw_matrix], - state_k_pair, - cfg.io_dtype, - ) - else: - diff_pair = raw_v_frag0[raw_matrix] - y_inp_pack0[reg_idx] = mul_f16x2( - beta_pack[reg_idx], - diff_pair, - cfg.io_dtype, - ) - y_inp_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) - for reg_idx in cutlass.range_constexpr(4): - raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1) - frag_pair = (reg_idx ^ 2) * 2 - if cutlass.const_expr(mState_init is not None): - state_k_val0, state_k_val1 = state_k_vec1[frag_pair], state_k_vec1[frag_pair + 1] - state_k_pair = fp32_to_fp16(state_k_val0, state_k_val1, dtype=cfg.io_dtype) - diff_pair = sub_f16x2( - raw_v_frag1[raw_matrix], - state_k_pair, - cfg.io_dtype, - ) + y_lo = [cutlass.Int32(0) for _ in range(4)] + y_hi = [cutlass.Int32(0) for _ in range(4)] + if cutlass.const_expr(mState_init is not None): + if seed_from_initial_state: + bars.mb_state_k_acc_ready.wait(state_k_acc_index.phase) + state_k_acc_index = advance(state_k_acc_index, 1) + state_k_vec_lo = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_lo_addr + state_k_col_id, cutlass.Float32), num=2) + state_k_vec_hi = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_hi_addr + state_k_col_id, cutlass.Float32), num=2) + for reg_idx in cutlass.range_constexpr(4): + frag_pair = reg_idx * 2 + state_k_lo = fp32_to_fp16(state_k_vec_lo[frag_pair], state_k_vec_lo[frag_pair + 1], dtype=cfg.io_dtype) + state_k_hi = fp32_to_fp16(state_k_vec_hi[frag_pair], state_k_vec_hi[frag_pair + 1], dtype=cfg.io_dtype) + y_lo[reg_idx] = mul_f16x2(beta_pack[reg_idx], sub_f16x2(raw_v_frag_lo[reg_idx], state_k_lo, cfg.io_dtype), cfg.io_dtype) + y_hi[reg_idx] = mul_f16x2(beta_pack[reg_idx], sub_f16x2(raw_v_frag_hi[reg_idx], state_k_hi, cfg.io_dtype), cfg.io_dtype) else: - diff_pair = raw_v_frag1[raw_matrix] - y_inp_pack1[reg_idx] = mul_f16x2( - beta_pack[reg_idx], - diff_pair, - cfg.io_dtype, - ) + for reg_idx in cutlass.range_constexpr(4): + y_lo[reg_idx] = mul_f16x2(beta_pack[reg_idx], raw_v_frag_lo[reg_idx], cfg.io_dtype) + y_hi[reg_idx] = mul_f16x2(beta_pack[reg_idx], raw_v_frag_hi[reg_idx], cfg.io_dtype) + else: + for reg_idx in cutlass.range_constexpr(4): + y_lo[reg_idx] = mul_f16x2(beta_pack[reg_idx], raw_v_frag_lo[reg_idx], cfg.io_dtype) + y_hi[reg_idx] = mul_f16x2(beta_pack[reg_idx], raw_v_frag_hi[reg_idx], cfg.io_dtype) - nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(st_row_addr + y_inp_col_id, cutlass.Int8), y_inp_pack0[0:4]) - nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(st_row16_addr + y_inp_col_id, cutlass.Int8), y_inp_pack1[0:4]) + y_input_pack_lo = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + y_input_pack_hi = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + for reg_idx in cutlass.range_constexpr(4): + y_input_pack_lo[reg_idx] = y_lo[reg_idx] + y_input_pack_hi[reg_idx] = y_hi[reg_idx] + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_lo_addr + y_input_col_id, cutlass.Int8), y_input_pack_lo[0:4]) + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_hi_addr + y_input_col_id, cutlass.Int8), y_input_pack_hi[0:4]) nvvm.tcgen05_wait("store") - if cutlass.const_expr(mState_init is not None): - state_k_acc_index = advance(state_k_acc_index, 1) bars.mb_v_done[raw_index.idx].arrive() bars.mb_beta_done[raw_bar_index.idx].arrive() - bars.mb_y_inp_ready.arrive() + bars.mb_y_input_ready.arrive() - # ---- U repack: u_acc TMEM -> packed b16 U input TMEM ---------------- + # ---- U stage: u acc TMEM -> packed b16 U input TMEM ---------------------- bars.mb_u_acc_ready.wait(u_acc_index.phase) u_vals = nvvm.tcgen05_ld( "32x32b", @@ -1709,21 +1589,20 @@ def compute1_warp_group( num=cfg.b_t, ) - u_inp_pack = cutlass.Array(cutlass.Int32, (cfg.b_t // 2), alignment=16) + u_input_pack = cutlass.Array(cutlass.Int32, (cfg.b_t // 2), alignment=16) for packed_col in cutlass.range_constexpr((cfg.b_t // 2)): - source_pair = packed_col ^ 4 - token0 = source_pair * 2 + token0 = packed_col * 2 token1 = token0 + 1 - u_inp_pack[packed_col] = fp32_to_fp16(u_vals[token0], u_vals[token1], dtype=cfg.io_dtype) + u_input_pack[packed_col] = fp32_to_fp16(u_vals[token0], u_vals[token1], dtype=cfg.io_dtype) nvvm.tcgen05_st( "32x32b", - nvvm.make_tmem_ptr(u_inp_addr, cutlass.Int8), - u_inp_pack[0 : (cfg.b_t // 2)], + nvvm.make_tmem_ptr(u_input_addr, cutlass.Int8), + u_input_pack[0 : (cfg.b_t // 2)], ) nvvm.tcgen05_wait("store") u_acc_index = advance(u_acc_index, 1) - bars.mb_u_inp_ready.arrive() + bars.mb_u_input_ready.arrive() if cutlass.const_expr(cfg.enable_checkpoints): bars.mb_state_acc_done.wait(state_upd_index.phase) state_upd_index = advance(state_upd_index, 1) @@ -1733,9 +1612,9 @@ def compute1_warp_group( if cutlass.const_expr(cfg.enable_checkpoints): cg1_checkpoint_chunks = checkpoint_every_n_tokens // cutlass.Int32(cfg.b_t) - cg1_checkpoint_mod = (cstart + cutlass.Int32(1)) % cg1_checkpoint_chunks + cg1_checkpoint_mod = (compute_start + cutlass.Int32(1)) % cg1_checkpoint_chunks for local_chunk_idx in cutlass.range(1, num_chunks_tile, 1, unroll=1): - chunk_idx = cstart + local_chunk_idx + chunk_idx = compute_start + local_chunk_idx cum_chunk = cum_chunk_base + local_chunk_idx sV_ptr = sV_raw.data_ptr() + raw_index.idx * (cfg.d_v * cfg.b_t) sBeta_ptr = sBeta_raw.data_ptr() + raw_bar_index.idx * cfg.b_t @@ -1750,42 +1629,39 @@ def compute1_warp_group( do_checkpoint = cg1_checkpoint_mod == 0 cg1_checkpoint_mod = cg1_checkpoint_mod + cutlass.Int32(1) cg1_checkpoint_mod = cutlass.Int32(0) if cg1_checkpoint_mod == cg1_checkpoint_chunks else cg1_checkpoint_mod - do_checkpoint = do_checkpoint and chunk_idx >= wstart + do_checkpoint = do_checkpoint and chunk_idx >= write_start - # ---- state repack: acc TMEM -> packed b16 TMEM ---------------------- + # ---- state stage: acc TMEM -> packed b16 TMEM ---------------------------- if cutlass.const_expr(not cfg.enable_checkpoints): bars.mb_state_acc_done.wait(state_upd_index.phase) state_upd_index = advance(state_upd_index, 1) state_vecs = [] - for sub in cutlass.range_constexpr(cfg.d_k // 16): - state_vecs.append(nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + state_col_id + sub * 16, cutlass.Float32), num=16)) + for i in cutlass.range_constexpr(cfg.d_k // 16): + state_vecs.append(nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_lo_addr + state_col_id + i * 16, cutlass.Float32), num=16)) - for sub in cutlass.range_constexpr(cfg.d_k // 16): + for i in cutlass.range_constexpr(cfg.d_k // 16): packed_state = cutlass.Array(cutlass.Int32, 8, alignment=16) for packed_col in cutlass.range_constexpr(8): - source_pair = packed_col ^ 4 - packed_state[packed_col] = fp32_to_fp16(state_vecs[sub][2 * source_pair], state_vecs[sub][2 * source_pair + 1], dtype=cfg.io_dtype) + packed_state[packed_col] = fp32_to_fp16(state_vecs[i][2 * packed_col], state_vecs[i][2 * packed_col + 1], dtype=cfg.io_dtype) nvvm.tcgen05_st( "32x32b", - nvvm.make_tmem_ptr((tmem_row << 16) + packed_col_id + sub * 8, cutlass.Int8), + nvvm.make_tmem_ptr(row_lo_addr + packed_col_id + i * 8, cutlass.Int8), packed_state[0:8], ) nvvm.tcgen05_wait("store") - bars.mb_state_inp_ready.arrive() + bars.mb_state_input_ready.arrive() - # ---- checkpoint store ----------------------------------------------- + # ---- checkpoint store ---------------------------------------------------- if cutlass.const_expr(cfg.enable_checkpoints): if do_checkpoint: checkpoint_stage = checkpoint_done_index.idx bars.mb_checkpoint_tmastg_done[checkpoint_stage].wait(checkpoint_done_index.phase) checkpoint_done_index = advance(checkpoint_done_index, cfg.smem_checkpoint_stages) checkpoint_stage_base = checkpoint_stage * (cfg.d_k * cfg.d_v) - for sub in cutlass.range_constexpr(cfg.d_k // 16): + for i in cutlass.range_constexpr(cfg.d_k // 16): for g in cutlass.range_constexpr(2): - packs = tuple( - fp32_to_fp16(state_vecs[sub][g * 8 + 2 * t], state_vecs[sub][g * 8 + 2 * t + 1], dtype=cfg.io_dtype) for t in range(4) - ) - dk = sub * 16 + g * 8 + packs = tuple(fp32_to_fp16(state_vecs[i][g * 8 + 2 * t], state_vecs[i][g * 8 + 2 * t + 1], dtype=cfg.io_dtype) for t in range(4)) + dk = i * 16 + g * 8 checkpoint_addr = ( checkpoint_stage_base + (dk // 64) * (cfg.d_v * 64) + value_dim * 64 + swizzle_xor_128b(value_dim, dk % 64, elem_bytes=2) ) @@ -1800,28 +1676,28 @@ def compute1_warp_group( bars.mb_o_acc_ready.wait(o_acc_index.phase) o_acc_index = advance(o_acc_index, 1) projection_col_id = q_state_col_base + prev_q_state_acc_stage * cfg.b_t - loaded_vec0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_addr + projection_col_id, cutlass.Float32), num=2) - loaded_vec1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row16_addr + projection_col_id, cutlass.Float32), num=2) + loaded_vec_lo = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_lo_addr + projection_col_id, cutlass.Float32), num=2) + loaded_vec_hi = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_hi_addr + projection_col_id, cutlass.Float32), num=2) - # ---- output drain: O acc TMEM -> scaled b16 SMEM -------------------- - stsm_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) - stsm_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + # ---- output store: O acc TMEM -> scaled b16 SMEM ------------------------- + o_pack_lo = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + o_pack_hi = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) for reg_idx in cutlass.range_constexpr(4): - scaled0_0, scaled0_1 = fmul2(loaded_vec0[2 * reg_idx], loaded_vec0[2 * reg_idx + 1], scale, scale) - scaled1_0, scaled1_1 = fmul2(loaded_vec1[2 * reg_idx], loaded_vec1[2 * reg_idx + 1], scale, scale) - stsm_pack0[reg_idx] = fp32_to_fp16(scaled0_0, scaled0_1, dtype=mO.element_type) - stsm_pack1[reg_idx] = fp32_to_fp16(scaled1_0, scaled1_1, dtype=mO.element_type) + scaled0_0, scaled0_1 = fmul2(loaded_vec_lo[2 * reg_idx], loaded_vec_lo[2 * reg_idx + 1], scale, scale) + scaled1_0, scaled1_1 = fmul2(loaded_vec_hi[2 * reg_idx], loaded_vec_hi[2 * reg_idx + 1], scale, scale) + o_pack_lo[reg_idx] = fp32_to_fp16(scaled0_0, scaled0_1, dtype=mO.element_type) + o_pack_hi[reg_idx] = fp32_to_fp16(scaled1_0, scaled1_1, dtype=mO.element_type) bars.mb_o_tmastg_done[prev_o_stage].wait(((prev_cum_chunk // cfg.smem_o_stages) + 1) % 2) nvvm.stmatrix( - sO_ptr + prev_o_stage_base + ov_swz_off0, - stsm_pack0.data_ptr().load(count=4, alignment=4), + sO_ptr + prev_o_stage_base + ov_swz_off_lo, + o_pack_lo.data_ptr().load(count=4, alignment=4), nvvm.MMALayout.COL, shape=nvvm.StoreShape.M8N8, ) nvvm.stmatrix( - sO_ptr + prev_o_stage_base + ov_swz_off, - stsm_pack1.data_ptr().load(count=4, alignment=4), + sO_ptr + prev_o_stage_base + ov_swz_off_hi, + o_pack_hi.data_ptr().load(count=4, alignment=4), nvvm.MMALayout.COL, shape=nvvm.StoreShape.M8N8, ) @@ -1829,73 +1705,71 @@ def compute1_warp_group( nvvm.fence_proxy("async.shared", space="cta") bars.mb_o_tmastg_ready[prev_o_stage].arrive() - # ---- Y staging: Y = Beta * (V - state*K) ----------------------------- + # ---- Y stage: Y = Beta * (V - k state) ----------------------------------- bars.mb_v_ready[raw_bar_index.idx].wait(raw_bar_index.phase) - raw_v_frag0 = nvvm.ldmatrix( - sV_ptr + ov_swz_off0, + raw_v_frag_lo = nvvm.ldmatrix( + sV_ptr + ov_swz_off_lo, 4, nvvm.MMALayout.COL, ) - raw_v_frag1 = nvvm.ldmatrix( - sV_ptr + ov_swz_off, + raw_v_frag_hi = nvvm.ldmatrix( + sV_ptr + ov_swz_off_hi, 4, nvvm.MMALayout.COL, ) bars.mb_beta_ready[raw_bar_index.idx].wait(raw_bar_index.phase) bars.mb_state_k_acc_ready.wait(state_k_acc_index.phase) - state_k_vec0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_addr + state_k_col_id, cutlass.Float32), num=2) - state_k_vec1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row16_addr + state_k_col_id, cutlass.Float32), num=2) + state_k_vec_lo = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_lo_addr + state_k_col_id, cutlass.Float32), num=2) + state_k_vec_hi = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_hi_addr + state_k_col_id, cutlass.Float32), num=2) beta_pack = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) for reg_idx in cutlass.range_constexpr(4): - token0 = (((reg_idx // 2) * 4 + (lane & 3)) ^ 4) * 2 + token0 = ((reg_idx // 2) * 4 + (lane_idx & 3)) * 2 beta0 = (sBeta_ptr + token0).load().to(cutlass.Float32) beta1 = (sBeta_ptr + token0 + 1).load().to(cutlass.Float32) beta_pack[reg_idx] = fp32_to_fp16(beta0, beta1, dtype=cfg.io_dtype) - y_inp_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + y_input_pack_lo = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) for reg_idx in cutlass.range_constexpr(4): - raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1) - frag_pair = (reg_idx ^ 2) * 2 - state_k_val0, state_k_val1 = state_k_vec0[frag_pair], state_k_vec0[frag_pair + 1] + frag_pair = reg_idx * 2 + state_k_val0, state_k_val1 = state_k_vec_lo[frag_pair], state_k_vec_lo[frag_pair + 1] state_k_pair = fp32_to_fp16(state_k_val0, state_k_val1, dtype=cfg.io_dtype) diff_pair = sub_f16x2( - raw_v_frag0[raw_matrix], + raw_v_frag_lo[reg_idx], state_k_pair, cfg.io_dtype, ) - y_inp_pack0[reg_idx] = mul_f16x2( + y_input_pack_lo[reg_idx] = mul_f16x2( beta_pack[reg_idx], diff_pair, cfg.io_dtype, ) - y_inp_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + y_input_pack_hi = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) for reg_idx in cutlass.range_constexpr(4): - raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1) - frag_pair = (reg_idx ^ 2) * 2 - state_k_val0, state_k_val1 = state_k_vec1[frag_pair], state_k_vec1[frag_pair + 1] + frag_pair = reg_idx * 2 + state_k_val0, state_k_val1 = state_k_vec_hi[frag_pair], state_k_vec_hi[frag_pair + 1] state_k_pair = fp32_to_fp16(state_k_val0, state_k_val1, dtype=cfg.io_dtype) diff_pair = sub_f16x2( - raw_v_frag1[raw_matrix], + raw_v_frag_hi[reg_idx], state_k_pair, cfg.io_dtype, ) - y_inp_pack1[reg_idx] = mul_f16x2( + y_input_pack_hi[reg_idx] = mul_f16x2( beta_pack[reg_idx], diff_pair, cfg.io_dtype, ) - nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(st_row_addr + y_inp_col_id, cutlass.Int8), y_inp_pack0[0:4]) - nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(st_row16_addr + y_inp_col_id, cutlass.Int8), y_inp_pack1[0:4]) + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_lo_addr + y_input_col_id, cutlass.Int8), y_input_pack_lo[0:4]) + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_hi_addr + y_input_col_id, cutlass.Int8), y_input_pack_hi[0:4]) nvvm.tcgen05_wait("store") state_k_acc_index = advance(state_k_acc_index, 1) bars.mb_v_done[raw_index.idx].arrive() bars.mb_beta_done[raw_bar_index.idx].arrive() - bars.mb_y_inp_ready.arrive() + bars.mb_y_input_ready.arrive() - # ---- U repack: u_acc TMEM -> packed b16 U input TMEM ---------------- + # ---- U stage: u acc TMEM -> packed b16 U input TMEM ---------------------- bars.mb_u_acc_ready.wait(u_acc_index.phase) u_vals = nvvm.tcgen05_ld( "32x32b", @@ -1903,21 +1777,20 @@ def compute1_warp_group( num=cfg.b_t, ) - u_inp_pack = cutlass.Array(cutlass.Int32, (cfg.b_t // 2), alignment=16) + u_input_pack = cutlass.Array(cutlass.Int32, (cfg.b_t // 2), alignment=16) for packed_col in cutlass.range_constexpr((cfg.b_t // 2)): - source_pair = packed_col ^ 4 - token0 = source_pair * 2 + token0 = packed_col * 2 token1 = token0 + 1 - u_inp_pack[packed_col] = fp32_to_fp16(u_vals[token0], u_vals[token1], dtype=cfg.io_dtype) + u_input_pack[packed_col] = fp32_to_fp16(u_vals[token0], u_vals[token1], dtype=cfg.io_dtype) nvvm.tcgen05_st( "32x32b", - nvvm.make_tmem_ptr(u_inp_addr, cutlass.Int8), - u_inp_pack[0 : (cfg.b_t // 2)], + nvvm.make_tmem_ptr(u_input_addr, cutlass.Int8), + u_input_pack[0 : (cfg.b_t // 2)], ) nvvm.tcgen05_wait("store") u_acc_index = advance(u_acc_index, 1) - bars.mb_u_inp_ready.arrive() + bars.mb_u_input_ready.arrive() if cutlass.const_expr(cfg.enable_checkpoints): bars.mb_state_acc_done.wait(state_upd_index.phase) @@ -1938,28 +1811,28 @@ def compute1_warp_group( o_acc_index = advance(o_acc_index, 1) projection_col_id = q_state_col_base + final_q_state_acc_stage * cfg.b_t - loaded_vec0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_addr + projection_col_id, cutlass.Float32), num=2) - loaded_vec1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row16_addr + projection_col_id, cutlass.Float32), num=2) + loaded_vec_lo = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_lo_addr + projection_col_id, cutlass.Float32), num=2) + loaded_vec_hi = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_hi_addr + projection_col_id, cutlass.Float32), num=2) - # ---- output drain: O acc TMEM -> scaled b16 SMEM -------------------- - stsm_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) - stsm_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + # ---- output store: O acc TMEM -> scaled b16 SMEM ------------------------- + o_pack_lo = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + o_pack_hi = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) for reg_idx in cutlass.range_constexpr(4): - scaled0_0, scaled0_1 = fmul2(loaded_vec0[2 * reg_idx], loaded_vec0[2 * reg_idx + 1], scale, scale) - scaled1_0, scaled1_1 = fmul2(loaded_vec1[2 * reg_idx], loaded_vec1[2 * reg_idx + 1], scale, scale) - stsm_pack0[reg_idx] = fp32_to_fp16(scaled0_0, scaled0_1, dtype=mO.element_type) - stsm_pack1[reg_idx] = fp32_to_fp16(scaled1_0, scaled1_1, dtype=mO.element_type) + scaled0_0, scaled0_1 = fmul2(loaded_vec_lo[2 * reg_idx], loaded_vec_lo[2 * reg_idx + 1], scale, scale) + scaled1_0, scaled1_1 = fmul2(loaded_vec_hi[2 * reg_idx], loaded_vec_hi[2 * reg_idx + 1], scale, scale) + o_pack_lo[reg_idx] = fp32_to_fp16(scaled0_0, scaled0_1, dtype=mO.element_type) + o_pack_hi[reg_idx] = fp32_to_fp16(scaled1_0, scaled1_1, dtype=mO.element_type) bars.mb_o_tmastg_done[final_o_stage].wait(((last_cum_chunk // cfg.smem_o_stages) + 1) % 2) nvvm.stmatrix( - sO_ptr + final_o_stage_base + ov_swz_off0, - stsm_pack0.data_ptr().load(count=4, alignment=4), + sO_ptr + final_o_stage_base + ov_swz_off_lo, + o_pack_lo.data_ptr().load(count=4, alignment=4), nvvm.MMALayout.COL, shape=nvvm.StoreShape.M8N8, ) nvvm.stmatrix( - sO_ptr + final_o_stage_base + ov_swz_off, - stsm_pack1.data_ptr().load(count=4, alignment=4), + sO_ptr + final_o_stage_base + ov_swz_off_hi, + o_pack_hi.data_ptr().load(count=4, alignment=4), nvvm.MMALayout.COL, shape=nvvm.StoreShape.M8N8, ) @@ -1967,18 +1840,18 @@ def compute1_warp_group( nvvm.fence_proxy("async.shared", space="cta") bars.mb_o_tmastg_ready[final_o_stage].arrive() - owns_final = wend == num_chunks_b + owns_final = write_end == batch_num_chunks - # ---- final-state drain: state acc TMEM -> GMEM --------------------------- + # ---- final state store: TMEM -> GMEM ----------------------------------------- if cutlass.const_expr(mState_out is not None): - if seqlen_b > 0: + if batch_seqlen > 0: if owns_final: state_vw = 16 // (mState_out.element_type.width // 8) state_dst = (mState_out.iterator + mState_out.layout((batch_idx, head_o, value_dim, 0))).raw_ptr() for key_block_start in cutlass.range_constexpr(0, cfg.d_k, 32): loaded = nvvm.tcgen05_ld( "32x32b", - nvvm.make_tmem_ptr((row_id << 16) + (tmem_col + cfg.tmem_state_acc_offset + key_block_start), cutlass.Float32), + nvvm.make_tmem_ptr(row_lo_addr + (tmem_col + cfg.tmem_state_acc_offset + key_block_start), cutlass.Float32), num=32, ) @@ -1999,7 +1872,7 @@ def compute1_warp_group( else: mState_out[batch_idx, head_o, value_dim, key_dim] = cutlass.Float32(0.0).to(mState_out.element_type) cum_chunk_base += num_chunks_tile - tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one) bars.mb_tmem_done[0].arrive() @@ -2020,7 +1893,7 @@ def host( final_state: cute.Tensor | None, work_items: cute.Tensor | None, work_count: cute.Tensor | None, - sched_ctr: cute.Tensor | None, + scheduler_counter: cute.Tensor | None, tensormap_workspace: cute.Tensor, checkpoint_every_n_tokens: cutlass.Int32, scale: cutlass.Float32, @@ -2046,7 +1919,7 @@ def host( final_state, work_items, work_count, - sched_ctr, + scheduler_counter, scale, checkpoint_every_n_tokens, ).launch( @@ -2075,7 +1948,7 @@ def kernel( mState_out: cute.Tensor | None, mWorkItems: cute.Tensor, mCount: cute.Tensor, - mSched: cute.Tensor | None, + mScheduler: cute.Tensor | None, scale: cutlass.Float32, checkpoint_every_n_tokens: cutlass.Int32, ) -> None: @@ -2085,11 +1958,11 @@ def kernel( bidx = cute.arch.block_idx()[0] num_ctas = cute.arch.grid_dim()[0] warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - lane = tidx % cfg.threads_per_warp + lane_idx = tidx % cfg.threads_per_warp total_tiles = mCount[0] - if cutlass.const_expr(cfg.dyn_sched): - assert mSched is not None and mSched.element_type == cutlass.Int32 + if cutlass.const_expr(cfg.dynamic_scheduling): + assert mScheduler is not None and mScheduler.element_type == cutlass.Int32 assert mQ.element_type == cfg.io_dtype and mK.element_type == cfg.io_dtype and mV.element_type == cfg.io_dtype assert mGate.element_type == cutlass.Float32 beta_expected = cfg.io_dtype if cutlass.const_expr(cfg.beta_sigmoid) else cutlass.Float32 @@ -2114,11 +1987,10 @@ def kernel( desc_o_base = desc_base_words + cutlass.Int32(4) * arr_words desc_checkpoint_base = desc_base_words + cutlass.Int32(5) * arr_words - # Buffers are declaration-ordered and intentionally non-aliased. SMEM = cutlass.AddressSpace.smem bars = make_kda_bars(cfg) tmem_base_slot = cutlass.Array(cutlass.Int32, 1, space=SMEM, alignment=4) - sSched = cutlass.Array(cutlass.Int32, cfg.sched_stages, space=SMEM, alignment=16) + sScheduler = cutlass.Array(cutlass.Int32, cfg.scheduler_stages, space=SMEM, alignment=16) sK_decay_raw = cutlass.Array(cfg.io_dtype, cfg.k_decay_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) sQ_decay_raw = cutlass.Array(cfg.io_dtype, cfg.q_decay_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) sK_restore_raw = cutlass.Array(cfg.io_dtype, cfg.k_restore_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) @@ -2157,7 +2029,7 @@ def kernel( stride_byte_offset=1024, layout=nvvm.Tcgen05SmemSwizzle.SWIZZLE_128B, ) - sK_restore = SmemTile( + sK_restore_trans = SmemTile( base=sK_restore_raw, elems_per_stage=(cfg.d_k * cfg.b_t), stages=cfg.smem_decay_stages, @@ -2205,15 +2077,15 @@ def kernel( bars.mb_state_k_acc_ready.init() bars.mb_u_acc_ready.init() bars.mb_state_acc_done.init() - bars.mb_state_inp_ready.init() + bars.mb_state_input_ready.init() for stage in cutlass.range_constexpr(cfg.smem_state_scale_diag_stages): bars.mb_state_scale_diag_done[stage].init() for stage in cutlass.range_constexpr(cfg.smem_decay_stages): bars.mb_decay_tcgen05_done[stage].init() bars.mb_decay_super_done[stage].init() bars.mb_k_restore_done[stage].init() - bars.mb_y_inp_ready.init() - bars.mb_u_inp_ready.init() + bars.mb_y_input_ready.init() + bars.mb_u_input_ready.init() bars.mb_tmem_done[0].init() elif warp_idx == cfg.super_mma_warp_id: if elect_one: @@ -2231,9 +2103,9 @@ def kernel( for stage in cutlass.range_constexpr(cfg.smem_o_stages): bars.mb_o_tmastg_ready[stage].init() bars.mb_o_tmastg_done[stage].init() - for stage in cutlass.range_constexpr(cfg.sched_stages): - bars.mb_sched_ready[stage].init() - bars.mb_sched_done[stage].init() + for stage in cutlass.range_constexpr(cfg.scheduler_stages): + bars.mb_scheduler_ready[stage].init() + bars.mb_scheduler_done[stage].init() if cutlass.const_expr(cfg.enable_checkpoints): for stage in cutlass.range_constexpr(cfg.smem_checkpoint_stages): bars.mb_checkpoint_tmastg_ready[stage].init() @@ -2252,9 +2124,9 @@ def kernel( num_ctas, cu_seqlens, mWorkItems, - mSched, - sSched, - lane, + mScheduler, + sScheduler, + lane_idx, sQ_raw, sK_raw, sV_raw, @@ -2273,8 +2145,8 @@ def kernel( num_ctas, cu_seqlens, mWorkItems, - sSched, - lane, + sScheduler, + lane_idx, sK_inv_raw, sIntermediate_raw, sBeta_raw, @@ -2289,11 +2161,11 @@ def kernel( num_ctas, cu_seqlens, mWorkItems, - sSched, + sScheduler, tmem_base_slot, sIntermediate, sK_decay, - sK_restore, + sK_restore_trans, sQ_decay, sState_scale_diag, bars, @@ -2306,8 +2178,8 @@ def kernel( num_ctas, cu_seqlens, mWorkItems, - sSched, - lane, + sScheduler, + lane_idx, mO, sK_inv_raw, sO_raw, @@ -2327,8 +2199,8 @@ def kernel( num_ctas, cu_seqlens, mWorkItems, - sSched, - lane, + sScheduler, + lane_idx, warp_idx, mQ, mA_log, @@ -2353,8 +2225,8 @@ def kernel( num_ctas, cu_seqlens, mWorkItems, - sSched, - lane, + sScheduler, + lane_idx, tmem_base_slot, warp_idx, mState_out, @@ -2391,8 +2263,8 @@ class KdaCfg: v_ratio: int n_heads_out: int max_active_clusters: int - dyn_sched: bool = False - sched_stages: int = CFG.SMEM_SCHED_STAGES + dynamic_scheduling: bool = False + scheduler_stages: int = CFG.SMEM_SCHEDULER_STAGES compute_group_0_warp_ids: tuple[int, ...] = CFG.COMPUTE_GROUP_0_WARP_IDS compute_group_1_warp_ids: tuple[int, ...] = CFG.COMPUTE_GROUP_1_WARP_IDS @@ -2409,7 +2281,7 @@ class KdaCfg: cg0_group_count: int = 2 cg0_warps_per_group: int = 4 cg0_threads_per_group: int = 0 - cg0_group_sync_barrier_base_id: int = 1 # CG0 group g syncs on nbar id 1 + g + cg0_group_sync_barrier_base_id: int = 1 # CG0 group g syncs on named-barrier id 1 + g cg0_tile_entry_barrier_id: int = 5 # CG0-wide (both groups) work-item entry sync tmem_user_threads: int = 0 tmem_lifecycle_barrier_id: int = 3 @@ -2417,9 +2289,9 @@ class KdaCfg: num_regs_compute_group_1: int = CFG.NUM_REGS_COMPUTE_GROUP_1 num_regs_other: int = CFG.NUM_REGS_OTHER - # ---- SMEM / TMEM ring stage counts ------------------------------------------- + # ---- SMEM / TMEM ring stage counts ----------------------------------------------- smem_raw_stages: int = CFG.SMEM_RAW_STAGES - smem_raw_bar_stages: int = 0 # ready/beta-ring mbar depth: raw rounded up to even (CG0 ping-pong parity) + smem_raw_bar_stages: int = 0 # ready/beta-ring mbar depth: raw rounded up to even smem_o_stages: int = CFG.SMEM_O_STAGES smem_checkpoint_stages: int = 1 smem_decay_stages: int = CFG.SMEM_DECAY_STAGES @@ -2428,16 +2300,16 @@ class KdaCfg: qk_scale_ready_stages: int = CFG.QK_SCALE_READY_STAGES tmem_q_state_acc_stages: int = CFG.TMEM_Q_STATE_ACC_STAGES - # ---- TMEM column offsets (state doubles as the final_state acc) -------------- + # ---- TMEM column offsets (state doubles as the final state acc) ------------------ tmem_state_acc_offset: int = 0 - tmem_state_inp_offset: int = 0 + tmem_state_input_offset: int = 0 tmem_q_state_acc_offset: int = 0 tmem_state_k_acc_offset: int = 0 tmem_u_acc_offset: int = 0 - tmem_y_inp_offset: int = 0 - tmem_u_inp_offset: int = 0 + tmem_y_input_offset: int = 0 + tmem_u_input_offset: int = 0 - # ---- SMEM buffer cosizes ----------------------------------------------------- + # ---- SMEM buffer cosizes --------------------------------------------------------- q_cosize: int = 0 k_cosize: int = 0 v_cosize: int = 0 @@ -2474,7 +2346,7 @@ def build_cfg( v_ratio: int, n_heads_out: int, max_active_clusters: int, - dyn_sched: bool = False, + dynamic_scheduling: bool = False, ) -> KdaCfg: """Build the per-compile ``KdaCfg`` (io_dtype in {Float16, BFloat16}); fills the derived TMEM column offsets and SMEM buffer cosizes.""" @@ -2495,7 +2367,7 @@ def build_cfg( v_ratio=v_ratio, n_heads_out=n_heads_out, max_active_clusters=max_active_clusters, - dyn_sched=dyn_sched, + dynamic_scheduling=dynamic_scheduling, ) if enable_checkpoints: cfg.smem_raw_stages = 5 @@ -2507,13 +2379,13 @@ def build_cfg( if cfg.smem_state_scale_diag_stages != cfg.qk_scale_ready_stages: raise ValueError("diag and qk-scale ready rings must share their rolling stage") - cfg.tmem_state_inp_offset = cfg.tmem_state_acc_offset + cfg.d_k - cfg.tmem_q_state_acc_offset = cfg.tmem_state_inp_offset + (cfg.d_k // 2) + cfg.tmem_state_input_offset = cfg.tmem_state_acc_offset + cfg.d_k + cfg.tmem_q_state_acc_offset = cfg.tmem_state_input_offset + (cfg.d_k // 2) cfg.tmem_state_k_acc_offset = cfg.tmem_q_state_acc_offset + cfg.tmem_q_state_acc_stages * cfg.b_t cfg.tmem_u_acc_offset = cfg.tmem_state_k_acc_offset + cfg.b_t - cfg.tmem_y_inp_offset = cfg.tmem_u_acc_offset + cfg.b_t - cfg.tmem_u_inp_offset = cfg.tmem_y_inp_offset + (cfg.b_t // 2) - assert (cfg.tmem_u_inp_offset + (cfg.b_t // 2)) <= 512 + cfg.tmem_y_input_offset = cfg.tmem_u_acc_offset + cfg.b_t + cfg.tmem_u_input_offset = cfg.tmem_y_input_offset + (cfg.b_t // 2) + assert (cfg.tmem_u_input_offset + (cfg.b_t // 2)) <= 512 cfg.q_cosize = cfg.smem_raw_stages * cfg.d_k * cfg.b_t cfg.k_cosize = cfg.smem_raw_stages * cfg.d_k * cfg.b_t @@ -2547,7 +2419,7 @@ def build_descs_body( base_gate, base_o, base_checkpoint, - desc_ws: cute.Tensor, + desc_workspace: cute.Tensor, cu_seqlens: cute.Tensor, q: cute.Tensor, k: cute.Tensor, @@ -2568,12 +2440,12 @@ def build_descs_body( prologue kernel after its order pass; warps past the array count fall through the widx guards.""" arr_words = n_batch * cutlass.Int32(TENSOR_MAP_QWORDS) - desc_q_arr = cute.make_tensor(desc_ws.iterator, cute.make_layout((arr_words,), stride=(1,))) - desc_k_arr = cute.make_tensor(desc_ws.iterator + arr_words, cute.make_layout((arr_words,), stride=(1,))) - desc_v_arr = cute.make_tensor(desc_ws.iterator + 2 * arr_words, cute.make_layout((arr_words,), stride=(1,))) - desc_gate_arr = cute.make_tensor(desc_ws.iterator + 3 * arr_words, cute.make_layout((arr_words,), stride=(1,))) - desc_o_arr = cute.make_tensor(desc_ws.iterator + 4 * arr_words, cute.make_layout((arr_words,), stride=(1,))) - desc_checkpoint_arr = cute.make_tensor(desc_ws.iterator + 5 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + desc_q_arr = cute.make_tensor(desc_workspace.iterator, cute.make_layout((arr_words,), stride=(1,))) + desc_k_arr = cute.make_tensor(desc_workspace.iterator + arr_words, cute.make_layout((arr_words,), stride=(1,))) + desc_v_arr = cute.make_tensor(desc_workspace.iterator + 2 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + desc_gate_arr = cute.make_tensor(desc_workspace.iterator + 3 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + desc_o_arr = cute.make_tensor(desc_workspace.iterator + 4 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + desc_checkpoint_arr = cute.make_tensor(desc_workspace.iterator + 5 * arr_words, cute.make_layout((arr_words,), stride=(1,))) if widx == 0: if nvvm.elect_sync(): @@ -2607,7 +2479,7 @@ def build_descs_body( @cute.kernel def prologue_kernel( order_gen: cutlass.Constexpr[bool], - has_sched: cutlass.Constexpr[bool], + has_scheduler: cutlass.Constexpr[bool], b_t: cutlass.Constexpr[int], base_q: cutlass.GridConstant[cuda.tensor_map.TensorMap], base_k: cutlass.GridConstant[cuda.tensor_map.TensorMap], @@ -2615,7 +2487,7 @@ def prologue_kernel( base_gate: cutlass.GridConstant[cuda.tensor_map.TensorMap], base_o: cutlass.GridConstant[cuda.tensor_map.TensorMap], base_checkpoint: cutlass.GridConstant[cuda.tensor_map.TensorMap], - desc_ws: cute.Tensor, + desc_workspace: cute.Tensor, cu_seqlens: cute.Tensor, q: cute.Tensor, k: cute.Tensor, @@ -2626,7 +2498,7 @@ def prologue_kernel( mStaging: cute.Tensor | None, mCount: cute.Tensor, mWorkItems: cute.Tensor, - mSched: cute.Tensor | None, + mScheduler: cute.Tensor | None, n_batch: cutlass.Int32, q_token_stride: cutlass.Int32, k_token_stride: cutlass.Int32, @@ -2636,7 +2508,7 @@ def prologue_kernel( checkpoint_entry_stride: cutlass.Int32, checkpoint_every_n: cutlass.Int32, ) -> None: - """Single-CTA prologue: LPT-order the work-item table and zero the sched + """Single-CTA prologue: LPT-order the work-item table and zero the scheduler rings via :func:`order_body`, then build the per-batch TMA-descriptor arrays via :func:`build_descs_body`, one warp per array (the extra warps only take part in the order phase).""" @@ -2649,10 +2521,10 @@ def prologue_kernel( n_heads_out = cutlass.Int32(gate.shape[1]) order_body( order_gen, - has_sched, + has_scheduler, b_t, ORDER_THREADS, - ORDER_ELEMS, + ORDER_ELEMENTS, tidx, n_heads_out, n_heads_out * n_batch, @@ -2660,7 +2532,7 @@ def prologue_kernel( mStaging, mCount, mWorkItems, - mSched, + mScheduler, sKey, sIdx, sSpread, @@ -2673,7 +2545,7 @@ def prologue_kernel( base_gate, base_o, base_checkpoint, - desc_ws, + desc_workspace, cu_seqlens, q, k, @@ -2697,7 +2569,7 @@ def prologue( io_dtype: cutlass.Constexpr, b_t: cutlass.Constexpr[int], order_gen: cutlass.Constexpr[bool], - has_sched: cutlass.Constexpr[bool], + has_scheduler: cutlass.Constexpr[bool], q: cute.Tensor, k: cute.Tensor, v: cute.Tensor, @@ -2708,7 +2580,7 @@ def prologue( work_item_staging: cute.Tensor | None, work_count: cute.Tensor, work_items: cute.Tensor, - sched_ctr: cute.Tensor | None, + scheduler_counter: cute.Tensor | None, tensormap_workspace: cute.Tensor, checkpoint_every_n: cutlass.Int32, stream: cuda_driver.CUstream, @@ -2752,7 +2624,7 @@ def prologue( base_checkpoint = cuda.create_tensor_map_tiled_from_view(checkpoint_view, box_dims=(tma_granu_elems, d_k, 1, 1), stride_order=(0, 1, 2, 3), swizzle=swz) prologue_kernel( order_gen, - has_sched, + has_scheduler, b_t, base_q, base_k, @@ -2771,7 +2643,7 @@ def prologue( work_item_staging, work_count, work_items, - sched_ctr, + scheduler_counter, cutlass.Int32(batch_size), cutlass.Int32(q.stride[0]), cutlass.Int32(k.stride[0]), @@ -2783,7 +2655,7 @@ def prologue( ).launch(grid=(1, 1, 1), block=(ORDER_THREADS, 1, 1), stream=stream) -# ---- Torch adapter / host-side compilation --------------------------------------- +# ---- Torch adapter / host-side compilation ------------------------------------------- @lru_cache(maxsize=None) @@ -2801,7 +2673,7 @@ def get_compiled_cache( safe_gate: bool, gate_lower_bound: float, beta_sigmoid: bool, - dyn_sched: bool, + dynamic_scheduling: bool, order_gen: bool, ): """Return a mutable dict that lazily stores the compiled kernel.""" @@ -2822,7 +2694,7 @@ def compile( k_ratio: int, v_ratio: int, n_heads_out: int, - dyn_sched: bool = False, + dynamic_scheduling: bool = False, *, num_sm: int, q_cute, @@ -2838,8 +2710,8 @@ def compile( state_out_cute, work_items_cute=None, work_count_cute=None, - sched_ctr_cute=None, - tensormap_ws_cute, + scheduler_counter_cute=None, + tensormap_workspace_cute, checkpoint_every_n_tokens, scale, stream, @@ -2860,7 +2732,7 @@ def compile( v_ratio=v_ratio, n_heads_out=n_heads_out, max_active_clusters=num_sm, - dyn_sched=dyn_sched, + dynamic_scheduling=dynamic_scheduling, ) return cute.compile( @@ -2879,8 +2751,8 @@ def compile( state_out_cute, work_items_cute, work_count_cute, - sched_ctr_cute, - tensormap_ws_cute, + scheduler_counter_cute, + tensormap_workspace_cute, checkpoint_every_n_tokens, scale, stream, @@ -2909,7 +2781,7 @@ def chunk_kda_sm100( use_beta_sigmoid_in_kernel: bool = False, work_items=None, work_count=None, - sched_ctr=None, + scheduler_counter=None, work_item_scratch=None, *, tensormap_workspace, @@ -2951,13 +2823,14 @@ def chunk_kda_sm100( use_beta_sigmoid_in_kernel: ``beta`` holds logits; sigmoid in-kernel work_items: ``(max_items, 8)`` int32 work-item table from ``common/split_k.py`` (REQUIRED; an uncut table row is the whole - (b, h) sequence). Each item computes chunks ``[cstart, wend)`` - and writes O/checkpoints only for ``[wstart, wend)``. + (b, h) sequence). Each item computes chunks ``[compute_start, write_end)`` + and writes O/checkpoints only for ``[write_start, write_end)``. work_count: ``(1,)`` int32 device-side item count (REQUIRED) - sched_ctr: ``(2,)`` int32 device scratch ``[ticket, done]`` enabling + scheduler_counter: ``(2,)`` int32 device scratch ``[ticket, done]`` enabling the dynamic (work-stealing) tile scheduler; must be zeroed before - every launch (``build_split_table`` does this when it is passed as - ``sched_ctr``). None keeps the static CTA stride. + every launch (the split-table stage and the order-generating + prologue both zero it when passed as ``scheduler_counter``). + None keeps the static CTA stride. """ HQ = q.shape[1] HK = k.shape[1] @@ -2974,8 +2847,10 @@ def chunk_kda_sm100( f"output_state_checkpoints dtype must match the io dtype (fp32 state belongs to output_state): got {output_state_checkpoints.dtype} with io {q.dtype}" ) if work_items is None or work_count is None: - raise ValueError("work_items/work_count are required (the split-table stage builds them for every launch)") - dyn_sched = sched_ctr is not None + raise ValueError( + "work_items/work_count are required (built by the split-table stage, or by an order-generating prologue when work_item_scratch is None)" + ) + dynamic_scheduling = scheduler_counter is not None order_gen = work_item_scratch is None if initial_state is not None: @@ -3014,7 +2889,7 @@ def chunk_kda_sm100( safe_gate, gate_lower_bound, use_beta_sigmoid_in_kernel, - dyn_sched, + dynamic_scheduling, order_gen, ) @@ -3043,11 +2918,11 @@ def chunk_kda_sm100( work_items_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1), divisibility=1) work_count_cute = from_dlpack(work_count, assumed_align=4).mark_layout_dynamic() - sched_ctr_cute = None - if dyn_sched: - sched_ctr_cute = from_dlpack(sched_ctr, assumed_align=4).mark_layout_dynamic() + scheduler_counter_cute = None + if dynamic_scheduling: + scheduler_counter_cute = from_dlpack(scheduler_counter, assumed_align=4).mark_layout_dynamic() - tensormap_ws_cute = from_dlpack(tensormap_workspace, assumed_align=128).mark_layout_dynamic() + tensormap_workspace_cute = from_dlpack(tensormap_workspace, assumed_align=128).mark_layout_dynamic() cache["compiled"] = compile( io_dtype, @@ -3063,7 +2938,7 @@ def chunk_kda_sm100( k_ratio, v_ratio, HO, - dyn_sched, + dynamic_scheduling, num_sm=multiprocessor_count(current_device()), q_cute=q_cute, k_cute=k_cute, @@ -3078,8 +2953,8 @@ def chunk_kda_sm100( state_out_cute=state_out_cute, work_items_cute=work_items_cute, work_count_cute=work_count_cute, - sched_ctr_cute=sched_ctr_cute, - tensormap_ws_cute=tensormap_ws_cute, + scheduler_counter_cute=scheduler_counter_cute, + tensormap_workspace_cute=tensormap_workspace_cute, checkpoint_every_n_tokens=checkpoint_every_n_tokens, scale=scale, stream=cu_stream, @@ -3095,7 +2970,7 @@ def chunk_kda_sm100( gate_pl = from_dlpack(gate, assumed_align=16).mark_layout_dynamic(leading_dim=2) o_pl = from_dlpack(output, assumed_align=16).mark_layout_dynamic(leading_dim=2) cu_pl = from_dlpack(cu_seqlens, assumed_align=8).mark_layout_dynamic() - ws_pl = from_dlpack(tensormap_workspace, assumed_align=128).mark_layout_dynamic() + workspace_pl = from_dlpack(tensormap_workspace, assumed_align=128).mark_layout_dynamic() state_checkpoints_pl = None if state_checkpoints_for_descs is not None: state_checkpoints_pl = from_dlpack(state_checkpoints_for_descs, assumed_align=16).mark_layout_dynamic(leading_dim=3) @@ -3106,15 +2981,15 @@ def chunk_kda_sm100( work_items_pl = from_dlpack(work_items, assumed_align=16) work_items_pl.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1), divisibility=1) work_count_pl = from_dlpack(work_count, assumed_align=4).mark_layout_dynamic() - sched_pl = None - if dyn_sched: - sched_pl = from_dlpack(sched_ctr, assumed_align=4).mark_layout_dynamic() + scheduler_pl = None + if dynamic_scheduling: + scheduler_pl = from_dlpack(scheduler_counter, assumed_align=4).mark_layout_dynamic() cache["prologue"] = cute.compile( prologue, io_dtype, CFG.B_T, order_gen, - dyn_sched, + dynamic_scheduling, q_pl, k_pl, v_pl, @@ -3125,8 +3000,8 @@ def chunk_kda_sm100( staging_pl, work_count_pl, work_items_pl, - sched_pl, - ws_pl, + scheduler_pl, + workspace_pl, cutlass.Int32(checkpoint_every_n_tokens), cu_stream, options="--enable-tvm-ffi", @@ -3142,7 +3017,7 @@ def chunk_kda_sm100( work_item_scratch if not order_gen else None, work_count, work_items, - sched_ctr, + scheduler_counter, tensormap_workspace, checkpoint_every_n_tokens, cu_stream, @@ -3161,7 +3036,7 @@ def chunk_kda_sm100( output_state if store_final_state else None, work_items, work_count, - sched_ctr, + scheduler_counter, tensormap_workspace, checkpoint_every_n_tokens, scale, @@ -3186,7 +3061,7 @@ def run_prefill( output_state_checkpoints, work_items, work_count, - sched_ctr, + scheduler_counter, work_item_scratch, tensormap_workspace, checkpoint_every_n_tokens, @@ -3208,7 +3083,7 @@ def run_prefill( work_item_scratch, work_count, work_items, - sched_ctr, + scheduler_counter, tensormap_workspace, checkpoint_every_n_tokens, cu_stream, @@ -3227,7 +3102,7 @@ def run_prefill( output_state, work_items, work_count, - sched_ctr, + scheduler_counter, tensormap_workspace, checkpoint_every_n_tokens, scale, diff --git a/python/cudnn/linear_attention/frost/kernel/kda_recompute_config.py b/python/cudnn/linear_attention/frost/kernel/kda_recompute_config.py index 0e39486e8..3dcc6a3bd 100644 --- a/python/cudnn/linear_attention/frost/kernel/kda_recompute_config.py +++ b/python/cudnn/linear_attention/frost/kernel/kda_recompute_config.py @@ -54,7 +54,7 @@ class Cfg: # --- SMEM / TMEM ring stage counts --- SMEM_RAW_STAGES: int = 8 - SMEM_SCHED_STAGES: int = 8 + SMEM_SCHEDULER_STAGES: int = 8 SMEM_DECAY_STAGES: int = 2 SMEM_INTERMEDIATE_STAGES: int = 2 SMEM_STATE_SCALE_DIAG_STAGES: int = 4 diff --git a/python/cudnn/linear_attention/frost/kernel/kda_recompute_f16.py b/python/cudnn/linear_attention/frost/kernel/kda_recompute_f16.py index 404707b2d..cd43c9ea5 100644 --- a/python/cudnn/linear_attention/frost/kernel/kda_recompute_f16.py +++ b/python/cudnn/linear_attention/frost/kernel/kda_recompute_f16.py @@ -15,74 +15,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Chunked Kimi Delta Attention (KDA) recompute (state/checkpoints-only) kernel for -Blackwell SM100/SM103 (Cutlass DSL), BT=16 tiling with a per-key-channel decay. -Framework-neutral entry ``chunk_kda_recompute_sm100``. - -Persistent kernel: the grid is the SM count and every warp role -runs a tile-scheduler loop (``decode_work_item``); a tile is one (batch, -head) sequence, or one split-K work item computing chunks ``[cstart, wend)`` -and writing checkpoints only for the owned ``[wstart, wend)`` (see -``common/split_k.py``; warmup chunks rebuild the incoming state from -zero). All ring stage/phase bookkeeping runs on cumulative per-CTA chunk -counters so pipelines flow seamlessly across tiles. - -Pipeline (direct CUTLASS primitives, chunk_idx-size 16 KDA schedule): - - load K/V/Gate/Beta - optional in-kernel L2-norm of K (L2NORM specialization) - exp2(g), exp2(-g), stage final-token exp2(g) as exp2(g_last) - super-MMA: KK/Neumann inverse + apply Beta - tcgen05-MMA: State*K / U / state update - store periodic state checkpoints, final state - -ABI: k `[T, HK, DK]`, v `[T, HV, DV]`, gate -`[T, HO, DK]` fp32 (natural-log decay unless SAFE_GATE, which applies the -safe-gate transform from raw gate + a_log/dt_bias), beta `[T, HO]` fp32 -post-sigmoid, cu_seqlens int32, states/checkpoints `[N, HO, DV, DK]` (VK, k -contiguous). -GQA/GVA head broadcast follows repeat_interleave: source head = -head_idx // (HO // H_x). State presence, L2NORM, SAFE_GATE, checkpoints, and the -head ratios are compile-time specializations. - -Warp assignments (16 warps = 512 threads): - warps 0-7 : compute group 0 - Gate prefix scan + decay/restore operands - warps 8-11 : compute group 1 - TMEM value side, state stores - warp 12 : super-MMA - register-MMA KK^T + Neumann inverse - warp 13 : tcgen05-MMA - the four state GEMMs + the TMEM lifecycle - warp 14 : TMA load - per-chunk input G->S loads - warp 15 : epilogue - the checkpoint TMA store - -SMEM layout: - Buffer Bytes Stages - K / V raw 32768 8 <-- SW128 TMA ring (io dtype) - Gate raw 65536 8 <-- fp32 prefix-scan source - Beta 512 8 <-- fp32 per-token scalars - K_inv 8192 2 <-- token-major ldmatrix/tcgen05 B operand - K decay 8192 2 <-- tcgen05 SW128 K-box-major A/B operands - K restore 8192 2 <-- tcgen05 B operand for the state update - state-scale diag 12288 3 <-- per-k-atom decay diagonal blocks - T_inv 2048 2 <-- SW32 16x16 register-MMA tiles - -TMEM layout (240 of 512 columns): - Buffer Cols Purpose - state 0-127 state[DK,DV] fp32 recurrent state - state inp 128-191 packed b16 A operand view of the state - state_k_acc 192-207 State*K fp32 accumulator - u_acc 208-223 U fp32 accumulator - Y 224-231 packed b16 A operand: Beta * (V - State*K) - U input 232-239 packed b16 A operand: the b16 U repack - -GEMM schedule (tcgen05-MMA warp, in issue order per chunk): - State*K -> state_k_acc - State decay (diag blocks) - U = Y(T) @ T_inv -> u_acc - final_state += U @ K_restore - -Requires a cutlass DSL build providing `cutlass.experimental.*`; not -available in the pip nvidia-cutlass-dsl releases. -""" - from dataclasses import dataclass from functools import lru_cache from typing import Callable, NamedTuple, Optional, Type @@ -94,7 +26,7 @@ import cutlass.cute as cute from cutlass.cute.runtime import from_dlpack -from ..common.split_k import ORDER_CAPACITY, ORDER_ELEMS, ORDER_THREADS, decode_work_item, order_body +from ..common.split_k import ORDER_CAPACITY, ORDER_ELEMENTS, ORDER_THREADS, decode_work_item, order_body from ..common.host import get_dtype from cudnn.frost.buffers import data_ptr from cudnn.frost.device import current_device, multiprocessor_count @@ -109,9 +41,10 @@ ) from cudnn.frost.tile_dsl.handles import GmemTileTma, MmaDesc, SmemTile, tma_slice_runtime_desc from cudnn.frost.tile_dsl.mma import mma_step, mma_ts_step -from cudnn.frost.tile_dsl.swizzle import swizzle_lin_128b, swizzle_lin_S, swizzle_xor_128b +from cudnn.frost.tile_dsl.swizzle import swizzle_xor_128b, swizzle_xor_32b from cudnn.frost.tile_dsl.tma import tma_load_tile, tma_store_commit, tma_store_tile, tma_store_wait, tma_tensormap_acquire from cudnn.frost.tile_dsl.pointwise import ( + sigmoid, opaque_f32_zero, f16x2_to_f32, fadd2, @@ -124,20 +57,12 @@ ) LOG2_E: float = 1.4426950408889634 - - DEFAULT_GATE_LOWER_BOUND: float = -5.0 - - -# Host-side API defaults. - - L2_NORM_EPS: float = 1.0e-12 class KdaBars(NamedTuple): - """Every inter-warp handoff as an ``MBarrier`` over its ring. Consumers track ``(idx, phase)`` inline; the producer tag selects - the arrive lowering (``TMA_LOAD``/``MMA_COMMIT``/``THREAD``).""" + """Every inter-warp handoff as an ``MBarrier`` over its ring.""" mb_k_ready: MBarrier mb_k_done: MBarrier @@ -152,9 +77,9 @@ class KdaBars(NamedTuple): mb_state_k_acc_ready: MBarrier mb_u_acc_ready: MBarrier - mb_state_inp_ready: MBarrier - mb_y_inp_ready: MBarrier - mb_u_inp_ready: MBarrier + mb_state_input_ready: MBarrier + mb_y_input_ready: MBarrier + mb_u_input_ready: MBarrier mb_t_inv_ready: MBarrier mb_t_inv_done: MBarrier @@ -172,13 +97,12 @@ class KdaBars(NamedTuple): mb_checkpoint_tmastg_ready: MBarrier mb_checkpoint_tmastg_done: MBarrier - mb_sched_ready: MBarrier - mb_sched_done: MBarrier + mb_scheduler_ready: MBarrier + mb_scheduler_done: MBarrier def make_kda_bars(cfg) -> KdaBars: - """Bars factory. MUST be called from inside ``kernel`` (allocates the - mbarrier rings in SMEM ahead of the data buffers).""" + """KdaBars factory.""" def alloc(n): return cutlass.Array(cutlass.Int64, n, space=cutlass.AddressSpace.smem, alignment=8) @@ -198,9 +122,9 @@ def alloc(n): mb_beta_done=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=WARP + CG1_THREADS, producer=Producer.THREAD), mb_state_k_acc_ready=MBarrier(alloc(1), stages=1, init_count=1, producer=Producer.MMA_COMMIT), mb_u_acc_ready=MBarrier(alloc(1), stages=1, init_count=1, producer=Producer.MMA_COMMIT), - mb_state_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), - mb_y_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), - mb_u_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_state_input_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_y_input_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_u_input_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), mb_t_inv_ready=MBarrier(alloc(cfg.smem_intermediate_stages), stages=cfg.smem_intermediate_stages, init_count=WARP, producer=Producer.THREAD), mb_t_inv_done=MBarrier(alloc(cfg.smem_intermediate_stages), stages=cfg.smem_intermediate_stages, init_count=1, producer=Producer.MMA_COMMIT), mb_qk_scale_ready=MBarrier( @@ -226,40 +150,40 @@ def alloc(n): alloc(cfg.smem_checkpoint_stages), stages=cfg.smem_checkpoint_stages, init_count=CG1_THREADS, producer=Producer.THREAD ), mb_checkpoint_tmastg_done=MBarrier(alloc(cfg.smem_checkpoint_stages), stages=cfg.smem_checkpoint_stages, init_count=WARP, producer=Producer.THREAD), - mb_sched_ready=MBarrier(alloc(cfg.sched_stages), stages=cfg.sched_stages, init_count=1, producer=Producer.THREAD), - mb_sched_done=MBarrier(alloc(cfg.sched_stages), stages=cfg.sched_stages, init_count=15, producer=Producer.THREAD), + mb_scheduler_ready=MBarrier(alloc(cfg.scheduler_stages), stages=cfg.scheduler_stages, init_count=1, producer=Producer.THREAD), + mb_scheduler_done=MBarrier(alloc(cfg.scheduler_stages), stages=cfg.scheduler_stages, init_count=15, producer=Producer.THREAD), ) -# ---- Dynamic tile scheduler ------------------------------------------------------ +# ---- Dynamic tile scheduler ---------------------------------------------------------- @cute.jit -def sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ctas): +def scheduler_publish_next(cfg, bars, sScheduler, mScheduler, scheduler_state, tile_idx, num_ctas, elect_one): """TMA-warp side: pull the next tile off the global ticket, publish it.""" - if cutlass.const_expr(cfg.dyn_sched): - bars.mb_sched_done[sched_state.idx].wait(sched_state.phase) - if nvvm.elect_sync(): - fetched = cutlass.Int32(nvvm.atomicrmw("add", mSched.iterator, cutlass.Int32(1), mem_order="relaxed", syncscope="gpu")) - sSched[sched_state.idx] = num_ctas + fetched + if cutlass.const_expr(cfg.dynamic_scheduling): + bars.mb_scheduler_done[scheduler_state.idx].wait(scheduler_state.phase) + if elect_one: + fetched = cutlass.Int32(nvvm.atomicrmw("add", mScheduler.iterator, cutlass.Int32(1), mem_order="relaxed", syncscope="gpu")) + sScheduler[scheduler_state.idx] = num_ctas + fetched nvvm.bar_warp_sync(cute.arch.FULL_MASK) - next_tile = sSched[sched_state.idx] - if nvvm.elect_sync(): - bars.mb_sched_ready[sched_state.idx].arrive() - return next_tile, advance(sched_state, cfg.sched_stages) - return tile_idx + num_ctas, sched_state + next_tile = sScheduler[scheduler_state.idx] + if elect_one: + bars.mb_scheduler_ready[scheduler_state.idx].arrive() + return next_tile, advance(scheduler_state, cfg.scheduler_stages) + return tile_idx + num_ctas, scheduler_state @cute.jit -def sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas): +def scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one): """Consumer side: read the TMA warp's published next tile.""" - if cutlass.const_expr(cfg.dyn_sched): - bars.mb_sched_ready[sched_state.idx].wait(sched_state.phase) - next_tile = sSched[sched_state.idx] - if nvvm.elect_sync(): - bars.mb_sched_done[sched_state.idx].arrive() - return next_tile, advance(sched_state, cfg.sched_stages) - return tile_idx + num_ctas, sched_state + if cutlass.const_expr(cfg.dynamic_scheduling): + bars.mb_scheduler_ready[scheduler_state.idx].wait(scheduler_state.phase) + next_tile = sScheduler[scheduler_state.idx] + if elect_one: + bars.mb_scheduler_done[scheduler_state.idx].arrive() + return next_tile, advance(scheduler_state, cfg.scheduler_stages) + return tile_idx + num_ctas, scheduler_state @cute.jit @@ -270,9 +194,9 @@ def tmaldg_warp( num_ctas, cu_seqlens, mWorkItems, - mSched, - sSched, - lane, + mScheduler, + sScheduler, + lane_idx, sK_raw, sV_raw, sGate_raw, @@ -283,8 +207,12 @@ def tmaldg_warp( ) -> None: """TMA-LDG warp role (warp 14): persistent scheduler loop issuing the per-chunk K/V/Gate G->S loads.""" - elect_one = nvvm.elect_sync() nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) + + raw_index = PipelineState.start(phase=1) + scheduler_state = PipelineState.start(phase=1) + + elect_one = nvvm.elect_sync() sK_tma = SmemTile( base=sK_raw, elems_per_stage=(cfg.d_k * cfg.b_t), @@ -318,11 +246,11 @@ def tmaldg_warp( tma_granu_elems=32, tma_subtile_stride_elems=(cfg.b_t * 32), ) - raw_index = PipelineState.start(phase=1) - sched_state = PipelineState.start(phase=1) tile_idx = cutlass.Int32(bidx) while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) head_o = head_idx head_k = head_idx if cfg.k_ratio == 1 else head_idx // cutlass.Int32(cfg.k_ratio) head_v = head_idx if cfg.v_ratio == 1 else head_idx // cutlass.Int32(cfg.v_ratio) @@ -334,24 +262,24 @@ def tmaldg_warp( tma_tensormap_acquire(desc_k_slot) tma_tensormap_acquire(desc_v_slot) tma_tensormap_acquire(desc_gate_slot) - for chunk_idx in cutlass.range(cstart, wend, 1, unroll=1): + for chunk_idx in cutlass.range(compute_start, write_end, 1, unroll=1): chunk_start = chunk_idx * cfg.b_t - # ---- K load ---------------------------------------------------------- + # ---- K load -------------------------------------------------------------- bars.mb_k_done[raw_index.idx].wait(raw_index.phase) if elect_one: bars.mb_k_ready[raw_index.idx].arrive(n_bytes=cfg.tma_k_bytes) k_slice = tma_slice_runtime_desc(desc_k_slot, cutlass.Int32(0), head_k, chunk_start) tma_load_tile(sK_tma[raw_index.idx], k_slice, bars.mb_k_ready[raw_index.idx].smem_ptr, acquire=False) - # ---- Gate load ------------------------------------------------------- + # ---- Gate load: GMEM -> SMEM --------------------------------------------- bars.mb_gate_done[raw_index.idx].wait(raw_index.phase) if elect_one: bars.mb_gate_ready[raw_index.idx].arrive(n_bytes=cfg.tma_gate_bytes) gate_slice = tma_slice_runtime_desc(desc_gate_slot, cutlass.Int32(0), head_o, chunk_start) tma_load_tile(sGate_tma[raw_index.idx], gate_slice, bars.mb_gate_ready[raw_index.idx].smem_ptr, acquire=False) - # ---- V load ---------------------------------------------------------- + # ---- V load -------------------------------------------------------------- bars.mb_v_done[raw_index.idx].wait(raw_index.phase) if elect_one: bars.mb_v_ready[raw_index.idx].arrive(n_bytes=cfg.tma_v_bytes) @@ -359,7 +287,7 @@ def tmaldg_warp( tma_load_tile(sV_tma[raw_index.idx], v_slice, bars.mb_v_ready[raw_index.idx].smem_ptr, acquire=False) raw_index = advance(raw_index, cfg.smem_raw_stages) - tile_idx, sched_state = sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_publish_next(cfg, bars, sScheduler, mScheduler, scheduler_state, tile_idx, num_ctas, elect_one) @cute.jit @@ -370,8 +298,8 @@ def super_mma_warp( num_ctas, cu_seqlens, mWorkItems, - sSched, - lane, + sScheduler, + lane_idx, sK_inv_raw, sIntermediate_raw, sBeta_raw, @@ -381,29 +309,31 @@ def super_mma_warp( """Super-MMA warp role (warp 12): persistent scheduler loop computing the register-MMA Neumann-series T_inv.""" nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) + elect_one = nvvm.elect_sync() raw_index = PipelineState.start(phase=0) t_inv_free = PipelineState.start(phase=1) k_decay_ready = PipelineState.start(phase=0) - # ---- ldmatrix/stmatrix lane decode ------------------------------------------- - rhs_row_coord = lane % 8 + (cutlass.Int32(8) if (lane // 16) else cutlass.Int32(0)) - rhs_col_offset = cutlass.Int32(8) if ((lane // 8) % 2) else cutlass.Int32(0) - lhs_row_coord = lane % 8 + (cutlass.Int32(8) if ((lane // 8) % 2) else cutlass.Int32(0)) - lhs_col_offset = cutlass.Int32(8) if ((lane // 8) // 2) else cutlass.Int32(0) - decay_key_mask = cutlass.Int32(8) - stsm_row_coord = lane & 7 - stsm_col_coord = cutlass.Int32(0) - if (lane // 8) & 1: - stsm_row_coord = stsm_row_coord + cutlass.Int32(8) - if lane // 8 >= 2: - stsm_col_coord = cutlass.Int32(8) - stsm_idx = swizzle_lin_S(stsm_row_coord * cfg.b_t + (stsm_col_coord ^ (cfg.b_t // 2)), bbits=1, mbase=3, sshift=3) + # ---- ldmatrix/stmatrix lane decode ----------------------------------------------- + k_inv_row_coord = lane_idx % 8 + (cutlass.Int32(8) if (lane_idx // 16) else cutlass.Int32(0)) + k_inv_col_offset = cutlass.Int32(8) if ((lane_idx // 8) % 2) else cutlass.Int32(0) + k_decay_row_coord = lane_idx % 8 + (cutlass.Int32(8) if ((lane_idx // 8) % 2) else cutlass.Int32(0)) + k_decay_col_offset = cutlass.Int32(8) if ((lane_idx // 8) // 2) else cutlass.Int32(0) + t_inv_row_coord = lane_idx & 7 + t_inv_col_coord = cutlass.Int32(0) + if (lane_idx // 8) & 1: + t_inv_row_coord = t_inv_row_coord + cutlass.Int32(8) + if lane_idx // 8 >= 2: + t_inv_col_coord = cutlass.Int32(8) + t_inv_idx = t_inv_row_coord * cfg.b_t + swizzle_xor_32b(t_inv_row_coord, t_inv_col_coord) cum_chunk_base = cutlass.Int32(0) - sched_state = PipelineState.start(phase=0) + scheduler_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) - num_chunks_tile = wend - cstart # processed chunks; ring bookkeeping runs on cum_chunk_base + local_chunk_idx + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) + num_chunks_tile = write_end - compute_start for local_chunk_idx in cutlass.range(num_chunks_tile, unroll=1): cum_chunk = cum_chunk_base + local_chunk_idx decay_stage = k_decay_ready.idx @@ -416,54 +346,52 @@ def super_mma_warp( bars.mb_k_decay_inv_cg0_ready[decay_stage].wait(k_decay_ready.phase) k_decay_ready = advance(k_decay_ready, cfg.smem_decay_stages) - # ---- KK = K_decay @ K_inv^T ------------------------------------------ + # ---- KK = K decay @ K inv^T ---------------------------------------------- kk_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): kk_acc[accum_idx] = cutlass.Float32(0.0) - for k_block in cutlass.range_constexpr((cfg.d_k // 16)): - # Load B operand - k_inv_col = k_block * 16 + rhs_col_offset + for i in cutlass.range_constexpr((cfg.d_k // 16)): + k_inv_col = i * 16 + k_inv_col_offset k_inv_segment = k_inv_col // 64 - rhs_frag = nvvm.ldmatrix( + k_inv_frag = nvvm.ldmatrix( sK_inv_ptr + k_inv_segment * (cfg.b_t * 64) - + rhs_row_coord * 64 - + swizzle_xor_128b(rhs_row_coord, k_inv_col - k_inv_segment * 64, elem_bytes=2), + + k_inv_row_coord * 64 + + swizzle_xor_128b(k_inv_row_coord, k_inv_col - k_inv_segment * 64, elem_bytes=2), 4, nvvm.MMALayout.ROW, ) - # Load A operand - storage_key = (k_block * 16 + lhs_col_offset) ^ decay_key_mask - storage_slice = storage_key // 64 - kk_lhs_frag = nvvm.ldmatrix( + k_decay_col = i * 16 + k_decay_col_offset + k_decay_segment = k_decay_col // 64 + k_decay_frag = nvvm.ldmatrix( sK_decay_ptr - + storage_slice * (cfg.b_t * 64) - + swizzle_xor_128b(lhs_row_coord, lhs_row_coord * 64 + storage_key - storage_slice * 64, elem_bytes=2), + + k_decay_segment * (cfg.b_t * 64) + + swizzle_xor_128b(k_decay_row_coord, k_decay_row_coord * 64 + k_decay_col - k_decay_segment * 64, elem_bytes=2), 4, nvvm.MMALayout.ROW, ) mma_step( kk_acc, - (kk_lhs_frag[0], kk_lhs_frag[1], kk_lhs_frag[2], kk_lhs_frag[3]), - (rhs_frag[0], rhs_frag[1], rhs_frag[2], rhs_frag[3]), + (k_decay_frag[0], k_decay_frag[1], k_decay_frag[2], k_decay_frag[3]), + (k_inv_frag[0], k_inv_frag[1], k_inv_frag[2], k_inv_frag[3]), k_step=0, M=16, N=16, ab_dtype=cfg.io_dtype, ) - # ---- L = Beta * tril(KK, -1) fragment -------------------------------- + # ---- L = Beta * tril(KK, -1) --------------------------------------------- bars.mb_beta_ready[raw_index.idx].wait(raw_index.phase) - row_lo = lane // 4 + row_lo = lane_idx // 4 row_hi = row_lo + cutlass.Int32(8) beta_lo = (sBeta_ptr + row_lo).load().to(cutlass.Float32) beta_hi = (sBeta_ptr + row_hi).load().to(cutlass.Float32) l_regs = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): row_coord = row_hi if cutlass.const_expr(accum_idx % 4 >= 2) else row_lo - col_coord = (accum_idx // 4) * 8 + 2 * (lane % 4) + col_coord = (accum_idx // 4) * 8 + 2 * (lane_idx % 4) if cutlass.const_expr(accum_idx % 2 == 1): col_coord = col_coord + cutlass.Int32(1) l_regs[accum_idx] = kk_acc[accum_idx] if row_coord > col_coord else cutlass.Float32(0.0) @@ -477,13 +405,13 @@ def super_mma_warp( l_a3 = fp32_to_fp16(l_regs[6], l_regs[7], dtype=cfg.io_dtype) l_values = cutlass.Vector.from_elements((l_a0, l_a1, l_a2, l_a3), cutlass.Int32).bitcast(cfg.io_dtype).to(cutlass.Float32) - # ---- T_inv = I - L, then three Neumann doubling rounds --------------- + # ---- T^-1 = I - L, then three Neumann doubling rounds -------------------- tinv_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): row_coord = row_lo if cutlass.const_expr(accum_idx % 4 >= 2): row_coord = row_hi - col_coord = (accum_idx // 4) * 8 + 2 * (lane % 4) + col_coord = (accum_idx // 4) * 8 + 2 * (lane_idx % 4) if cutlass.const_expr(accum_idx % 2 == 1): col_coord = col_coord + cutlass.Int32(1) eye = cutlass.Float32(1.0) if row_coord == col_coord else cutlass.Float32(0.0) @@ -491,8 +419,8 @@ def super_mma_warp( lpow_a0, lpow_a1, lpow_a2, lpow_a3 = l_a0, l_a1, l_a2, l_a3 mov_lpow0, mov_lpow1, mov_lpow2, mov_lpow3 = movmatrix_16b(l_a0), movmatrix_16b(l_a1), movmatrix_16b(l_a2), movmatrix_16b(l_a3) - for _round in cutlass.range_constexpr(3): - # ---- Lpow = Lpow @ Lpow ------------------------------------------ + for neumann_round in cutlass.range_constexpr(3): + # ---- Lpow = Lpow @ Lpow ---------------------------------------------- sq_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): sq_acc[accum_idx] = cutlass.Float32(0.0) @@ -510,7 +438,7 @@ def super_mma_warp( lpow_a2 = fp32_to_fp16(sq_acc[4], sq_acc[5], dtype=cfg.io_dtype) lpow_a3 = fp32_to_fp16(sq_acc[6], sq_acc[7], dtype=cfg.io_dtype) mov_lpow0, mov_lpow1, mov_lpow2, mov_lpow3 = movmatrix_16b(lpow_a0), movmatrix_16b(lpow_a1), movmatrix_16b(lpow_a2), movmatrix_16b(lpow_a3) - # ---- T_inv += T_inv @ Lpow --------------------------------------- + # ---- T^-1 += T^-1 @ Lpow --------------------------------------------- upd_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): upd_acc[accum_idx] = cutlass.Float32(0.0) @@ -539,7 +467,7 @@ def super_mma_warp( bars.mb_t_inv_done[intermediate_stage].wait(t_inv_free.phase) t_inv_free = advance(t_inv_free, cfg.smem_intermediate_stages) nvvm.stmatrix( - sIntermediate_ptr + (cfg.b_t * cfg.b_t) + stsm_idx, + sIntermediate_ptr + (cfg.b_t * cfg.b_t) + t_inv_idx, [ fp32_to_fp16(tinv_acc[0], tinv_acc[1], dtype=cfg.io_dtype), fp32_to_fp16(tinv_acc[2], tinv_acc[3], dtype=cfg.io_dtype), @@ -554,7 +482,7 @@ def super_mma_warp( bars.mb_decay_super_done[decay_stage].arrive() raw_index = advance(raw_index, cfg.smem_raw_stages) cum_chunk_base += num_chunks_tile - tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one) @cute.jit @@ -565,11 +493,11 @@ def tcgen05_mma_warp( num_ctas, cu_seqlens, mWorkItems, - sSched, + sScheduler, sTmem_base, sIntermediate, sK_decay, - sK_restore, + sK_restore_trans, sState_scale_diag, bars, ) -> None: @@ -580,22 +508,22 @@ def tcgen05_mma_warp( nvvm.tcgen05_alloc(sTmem_base, cutlass.Int32(512), group=nvvm.CTAGroup.CTA_1) nvvm.barrier_cta_sync(cfg.tmem_lifecycle_barrier_id, thread_count=cfg.tmem_user_threads) tmem_base = sTmem_base.load() - state_inp_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_state_inp_offset, cutlass.Int8) + state_input_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_state_input_offset, cutlass.Int8) state_dsts = tuple(nvvm.make_tmem_ptr(tmem_base + cfg.tmem_state_acc_offset + k * 16, cutlass.Float32) for k in range(cfg.d_k // 16)) state_k_acc_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_state_k_acc_offset, cutlass.Float32) u_acc_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_u_acc_offset, cutlass.Float32) - y_inp_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_y_inp_offset, cutlass.Int8) - u_inp_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_u_inp_offset, cutlass.Int8) + y_input_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_y_input_offset, cutlass.Int8) + u_input_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_u_input_offset, cutlass.Int8) state_dst_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_state_acc_offset, cutlass.Float32) - state_inp_index = PipelineState.start(phase=0) + state_input_index = PipelineState.start(phase=0) state_read_index = PipelineState.start(phase=0) - y_inp_index = PipelineState.start(phase=0) - u_inp_index = PipelineState.start(phase=0) + y_input_index = PipelineState.start(phase=0) + u_input_index = PipelineState.start(phase=0) qk_scale_index = PipelineState.start(phase=0) k_decay_ready = PipelineState.start(phase=0) t_inv_ready = PipelineState.start(phase=0) - # ---- chunk-invariant GEMM descriptors ---------------------------------------- + # ---- chunk-invariant GEMM descriptors -------------------------------------------- bpe = cfg.io_dtype.width // 8 idesc_acc = nvvm.Tcgen05InstrDesc.build( c_dtype=cutlass.Float32, @@ -621,7 +549,7 @@ def tcgen05_mma_warp( m_dim=cfg.d_v, b_major=1, ) - bmm_state_desc = MmaDesc( + bmm_state_k_decay_desc = MmaDesc( M=cfg.d_v, N=cfg.b_t, K=cfg.d_k, @@ -633,7 +561,7 @@ def tcgen05_mma_warp( idesc=idesc_acc, kind=nvvm.Tcgen05MMAKind.F16, ) - bmm_diag_desc = MmaDesc( + bmm_state_diag_desc = MmaDesc( M=cfg.d_v, N=16, K=16, @@ -645,7 +573,7 @@ def tcgen05_mma_warp( idesc=idesc_diag, kind=nvvm.Tcgen05MMAKind.F16, ) - bmm_qk_desc = MmaDesc( + bmm_y_t_inv_desc = MmaDesc( M=cfg.d_v, N=cfg.b_t, K=cfg.b_t, @@ -657,7 +585,7 @@ def tcgen05_mma_warp( idesc=idesc_acc, kind=nvvm.Tcgen05MMAKind.F16, ) - bmm_final_state_desc = MmaDesc( + bmm_u_k_restore_desc = MmaDesc( M=cfg.d_v, N=cfg.d_k, K=cfg.b_t, @@ -669,38 +597,45 @@ def tcgen05_mma_warp( idesc=idesc_final_state, kind=nvvm.Tcgen05MMAKind.F16, ) - STATE_A_SEG = bmm_state_desc.sps_B * bmm_state_desc.tmem_advance_A - STATE_B_SEG = bmm_state_desc.smem_subtile_B >> 4 + STATE_A_SEG = bmm_state_k_decay_desc.sps_B * bmm_state_k_decay_desc.tmem_advance_A + STATE_B_SEG = bmm_state_k_decay_desc.smem_subtile_B >> 4 cum_chunk_base = cutlass.Int32(0) - sched_state = PipelineState.start(phase=0) + scheduler_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) - num_chunks_tile = wend - cstart + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) + num_chunks_tile = write_end - compute_start + if cutlass.const_expr(cfg.use_initial_state): + seed_state = compute_start == 0 for local_chunk_idx in cutlass.range(num_chunks_tile, unroll=1): cum_chunk = cum_chunk_base + local_chunk_idx - have_state = cutlass.Boolean(True) if cutlass.const_expr(cfg.use_initial_state) else local_chunk_idx > 0 + if cutlass.const_expr(cfg.use_initial_state): + have_state = local_chunk_idx > 0 or seed_state + else: + have_state = local_chunk_idx > 0 decay_stage = k_decay_ready.idx state_scale_diag_stage = qk_scale_index.idx intermediate_stage = t_inv_ready.idx sK_decay_stage = sK_decay[decay_stage] - sK_restore_stage = sK_restore[decay_stage] + sK_restore_stage = sK_restore_trans[decay_stage] sState_scale_diag_stage = sState_scale_diag[state_scale_diag_stage] sIntermediate_stage = sIntermediate[intermediate_stage] - # ---- State*K = State(T) @ K_decay^T ---------------------------------- + # ---- k state = state(T) @ K decay^T -------------------------------------- bars.mb_k_decay_inv_cg0_ready[decay_stage].wait(k_decay_ready.phase) k_decay_ready = advance(k_decay_ready, cfg.smem_decay_stages) if have_state: - bars.mb_state_inp_ready.wait(state_inp_index.phase) - state_inp_index = advance(state_inp_index, 1) + bars.mb_state_input_ready.wait(state_input_index.phase) + state_input_index = advance(state_input_index, 1) desc_k_decay = sK_decay_stage.desc() - for s in cutlass.range_constexpr(bmm_state_desc.num_subtiles_B): - for k in cutlass.range_constexpr(bmm_state_desc.sps_B): + for s in cutlass.range_constexpr(bmm_state_k_decay_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_state_k_decay_desc.sps_B): mma_ts_step( - bmm_state_desc, - state_inp_ptr.subview(s * STATE_A_SEG), + bmm_state_k_decay_desc, + state_input_ptr.subview(s * STATE_A_SEG), desc_k_decay + s * STATE_B_SEG, state_k_acc_ptr, k, @@ -719,15 +654,15 @@ def tcgen05_mma_warp( bars.mb_state_acc_read_done.wait(state_read_index.phase) state_read_index = advance(state_read_index, 1) - # ---- State decay = State(T) @ exp2(g_last) diag (per-k-atom blocks) ---- + # ---- state decay = state(T) @ diag(exp2(g last)) (per-k-atom blocks) ----- if have_state: desc_diag = sState_scale_diag_stage.desc() - for k_block in cutlass.range_constexpr(cfg.d_k // 16): + for i in cutlass.range_constexpr(cfg.d_k // 16): mma_ts_step( - bmm_diag_desc, - state_inp_ptr.subview(k_block * bmm_diag_desc.tmem_advance_A), - desc_diag.advance_start_address(k_block * 256 * 2), - state_dsts[k_block], + bmm_state_diag_desc, + state_input_ptr.subview(i * bmm_state_diag_desc.tmem_advance_A), + desc_diag.advance_start_address(i * 256 * 2), + state_dsts[i], 0, cutlass.Boolean(False), ) @@ -735,22 +670,22 @@ def tcgen05_mma_warp( if elect_one: bars.mb_state_scale_diag_done[state_scale_diag_stage].arrive(cta_group=1) - # ---- U = Y(T) @ T_inv ------------------------------------------------ + # ---- U = Y(T) @ T^-1 ----------------------------------------------------- bars.mb_t_inv_ready[intermediate_stage].wait(t_inv_ready.phase) - bars.mb_y_inp_ready.wait(y_inp_index.phase) - y_inp_index = advance(y_inp_index, 1) + bars.mb_y_input_ready.wait(y_input_index.phase) + y_input_index = advance(y_input_index, 1) desc_t_inv = sIntermediate_stage.shifted((cfg.b_t * cfg.b_t)).desc() - mma_ts_step(bmm_qk_desc, y_inp_ptr, desc_t_inv, u_acc_ptr, 0, cutlass.Boolean(False)) + mma_ts_step(bmm_y_t_inv_desc, y_input_ptr, desc_t_inv, u_acc_ptr, 0, cutlass.Boolean(False)) if elect_one: bars.mb_t_inv_done[intermediate_stage].arrive(cta_group=1) bars.mb_u_acc_ready.arrive(cta_group=1) - # ---- final_state += U(T) @ K_restore --------------------------------- - bars.mb_u_inp_ready.wait(u_inp_index.phase) - u_inp_index = advance(u_inp_index, 1) + # ---- final state += U(T) @ K restore ------------------------------------- + bars.mb_u_input_ready.wait(u_input_index.phase) + u_input_index = advance(u_input_index, 1) desc_k_restore = sK_restore_stage.desc() - mma_ts_step(bmm_final_state_desc, u_inp_ptr, desc_k_restore, state_dst_ptr, 0, have_state) + mma_ts_step(bmm_u_k_restore_desc, u_input_ptr, desc_k_restore, state_dst_ptr, 0, have_state) if elect_one: bars.mb_k_restore_done[decay_stage].arrive(cta_group=1) bars.mb_state_acc_done.arrive(cta_group=1) @@ -759,7 +694,7 @@ def tcgen05_mma_warp( qk_scale_index = advance(qk_scale_index, cfg.smem_state_scale_diag_stages) cum_chunk_base += num_chunks_tile - tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one) bars.mb_tmem_done[0].wait(0) nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_1) nvvm.tcgen05_dealloc( @@ -777,7 +712,7 @@ def epilogue_warp( num_ctas, cu_seqlens, mWorkItems, - sSched, + sScheduler, sCheckpoint_raw, desc_checkpoint_base, checkpoint_every_n_tokens, @@ -800,21 +735,23 @@ def epilogue_warp( tma_subtile_stride_elems=cfg.d_k * 64, ) checkpoint_ready_index = PipelineState.start(phase=0) - sched_state = PipelineState.start(phase=0) + scheduler_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) if cutlass.const_expr(cfg.enable_checkpoints): head_o = head_idx checkpoint_slot = batch_idx * cutlass.Int32(TENSOR_MAP_QWORDS) desc_checkpoint_slot = (desc_checkpoint_base + checkpoint_slot).tospace(cutlass.AddressSpace.generic) checkpoint_chunks = checkpoint_every_n_tokens // cutlass.Int32(cfg.b_t) - checkpoint_quot = (cstart + cutlass.Int32(1)) // checkpoint_chunks - checkpoint_mod = (cstart + cutlass.Int32(1)) % checkpoint_chunks + checkpoint_quot = (compute_start + cutlass.Int32(1)) // checkpoint_chunks + checkpoint_mod = (compute_start + cutlass.Int32(1)) % checkpoint_chunks if elect_one: tma_tensormap_acquire(desc_checkpoint_slot) - num_chunks_tile = wend - cstart - if num_chunks_tile > 0 and wstart == 0: + num_chunks_tile = write_end - compute_start + if num_chunks_tile > 0 and write_start == 0: checkpoint_stage = checkpoint_ready_index.idx bars.mb_checkpoint_tmastg_ready[checkpoint_stage].wait(checkpoint_ready_index.phase) checkpoint_ready_index = advance(checkpoint_ready_index, cfg.smem_checkpoint_stages) @@ -824,11 +761,11 @@ def epilogue_warp( tma_store_wait(0) bars.mb_checkpoint_tmastg_done[checkpoint_stage].arrive() for local_chunk_idx in cutlass.range(num_chunks_tile, unroll=1): - chunk_idx = cstart + local_chunk_idx + chunk_idx = compute_start + local_chunk_idx if local_chunk_idx > 0: - # ---- checkpoint store ---------------------------------------- + # ---- checkpoint store -------------------------------------------- do_checkpoint = checkpoint_mod == 0 - do_checkpoint = do_checkpoint and chunk_idx >= wstart + do_checkpoint = do_checkpoint and chunk_idx >= write_start if do_checkpoint: checkpoint_stage = checkpoint_ready_index.idx bars.mb_checkpoint_tmastg_ready[checkpoint_stage].wait(checkpoint_ready_index.phase) @@ -843,7 +780,7 @@ def epilogue_warp( if checkpoint_mod == checkpoint_chunks: checkpoint_mod = cutlass.Int32(0) checkpoint_quot = checkpoint_quot + cutlass.Int32(1) - tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one) @cute.jit @@ -851,9 +788,7 @@ def gate_scale(cfg, raw_gate: cutlass.Float32) -> cutlass.Float32: """Map raw gate to the log2-domain decay increment used by KDA.""" if cutlass.const_expr(cfg.safe_gate): - half = cutlass.Float32(0.5) - sigmoid = cute.math.tanh(raw_gate * half, approx=True) * half + half - return cfg.gate_scale_log2 * sigmoid + return cfg.gate_scale_log2 * sigmoid(raw_gate) # Default ABI: Gate arrives in natural-log space return raw_gate * cutlass.Float32(LOG2_E) @@ -866,8 +801,8 @@ def compute0_warp_group( num_ctas, cu_seqlens, mWorkItems, - sSched, - lane, + sScheduler, + lane_idx, warp_idx, mA_log, mDt_bias, @@ -884,31 +819,36 @@ def compute0_warp_group( """CG0 warp-group role (warps 0-7): persistent scheduler loop running the Gate prefix scan and staging the decay/restore operands.""" nvvm.setmaxregister(cfg.num_regs_compute_group_0, nvvm.SetMaxRegisterAction.INCREASE) + elect_one = nvvm.elect_sync() + + scheduler_state = PipelineState.start(phase=0) + cg0_warp = warp_idx - cfg.compute_group_0_warp_ids[0] - cg0_group_id = cg0_warp // cfg.cg0_warps_per_group cg0_local_warp = cg0_warp % cfg.cg0_warps_per_group - prefix_dim = cg0_local_warp * cfg.threads_per_warp + lane + + cg0_group_id = cg0_warp // cfg.cg0_warps_per_group + channel_dim = cg0_local_warp * cfg.threads_per_warp + lane_idx cg0_a_log_exp = cutlass.Float32(1.0) cg0_dt_bias_value = cutlass.Float32(0.0) cum_chunk_base = cutlass.Int32(0) - sched_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) opaque_one = opaque_f32_zero() + cutlass.Float32(1.0) while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) head_o = head_idx - num_chunks_tile = wend - cstart + num_chunks_tile = write_end - compute_start if cutlass.const_expr(cfg.safe_gate): if num_chunks_tile > 0: cg0_a_log_exp = cute.math.exp2(mA_log[head_o].to(cutlass.Float32) * LOG2_E, fastmath=True) - cg0_dt_bias_value = mDt_bias[head_o, prefix_dim].to(cutlass.Float32) - # tile entry: both ping-pong groups inherit each other's delivery proofs (parity-swap guard) + cg0_dt_bias_value = mDt_bias[head_o, channel_dim].to(cutlass.Float32) nvvm.barrier_cta_sync(cfg.cg0_tile_entry_barrier_id, thread_count=cfg.cg0_group_count * cfg.cg0_threads_per_group) group_cum_chunk_start = cum_chunk_base + cutlass.Int32(cg0_group_id) diag_ring_stage = group_cum_chunk_start % cutlass.Int32(cfg.smem_state_scale_diag_stages) diag_ring_phase = (group_cum_chunk_start // cutlass.Int32(cfg.smem_state_scale_diag_stages)) % cutlass.Int32(2) for local_chunk_idx in cutlass.range(cg0_group_id, num_chunks_tile, cfg.cg0_group_count, unroll=1): - chunk_idx = cstart + local_chunk_idx + chunk_idx = compute_start + local_chunk_idx cum_chunk = cum_chunk_base + local_chunk_idx chunk_start = chunk_idx * cfg.b_t decay_stage = cum_chunk % cfg.smem_decay_stages @@ -922,63 +862,53 @@ def compute0_warp_group( sK_restore_ptr = sK_restore_raw.data_ptr() + decay_stage * (cfg.d_k * cfg.b_t) sState_scale_diag_ptr = sState_scale_diag_raw.data_ptr() + state_scale_diag_stage * ((cfg.d_k // 16) * 256) - # ---- Beta scalars --------------------------------------------------- + # ---- Beta scalars -------------------------------------------------------- if cg0_local_warp == 0: bars.mb_beta_done[raw_stage].wait(((cum_chunk // cfg.smem_raw_stages) + 1) % 2) - if lane < cfg.b_t: - token_idx = chunk_idx * cfg.b_t + lane + if lane_idx < cfg.b_t: + token_idx = chunk_idx * cfg.b_t + lane_idx beta_value = cutlass.Float32(0.0) - if token_idx < seqlen_b: + if token_idx < batch_seqlen: beta_value = mBeta[batch_start + token_idx, head_o].to(cutlass.Float32) if cutlass.const_expr(cfg.beta_sigmoid): - half = cutlass.Float32(0.5) - beta_value = (cute.math.tanh(beta_value * half, approx=True) * half + half).to(mBeta.element_type).to(cutlass.Float32) - sBeta_raw[raw_stage * cfg.b_t + lane] = beta_value + beta_value = sigmoid(beta_value).to(mBeta.element_type).to(cutlass.Float32) + sBeta_raw[raw_stage * cfg.b_t + lane_idx] = beta_value bars.mb_beta_ready[raw_stage].arrive() bars.mb_gate_ready[raw_stage].wait((cum_chunk // cfg.smem_raw_stages) % 2) row_group_start = cg0_local_warp * (cfg.b_t // cfg.cg0_warps_per_group) - lane_row_group = lane // 8 - lane_in_row_group = lane - lane_row_group * 8 + lane_row_group = lane_idx // 8 + lane_in_row_group = lane_idx - lane_row_group * 8 decay_row = row_group_start + lane_row_group - decay_key_mask = cutlass.Int32(8) - prefix_dim = cg0_local_warp * cfg.threads_per_warp + lane + channel_dim = cg0_local_warp * cfg.threads_per_warp + lane_idx - # ---- Gate prefix scan ----------------------------------------------- - f32_segment = prefix_dim // 32 + # ---- Gate prefix scan ---------------------------------------------------- + f32_segment = channel_dim // 32 prefix_seg_base = f32_segment * (cfg.b_t * 32) - prefix_col = prefix_dim - f32_segment * 32 + prefix_col = channel_dim - f32_segment * 32 gate_raw = cutlass.Array(cutlass.Float32, cfg.b_t, alignment=16) for row in cutlass.range_constexpr(cfg.b_t): prefix_idx = prefix_seg_base + swizzle_xor_128b(row, row * 32 + prefix_col, elem_bytes=4) gate_raw[row] = (sGate_ptr + prefix_idx).load() g_prefix_regs = cutlass.Array(cutlass.Float32, cfg.b_t, alignment=16) if cutlass.const_expr(cfg.safe_gate): - valid_rows = seqlen_b - chunk_idx * cutlass.Int32(cfg.b_t) - valid_mask = cutlass.vector.create_mask([cfg.b_t], [valid_rows]) - for row_pair in cutlass.range_constexpr(cfg.b_t // 2): - row0 = row_pair * 2 - row1 = row0 + 1 - gate0 = cg0_a_log_exp * (gate_raw[row0] + cg0_dt_bias_value) - gate1 = cg0_a_log_exp * (gate_raw[row1] + cg0_dt_bias_value) - gate0 = gate_scale( - cfg, - gate0, - ) - gate1 = gate_scale( - cfg, - gate1, - ) - gate_pair = cutlass.Vector.from_elements((gate0, gate1), cutlass.Float32) - gate_pair = cutlass.vector.where(valid_mask[row0 : row1 + 1], gate_pair, 0.0) - g_prefix_regs[row0] = gate_pair[0] - g_prefix_regs[row1] = gate_pair[1] + for row in cutlass.range_constexpr(cfg.b_t): + gate = gate_raw[row] + token_idx = chunk_idx * cutlass.Int32(cfg.b_t) + cutlass.Int32(row) + if token_idx < batch_seqlen: + gate = gate_scale( + cfg, + cg0_a_log_exp * (gate + cg0_dt_bias_value), + ) + else: + gate = cutlass.Float32(0.0) + g_prefix_regs[row] = gate else: for row in cutlass.range_constexpr(cfg.b_t): gate = gate_raw[row] token_idx = chunk_idx * cutlass.Int32(cfg.b_t) + cutlass.Int32(row) - if token_idx < seqlen_b: + if token_idx < batch_seqlen: gate = gate_scale( cfg, gate, @@ -999,7 +929,7 @@ def compute0_warp_group( g_prefix_regs[row1] = prefix1 prefix_acc = prefix1 - # ---- exp2(g): stage prefixes + final-token decay --------------------- + # ---- exp2(g): stage prefixes + final-token decay ------------------------- for row in cutlass.range_constexpr(cfg.b_t): g_prefix_regs[row] = cute.math.exp2(g_prefix_regs[row], fastmath=True) @@ -1008,13 +938,11 @@ def compute0_warp_group( prefix_idx = prefix_seg_base + swizzle_xor_128b(row, row * 32 + prefix_col, elem_bytes=4) (sGate_ptr + prefix_idx).store(g_prefix_regs[row]) - # ---- state-scale diag: stage exp2(g_last) decay blocks --------------- + # ---- state-scale diag: stage exp2(g last) decay blocks ------------------- bars.mb_state_scale_diag_done[state_scale_diag_stage].wait(diag_ring_phase ^ cutlass.Int32(1)) - block = prefix_dim // cutlass.Int32(16) - coord = prefix_dim - block * cutlass.Int32(16) - storage_col = coord ^ cutlass.Int32((cfg.b_t // 2)) - linear_idx = block * cutlass.Int32(256) + coord * cutlass.Int32(16) + storage_col - diag_idx = swizzle_lin_S(linear_idx, bbits=1, mbase=3, sshift=3) + block = channel_dim // cutlass.Int32(16) + coord = channel_dim - block * cutlass.Int32(16) + diag_idx = block * cutlass.Int32(256) + coord * cutlass.Int32(16) + swizzle_xor_32b(channel_dim, coord) sState_scale_diag_ptr[diag_idx] = exp_g_last.to(cfg.io_dtype) nvvm.barrier_cta_sync(cfg.cg0_group_sync_barrier_base_id + cg0_group_id, thread_count=cfg.cg0_threads_per_group) @@ -1023,7 +951,7 @@ def compute0_warp_group( k_inv_pack = cutlass.Array(cutlass.Int32, 2 * 4, alignment=16) raw_k_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) - # ---- optional K L2-norm + K_inv staging ------------------------------ + # ---- optional K L2-norm + K inv stage ------------------------------------ if cutlass.const_expr(cfg.l2norm): kk_lo = opaque_f32_zero() kk_hi = opaque_f32_zero() @@ -1053,7 +981,7 @@ def compute0_warp_group( norm_floor_sq = cutlass.Float32(L2_NORM_EPS * L2_NORM_EPS) k_inv_norm = cute.math.rsqrt(cute.math.max(k_sum_sq, norm_floor_sq), fastmath=True) - # ---- decay/restore operands: exp2(+-g) applied per key channel ------- + # ---- decay/restore operands: exp2(+-g) applied per key channel ----------- exp_g_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) exp_g_last_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) for dim_half in cutlass.range_constexpr(2): @@ -1076,7 +1004,7 @@ def compute0_warp_group( dim_base = dim_half * (cfg.d_k // 2) + lane_in_row_group * 8 reg_base = dim_half * 8 - # ---- K decay + K_inv operands: K * exp2(+g) and K * exp2(-g) ----- + # ---- K decay + K inv operands: K * exp2(+g) and K * exp2(-g) --------- k_decay_pack = cutlass.Array(cutlass.Int32, 4, alignment=16) for pair_idx in cutlass.range_constexpr(4): dim0 = pair_idx * 2 @@ -1118,16 +1046,14 @@ def compute0_warp_group( f16_segment_dim = dim_base - f16_segment * 64 k_inv_swizzled_idx = f16_segment * (cfg.b_t * 64) + decay_row * 64 + swizzle_xor_128b(decay_row, f16_segment_dim, elem_bytes=2) (sK_inv_ptr + k_inv_swizzled_idx).store(k_inv_vec, alignment=16) - storage_key = dim_base ^ decay_key_mask - storage_slice = storage_key // 64 - decay_swizzled_idx = storage_slice * (cfg.b_t * 64) + swizzle_xor_128b( - decay_row, decay_row * 64 + storage_key - storage_slice * 64, elem_bytes=2 - ) + decay_col = dim_base + decay_segment = decay_col // 64 + decay_swizzled_idx = decay_segment * (cfg.b_t * 64) + swizzle_xor_128b(decay_row, decay_row * 64 + decay_col - decay_segment * 64, elem_bytes=2) (sK_decay_ptr + decay_swizzled_idx).store(k_decay_vec, alignment=16) nvvm.fence_proxy("async.shared", space="cta") bars.mb_k_decay_inv_cg0_ready[decay_stage].arrive() - # ---- K_restore operand: K_inv * exp_g_last -------------------------- + # ---- K restore operand: K inv * exp2(g last) ----------------------------- bars.mb_k_restore_done[decay_stage].wait(((cum_chunk // cfg.smem_decay_stages + 1) % 2)) for dim_half in cutlass.range_constexpr(2): dim_base = dim_half * (cfg.d_k // 2) + lane_in_row_group * 8 @@ -1138,10 +1064,9 @@ def compute0_warp_group( dim1 = dim0 + 1 exp_g_last_pair = fp32_to_fp16(exp_g_last_regs[reg_base + dim0], exp_g_last_regs[reg_base + dim1], dtype=cfg.io_dtype) k_restore_pack[pair_idx] = mul_f16x2(k_inv_pack[dim_half * 4 + pair_idx], exp_g_last_pair, cfg.io_dtype) - storage_row = decay_row ^ (cfg.b_t // 2) f16_segment = dim_base // 64 f16_segment_dim = dim_base - f16_segment * 64 - k_restore_idx = f16_segment * (cfg.b_t * 64) + storage_row * 64 + swizzle_xor_128b(storage_row, f16_segment_dim, elem_bytes=2) + k_restore_idx = f16_segment * (cfg.b_t * 64) + decay_row * 64 + swizzle_xor_128b(decay_row, f16_segment_dim, elem_bytes=2) k_restore_vec = cutlass.Vector.from_elements( ( k_restore_pack[0], @@ -1161,7 +1086,7 @@ def compute0_warp_group( diag_ring_stage = diag_ring_stage - cutlass.Int32(cfg.smem_state_scale_diag_stages) if wrapped else diag_ring_stage diag_ring_phase = diag_ring_phase ^ (cutlass.Int32(1) if wrapped else cutlass.Int32(0)) cum_chunk_base += num_chunks_tile - tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one) @cute.jit @@ -1172,8 +1097,8 @@ def compute1_warp_group( num_ctas, cu_seqlens, mWorkItems, - sSched, - lane, + sScheduler, + lane_idx, sTmem_base, warp_idx, mState_out, @@ -1187,141 +1112,123 @@ def compute1_warp_group( """CG1 warp-group role (warps 8-11): persistent scheduler loop staging the value-side TMEM operands and storing the checkpoint/final states.""" nvvm.setmaxregister(cfg.num_regs_compute_group_1, nvvm.SetMaxRegisterAction.INCREASE) - sCheckpoint_ptr = sCheckpoint_raw.data_ptr() if cutlass.const_expr(cfg.enable_checkpoints) else sV_raw.data_ptr() + elect_one = nvvm.elect_sync() + checkpoint_done_index = PipelineState.start(phase=1) + + sCheckpoint_ptr = sCheckpoint_raw.data_ptr() if cutlass.const_expr(cfg.enable_checkpoints) else sV_raw.data_ptr() nvvm.barrier_cta_sync(cfg.tmem_lifecycle_barrier_id, thread_count=cfg.tmem_user_threads) tmem_base = sTmem_base.load() tmem_col = tmem_base & 0xFFFF tmem_row = tmem_base >> 16 tmem_subpartition = warp_idx % (cfg.d_v // cfg.threads_per_warp) # ldmatrix.x4 COL lane decode for the V loads - frag_row_coord = (lane // 16) * 8 + (lane & 7) - frag_col_offset = ((lane // 8) & 1) * 8 - row_id = tmem_row + tmem_subpartition * cfg.threads_per_warp - value_dim = tmem_subpartition * cfg.threads_per_warp + lane + ov_row_coord = (lane_idx // 16) * 8 + (lane_idx & 7) + ov_col_offset = ((lane_idx // 8) & 1) * 8 + value_dim = tmem_subpartition * cfg.threads_per_warp + lane_idx value_dim_base = tmem_subpartition * cfg.threads_per_warp - row_addr = row_id << 16 - row16_addr = (row_id + 16) << 16 - st_row_addr = tmem_row << 16 - st_row16_addr = (tmem_row + 16) << 16 + row_lo_addr = tmem_row << 16 + row_hi_addr = (tmem_row + 16) << 16 state_col_id = tmem_col + cfg.tmem_state_acc_offset - packed_col_id = tmem_col + cfg.tmem_state_inp_offset + packed_col_id = tmem_col + cfg.tmem_state_input_offset statek_col_id = tmem_col + cfg.tmem_state_k_acc_offset - y_inp_col_id = tmem_col + cfg.tmem_y_inp_offset - u_acc_addr = row_addr + tmem_col + cfg.tmem_u_acc_offset - u_inp_addr = st_row_addr + tmem_col + cfg.tmem_u_inp_offset - v_swz_off0 = ( - (value_dim_base + frag_col_offset) // 64 * (cfg.b_t * 64) - + frag_row_coord * 64 - + swizzle_xor_128b(frag_row_coord, (value_dim_base + frag_col_offset) % 64, elem_bytes=2) + y_input_col_id = tmem_col + cfg.tmem_y_input_offset + u_acc_addr = row_lo_addr + tmem_col + cfg.tmem_u_acc_offset + u_input_addr = row_lo_addr + tmem_col + cfg.tmem_u_input_offset + v_swz_off_lo = ( + (value_dim_base + ov_col_offset) // 64 * (cfg.b_t * 64) + + ov_row_coord * 64 + + swizzle_xor_128b(ov_row_coord, (value_dim_base + ov_col_offset) % 64, elem_bytes=2) ) - v_swz_off = ( - (value_dim_base + 16 + frag_col_offset) // 64 * (cfg.b_t * 64) - + frag_row_coord * 64 - + swizzle_xor_128b(frag_row_coord, (value_dim_base + 16 + frag_col_offset) % 64, elem_bytes=2) + v_swz_off_hi = ( + (value_dim_base + 16 + ov_col_offset) // 64 * (cfg.b_t * 64) + + ov_row_coord * 64 + + swizzle_xor_128b(ov_row_coord, (value_dim_base + 16 + ov_col_offset) % 64, elem_bytes=2) ) - checkpoint_swz_off0 = (value_dim_base + frag_col_offset) // 64 * (cfg.d_k * 64) - checkpoint_swz_col0 = (value_dim_base + frag_col_offset) % 64 - checkpoint_swz_off = (value_dim_base + 16 + frag_col_offset) // 64 * (cfg.d_k * 64) - checkpoint_swz_col = (value_dim_base + 16 + frag_col_offset) % 64 + checkpoint_swz_off_lo = (value_dim_base + ov_col_offset) // 64 * (cfg.d_k * 64) + checkpoint_swz_col_lo = (value_dim_base + ov_col_offset) % 64 + checkpoint_swz_off_hi = (value_dim_base + 16 + ov_col_offset) // 64 * (cfg.d_k * 64) + checkpoint_swz_col_hi = (value_dim_base + 16 + ov_col_offset) % 64 state_k_acc_index = PipelineState.start(phase=0) u_acc_index = PipelineState.start(phase=0) state_upd_index = PipelineState.start(phase=0) raw_index = PipelineState.start(phase=0) - sched_state = PipelineState.start(phase=0) + scheduler_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) while tile_idx < total_tiles: - batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) + batch_idx, head_idx, batch_start, batch_end, batch_seqlen, batch_num_chunks, write_start, write_end, compute_start, compute_end = decode_work_item( + cfg, tile_idx, mWorkItems + ) head_o = head_idx - num_chunks_tile = wend - cstart + num_chunks_tile = write_end - compute_start if num_chunks_tile > 0: - # ---- first chunk: seed state TMEM from mState_init ---------- - seed_from_initial_state = cstart == 0 + # ---- first chunk: seed state TMEM from mState init ----------------------- + seed_from_initial_state = compute_start == 0 + sV_ptr = sV_raw.data_ptr() + raw_index.idx * (cfg.d_v * cfg.b_t) + sBeta_ptr = sBeta_raw.data_ptr() + raw_index.idx * cfg.b_t + + # ---- state seed: initial state GMEM -> state TMEM ------------------------ if cutlass.const_expr(mState_init is not None): if seed_from_initial_state: seed_vw = 16 // (mState_init.element_type.width // 8) seed_src = (mState_init.iterator + mState_init.layout((batch_idx, head_o, value_dim, 0))).raw_ptr() - for key_block_start in cutlass.range_constexpr(0, cfg.d_k, 32): - state_block = cutlass.Array(cutlass.Float32, 32, alignment=16) - for g in cutlass.range_constexpr(32 // seed_vw): - seed_chunk = (seed_src + key_block_start + g * seed_vw).load(count=seed_vw, alignment=16) + state_vecs = [] + for i in cutlass.range_constexpr(cfg.d_k // 16): + state_block = [] + for g in cutlass.range_constexpr(16 // seed_vw): + seed_chunk = (seed_src + i * 16 + g * seed_vw).load(count=seed_vw, alignment=16) for t in cutlass.range_constexpr(seed_vw): - state_block[g * seed_vw + t] = seed_chunk[t].to(cutlass.Float32) - - nvvm.tcgen05_st( - "32x32b", - nvvm.make_tmem_ptr((row_id << 16) + (tmem_col + cfg.tmem_state_acc_offset + key_block_start), cutlass.Float32), - state_block[0:32], - ) - else: - for key_block_start in cutlass.range_constexpr(0, cfg.d_k, 32): - state_block = cutlass.Array(cutlass.Float32, 32, alignment=16) - for col in cutlass.range_constexpr(32): - state_block[col] = cutlass.Float32(0.0) + state_block.append(seed_chunk[t].to(cutlass.Float32)) + state_vecs.append(state_block) + for i in cutlass.range_constexpr(cfg.d_k // 16): + state_pack = cutlass.Array(cutlass.Int32, 8, alignment=16) + for packed_col in cutlass.range_constexpr(8): + state_pack[packed_col] = fp32_to_fp16(state_vecs[i][2 * packed_col], state_vecs[i][2 * packed_col + 1], dtype=cfg.io_dtype) nvvm.tcgen05_st( "32x32b", - nvvm.make_tmem_ptr((row_id << 16) + (tmem_col + cfg.tmem_state_acc_offset + key_block_start), cutlass.Float32), - state_block[0:32], + nvvm.make_tmem_ptr(row_lo_addr + packed_col_id + i * 8, cutlass.Int8), + state_pack[0:8], ) - if cutlass.const_expr(mState_init is not None): - nvvm.tcgen05_wait("store") - sV_ptr = sV_raw.data_ptr() + raw_index.idx * (cfg.d_v * cfg.b_t) - sBeta_ptr = sBeta_raw.data_ptr() + raw_index.idx * cfg.b_t - # ---- state repack: acc TMEM -> packed b16 TMEM ---------------------- - if cutlass.const_expr(mState_init is not None): - state_vecs = [] - for sub in cutlass.range_constexpr(cfg.d_k // 16): - state_vecs.append(nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + state_col_id + sub * 16, cutlass.Float32), num=16)) - - for sub in cutlass.range_constexpr(cfg.d_k // 16): - state_pack = cutlass.Array(cutlass.Int32, 8, alignment=16) - for packed_col in cutlass.range_constexpr(8): - source_pair = packed_col ^ 4 - state_pack[packed_col] = fp32_to_fp16(state_vecs[sub][2 * source_pair], state_vecs[sub][2 * source_pair + 1], dtype=cfg.io_dtype) - nvvm.tcgen05_st( - "32x32b", - nvvm.make_tmem_ptr((tmem_row << 16) + packed_col_id + sub * 8, cutlass.Int8), - state_pack[0:8], - ) - - nvvm.tcgen05_wait("store") - bars.mb_state_inp_ready.arrive() - if cutlass.const_expr(cfg.enable_checkpoints): - if wstart == 0: - checkpoint_stage = checkpoint_done_index.idx - bars.mb_checkpoint_tmastg_done[checkpoint_stage].wait(checkpoint_done_index.phase) - checkpoint_done_index = advance(checkpoint_done_index, cfg.smem_checkpoint_stages) - checkpoint_stage_base = checkpoint_stage * (cfg.d_k * cfg.d_v) - for sub in cutlass.range_constexpr(cfg.d_k // 16): - for g in cutlass.range_constexpr(2): - packs = tuple( - fp32_to_fp16(state_vecs[sub][g * 8 + 2 * t], state_vecs[sub][g * 8 + 2 * t + 1], dtype=cfg.io_dtype) for t in range(4) - ) - dk = sub * 16 + g * 8 - checkpoint_addr = ( - checkpoint_stage_base + (dk // 64) * (cfg.d_v * 64) + value_dim * 64 + swizzle_xor_128b(value_dim, dk % 64, elem_bytes=2) - ) - (sCheckpoint_ptr + checkpoint_addr).store( - cutlass.Vector.from_elements(packs, cutlass.Int32).bitcast(cfg.io_dtype), alignment=16 - ) - nvvm.tcgen05_wait("load") - nvvm.fence_proxy("async.shared", space="cta") - bars.mb_checkpoint_tmastg_ready[checkpoint_stage].arrive() - bars.mb_state_acc_read_done.arrive() + nvvm.tcgen05_wait("store") + bars.mb_state_input_ready.arrive() + if cutlass.const_expr(cfg.enable_checkpoints): + if write_start == 0: + checkpoint_stage = checkpoint_done_index.idx + bars.mb_checkpoint_tmastg_done[checkpoint_stage].wait(checkpoint_done_index.phase) + checkpoint_done_index = advance(checkpoint_done_index, cfg.smem_checkpoint_stages) + checkpoint_stage_base = checkpoint_stage * (cfg.d_k * cfg.d_v) + for i in cutlass.range_constexpr(cfg.d_k // 16): + for g in cutlass.range_constexpr(2): + packs = tuple( + fp32_to_fp16(state_vecs[i][g * 8 + 2 * t], state_vecs[i][g * 8 + 2 * t + 1], dtype=cfg.io_dtype) for t in range(4) + ) + dk = i * 16 + g * 8 + checkpoint_addr = ( + checkpoint_stage_base + + (dk // 64) * (cfg.d_v * 64) + + value_dim * 64 + + swizzle_xor_128b(value_dim, dk % 64, elem_bytes=2) + ) + (sCheckpoint_ptr + checkpoint_addr).store( + cutlass.Vector.from_elements(packs, cutlass.Int32).bitcast(cfg.io_dtype), alignment=16 + ) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_checkpoint_tmastg_ready[checkpoint_stage].arrive() + bars.mb_state_acc_read_done.arrive() if cutlass.const_expr(cfg.enable_checkpoints and mState_init is None): - if wstart == 0: + if write_start == 0: checkpoint_stage = checkpoint_done_index.idx bars.mb_checkpoint_tmastg_done[checkpoint_stage].wait(checkpoint_done_index.phase) checkpoint_done_index = advance(checkpoint_done_index, cfg.smem_checkpoint_stages) checkpoint_stage_base = checkpoint_stage * (cfg.d_k * cfg.d_v) zero_packs = tuple(cutlass.Int32(0) for _ in range(4)) - for sub in cutlass.range_constexpr(cfg.d_k // 16): + for i in cutlass.range_constexpr(cfg.d_k // 16): for g in cutlass.range_constexpr(2): - dk = sub * 16 + g * 8 + dk = i * 16 + g * 8 checkpoint_addr = ( checkpoint_stage_base + (dk // 64) * (cfg.d_v * 64) + value_dim * 64 + swizzle_xor_128b(value_dim, dk % 64, elem_bytes=2) ) @@ -1331,81 +1238,62 @@ def compute1_warp_group( nvvm.fence_proxy("async.shared", space="cta") bars.mb_checkpoint_tmastg_ready[checkpoint_stage].arrive() - # ---- Y staging: Y = Beta * (V - State*K) ----------------------------- + # ---- Y stage: Y = Beta * (V - k state) ----------------------------------- bars.mb_v_ready[raw_index.idx].wait(raw_index.phase) - raw_v_frag0 = nvvm.ldmatrix( - sV_ptr + v_swz_off0, + raw_v_frag_lo = nvvm.ldmatrix( + sV_ptr + v_swz_off_lo, 4, nvvm.MMALayout.COL, ) - raw_v_frag1 = nvvm.ldmatrix( - sV_ptr + v_swz_off, + raw_v_frag_hi = nvvm.ldmatrix( + sV_ptr + v_swz_off_hi, 4, nvvm.MMALayout.COL, ) bars.mb_beta_ready[raw_index.idx].wait(raw_index.phase) - if cutlass.const_expr(mState_init is not None): - bars.mb_state_k_acc_ready.wait(state_k_acc_index.phase) - - state_k_vec0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_addr + statek_col_id, cutlass.Float32), num=2) - state_k_vec1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row16_addr + statek_col_id, cutlass.Float32), num=2) - beta_pack = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) for reg_idx in cutlass.range_constexpr(4): - token0 = (((reg_idx // 2) * 4 + (lane & 3)) ^ 4) * 2 + token0 = ((reg_idx // 2) * 4 + (lane_idx & 3)) * 2 beta0 = (sBeta_ptr + token0).load().to(cutlass.Float32) beta1 = (sBeta_ptr + token0 + 1).load().to(cutlass.Float32) beta_pack[reg_idx] = fp32_to_fp16(beta0, beta1, dtype=cfg.io_dtype) - y_inp_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) - for reg_idx in cutlass.range_constexpr(4): - raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1) - frag_pair = (reg_idx ^ 2) * 2 - if cutlass.const_expr(mState_init is not None): - state_k_val0, state_k_val1 = state_k_vec0[frag_pair], state_k_vec0[frag_pair + 1] - state_k_pair = fp32_to_fp16(state_k_val0, state_k_val1, dtype=cfg.io_dtype) - diff_pair = sub_f16x2( - raw_v_frag0[raw_matrix], - state_k_pair, - cfg.io_dtype, - ) - else: - diff_pair = raw_v_frag0[raw_matrix] - y_inp_pack0[reg_idx] = mul_f16x2( - beta_pack[reg_idx], - diff_pair, - cfg.io_dtype, - ) - y_inp_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) - for reg_idx in cutlass.range_constexpr(4): - raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1) - frag_pair = (reg_idx ^ 2) * 2 - if cutlass.const_expr(mState_init is not None): - state_k_val0, state_k_val1 = state_k_vec1[frag_pair], state_k_vec1[frag_pair + 1] - state_k_pair = fp32_to_fp16(state_k_val0, state_k_val1, dtype=cfg.io_dtype) - diff_pair = sub_f16x2( - raw_v_frag1[raw_matrix], - state_k_pair, - cfg.io_dtype, - ) + y_lo = [cutlass.Int32(0) for _ in range(4)] + y_hi = [cutlass.Int32(0) for _ in range(4)] + if cutlass.const_expr(mState_init is not None): + if seed_from_initial_state: + bars.mb_state_k_acc_ready.wait(state_k_acc_index.phase) + state_k_acc_index = advance(state_k_acc_index, 1) + state_k_vec_lo = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_lo_addr + statek_col_id, cutlass.Float32), num=2) + state_k_vec_hi = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_hi_addr + statek_col_id, cutlass.Float32), num=2) + for reg_idx in cutlass.range_constexpr(4): + frag_pair = reg_idx * 2 + state_k_lo = fp32_to_fp16(state_k_vec_lo[frag_pair], state_k_vec_lo[frag_pair + 1], dtype=cfg.io_dtype) + state_k_hi = fp32_to_fp16(state_k_vec_hi[frag_pair], state_k_vec_hi[frag_pair + 1], dtype=cfg.io_dtype) + y_lo[reg_idx] = mul_f16x2(beta_pack[reg_idx], sub_f16x2(raw_v_frag_lo[reg_idx], state_k_lo, cfg.io_dtype), cfg.io_dtype) + y_hi[reg_idx] = mul_f16x2(beta_pack[reg_idx], sub_f16x2(raw_v_frag_hi[reg_idx], state_k_hi, cfg.io_dtype), cfg.io_dtype) else: - diff_pair = raw_v_frag1[raw_matrix] - y_inp_pack1[reg_idx] = mul_f16x2( - beta_pack[reg_idx], - diff_pair, - cfg.io_dtype, - ) + for reg_idx in cutlass.range_constexpr(4): + y_lo[reg_idx] = mul_f16x2(beta_pack[reg_idx], raw_v_frag_lo[reg_idx], cfg.io_dtype) + y_hi[reg_idx] = mul_f16x2(beta_pack[reg_idx], raw_v_frag_hi[reg_idx], cfg.io_dtype) + else: + for reg_idx in cutlass.range_constexpr(4): + y_lo[reg_idx] = mul_f16x2(beta_pack[reg_idx], raw_v_frag_lo[reg_idx], cfg.io_dtype) + y_hi[reg_idx] = mul_f16x2(beta_pack[reg_idx], raw_v_frag_hi[reg_idx], cfg.io_dtype) - nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(st_row_addr + y_inp_col_id, cutlass.Int8), y_inp_pack0[0:4]) - nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(st_row16_addr + y_inp_col_id, cutlass.Int8), y_inp_pack1[0:4]) + y_input_pack_lo = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + y_input_pack_hi = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + for reg_idx in cutlass.range_constexpr(4): + y_input_pack_lo[reg_idx] = y_lo[reg_idx] + y_input_pack_hi[reg_idx] = y_hi[reg_idx] + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_lo_addr + y_input_col_id, cutlass.Int8), y_input_pack_lo[0:4]) + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_hi_addr + y_input_col_id, cutlass.Int8), y_input_pack_hi[0:4]) nvvm.tcgen05_wait("store") - if cutlass.const_expr(mState_init is not None): - state_k_acc_index = advance(state_k_acc_index, 1) bars.mb_v_done[raw_index.idx].arrive() bars.mb_beta_done[raw_index.idx].arrive() - bars.mb_y_inp_ready.arrive() + bars.mb_y_input_ready.arrive() - # ---- U repack: u_acc TMEM -> packed b16 (U input) TMEM -------------- + # ---- U stage: u acc TMEM -> packed b16 U input TMEM ---------------------- bars.mb_u_acc_ready.wait(u_acc_index.phase) u_acc_vals = nvvm.tcgen05_ld( "32x32b", @@ -1413,21 +1301,20 @@ def compute1_warp_group( num=cfg.b_t, ) - u_inp_pack = cutlass.Array(cutlass.Int32, (cfg.b_t // 2), alignment=16) + u_input_pack = cutlass.Array(cutlass.Int32, (cfg.b_t // 2), alignment=16) for packed_col in cutlass.range_constexpr((cfg.b_t // 2)): - source_pair = packed_col ^ 4 - token0 = source_pair * 2 + token0 = packed_col * 2 token1 = token0 + 1 - u_inp_pack[packed_col] = fp32_to_fp16(u_acc_vals[token0], u_acc_vals[token1], dtype=cfg.io_dtype) + u_input_pack[packed_col] = fp32_to_fp16(u_acc_vals[token0], u_acc_vals[token1], dtype=cfg.io_dtype) nvvm.tcgen05_st( "32x32b", - nvvm.make_tmem_ptr(u_inp_addr, cutlass.Int8), - u_inp_pack[0 : (cfg.b_t // 2)], + nvvm.make_tmem_ptr(u_input_addr, cutlass.Int8), + u_input_pack[0 : (cfg.b_t // 2)], ) nvvm.tcgen05_wait("store") u_acc_index = advance(u_acc_index, 1) - bars.mb_u_inp_ready.arrive() + bars.mb_u_input_ready.arrive() if cutlass.const_expr(cfg.enable_checkpoints): bars.mb_state_acc_done.wait(state_upd_index.phase) state_upd_index = advance(state_upd_index, 1) @@ -1436,9 +1323,9 @@ def compute1_warp_group( if cutlass.const_expr(cfg.enable_checkpoints): cg1_checkpoint_chunks = checkpoint_every_n_tokens // cutlass.Int32(cfg.b_t) - cg1_checkpoint_mod = (cstart + cutlass.Int32(1)) % cg1_checkpoint_chunks + cg1_checkpoint_mod = (compute_start + cutlass.Int32(1)) % cg1_checkpoint_chunks for local_chunk_idx in cutlass.range(1, num_chunks_tile, 1, unroll=1): - chunk_idx = cstart + local_chunk_idx + chunk_idx = compute_start + local_chunk_idx sV_ptr = sV_raw.data_ptr() + raw_index.idx * (cfg.d_v * cfg.b_t) sBeta_ptr = sBeta_raw.data_ptr() + raw_index.idx * cfg.b_t @@ -1447,30 +1334,29 @@ def compute1_warp_group( do_checkpoint = cg1_checkpoint_mod == 0 cg1_checkpoint_mod = cg1_checkpoint_mod + cutlass.Int32(1) cg1_checkpoint_mod = cutlass.Int32(0) if cg1_checkpoint_mod == cg1_checkpoint_chunks else cg1_checkpoint_mod - do_checkpoint = do_checkpoint and chunk_idx >= wstart + do_checkpoint = do_checkpoint and chunk_idx >= write_start - # ---- state repack: acc TMEM -> packed b16 TMEM ---------------------- + # ---- state stage: acc TMEM -> packed b16 TMEM ---------------------------- if cutlass.const_expr(not cfg.enable_checkpoints): bars.mb_state_acc_done.wait(state_upd_index.phase) state_upd_index = advance(state_upd_index, 1) state_vecs = [] - for sub in cutlass.range_constexpr(cfg.d_k // 16): - state_vecs.append(nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + state_col_id + sub * 16, cutlass.Float32), num=16)) + for i in cutlass.range_constexpr(cfg.d_k // 16): + state_vecs.append(nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_lo_addr + state_col_id + i * 16, cutlass.Float32), num=16)) - for sub in cutlass.range_constexpr(cfg.d_k // 16): + for i in cutlass.range_constexpr(cfg.d_k // 16): state_pack = cutlass.Array(cutlass.Int32, 8, alignment=16) for packed_col in cutlass.range_constexpr(8): - source_pair = packed_col ^ 4 - state_pack[packed_col] = fp32_to_fp16(state_vecs[sub][2 * source_pair], state_vecs[sub][2 * source_pair + 1], dtype=cfg.io_dtype) + state_pack[packed_col] = fp32_to_fp16(state_vecs[i][2 * packed_col], state_vecs[i][2 * packed_col + 1], dtype=cfg.io_dtype) nvvm.tcgen05_st( "32x32b", - nvvm.make_tmem_ptr((tmem_row << 16) + packed_col_id + sub * 8, cutlass.Int8), + nvvm.make_tmem_ptr(row_lo_addr + packed_col_id + i * 8, cutlass.Int8), state_pack[0:8], ) nvvm.tcgen05_wait("store") - bars.mb_state_inp_ready.arrive() + bars.mb_state_input_ready.arrive() - # ---- checkpoint store ----------------------------------------------- + # ---- checkpoint store ---------------------------------------------------- if cutlass.const_expr(cfg.enable_checkpoints): if do_checkpoint: checkpoint_stage = checkpoint_done_index.idx @@ -1478,7 +1364,7 @@ def compute1_warp_group( checkpoint_done_index = advance(checkpoint_done_index, cfg.smem_checkpoint_stages) checkpoint_stage_base = checkpoint_stage * (cfg.d_k * cfg.d_v) for k_base in cutlass.range_constexpr(0, cfg.d_k, 32): - checkpoint_vec = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + state_col_id + k_base, cutlass.Float32), num=32) + checkpoint_vec = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_lo_addr + state_col_id + k_base, cutlass.Float32), num=32) for g in cutlass.range_constexpr(4): packs = tuple(fp32_to_fp16(checkpoint_vec[g * 8 + 2 * t], checkpoint_vec[g * 8 + 2 * t + 1], dtype=cfg.io_dtype) for t in range(4)) dk = k_base + g * 8 @@ -1493,74 +1379,72 @@ def compute1_warp_group( else: bars.mb_state_acc_read_done.arrive() - # ---- Y staging: Y = Beta * (V - State*K) ----------------------------- + # ---- Y stage: Y = Beta * (V - k state) ----------------------------------- bars.mb_v_ready[raw_index.idx].wait(raw_index.phase) - raw_v_frag0 = nvvm.ldmatrix( - sV_ptr + v_swz_off0, + raw_v_frag_lo = nvvm.ldmatrix( + sV_ptr + v_swz_off_lo, 4, nvvm.MMALayout.COL, ) - raw_v_frag1 = nvvm.ldmatrix( - sV_ptr + v_swz_off, + raw_v_frag_hi = nvvm.ldmatrix( + sV_ptr + v_swz_off_hi, 4, nvvm.MMALayout.COL, ) bars.mb_beta_ready[raw_index.idx].wait(raw_index.phase) - # ---- read back State*K acc ------------------------------------------- + # ---- read k state acc ---------------------------------------------------- bars.mb_state_k_acc_ready.wait(state_k_acc_index.phase) - state_k_vec0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_addr + statek_col_id, cutlass.Float32), num=2) - state_k_vec1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row16_addr + statek_col_id, cutlass.Float32), num=2) + state_k_vec_lo = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_lo_addr + statek_col_id, cutlass.Float32), num=2) + state_k_vec_hi = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_hi_addr + statek_col_id, cutlass.Float32), num=2) beta_pack = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) for reg_idx in cutlass.range_constexpr(4): - token0 = (((reg_idx // 2) * 4 + (lane & 3)) ^ 4) * 2 + token0 = ((reg_idx // 2) * 4 + (lane_idx & 3)) * 2 beta0 = (sBeta_ptr + token0).load().to(cutlass.Float32) beta1 = (sBeta_ptr + token0 + 1).load().to(cutlass.Float32) beta_pack[reg_idx] = fp32_to_fp16(beta0, beta1, dtype=cfg.io_dtype) - y_inp_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + y_input_pack_lo = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) for reg_idx in cutlass.range_constexpr(4): - raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1) - frag_pair = (reg_idx ^ 2) * 2 - state_k_val0, state_k_val1 = state_k_vec0[frag_pair], state_k_vec0[frag_pair + 1] + frag_pair = reg_idx * 2 + state_k_val0, state_k_val1 = state_k_vec_lo[frag_pair], state_k_vec_lo[frag_pair + 1] state_k_pair = fp32_to_fp16(state_k_val0, state_k_val1, dtype=cfg.io_dtype) diff_pair = sub_f16x2( - raw_v_frag0[raw_matrix], + raw_v_frag_lo[reg_idx], state_k_pair, cfg.io_dtype, ) - y_inp_pack0[reg_idx] = mul_f16x2( + y_input_pack_lo[reg_idx] = mul_f16x2( beta_pack[reg_idx], diff_pair, cfg.io_dtype, ) - y_inp_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + y_input_pack_hi = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) for reg_idx in cutlass.range_constexpr(4): - raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1) - frag_pair = (reg_idx ^ 2) * 2 - state_k_val0, state_k_val1 = state_k_vec1[frag_pair], state_k_vec1[frag_pair + 1] + frag_pair = reg_idx * 2 + state_k_val0, state_k_val1 = state_k_vec_hi[frag_pair], state_k_vec_hi[frag_pair + 1] state_k_pair = fp32_to_fp16(state_k_val0, state_k_val1, dtype=cfg.io_dtype) diff_pair = sub_f16x2( - raw_v_frag1[raw_matrix], + raw_v_frag_hi[reg_idx], state_k_pair, cfg.io_dtype, ) - y_inp_pack1[reg_idx] = mul_f16x2( + y_input_pack_hi[reg_idx] = mul_f16x2( beta_pack[reg_idx], diff_pair, cfg.io_dtype, ) - nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(st_row_addr + y_inp_col_id, cutlass.Int8), y_inp_pack0[0:4]) - nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(st_row16_addr + y_inp_col_id, cutlass.Int8), y_inp_pack1[0:4]) + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_lo_addr + y_input_col_id, cutlass.Int8), y_input_pack_lo[0:4]) + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_hi_addr + y_input_col_id, cutlass.Int8), y_input_pack_hi[0:4]) nvvm.tcgen05_wait("store") state_k_acc_index = advance(state_k_acc_index, 1) bars.mb_v_done[raw_index.idx].arrive() bars.mb_beta_done[raw_index.idx].arrive() - bars.mb_y_inp_ready.arrive() + bars.mb_y_input_ready.arrive() - # ---- U repack: u_acc TMEM -> packed b16 (U input) TMEM -------------- + # ---- U stage: u acc TMEM -> packed b16 U input TMEM ---------------------- bars.mb_u_acc_ready.wait(u_acc_index.phase) u_acc_vals = nvvm.tcgen05_ld( "32x32b", @@ -1568,21 +1452,20 @@ def compute1_warp_group( num=cfg.b_t, ) - u_inp_pack = cutlass.Array(cutlass.Int32, (cfg.b_t // 2), alignment=16) + u_input_pack = cutlass.Array(cutlass.Int32, (cfg.b_t // 2), alignment=16) for packed_col in cutlass.range_constexpr((cfg.b_t // 2)): - source_pair = packed_col ^ 4 - token0 = source_pair * 2 + token0 = packed_col * 2 token1 = token0 + 1 - u_inp_pack[packed_col] = fp32_to_fp16(u_acc_vals[token0], u_acc_vals[token1], dtype=cfg.io_dtype) + u_input_pack[packed_col] = fp32_to_fp16(u_acc_vals[token0], u_acc_vals[token1], dtype=cfg.io_dtype) nvvm.tcgen05_st( "32x32b", - nvvm.make_tmem_ptr(u_inp_addr, cutlass.Int8), - u_inp_pack[0 : (cfg.b_t // 2)], + nvvm.make_tmem_ptr(u_input_addr, cutlass.Int8), + u_input_pack[0 : (cfg.b_t // 2)], ) nvvm.tcgen05_wait("store") u_acc_index = advance(u_acc_index, 1) - bars.mb_u_inp_ready.arrive() + bars.mb_u_input_ready.arrive() if cutlass.const_expr(cfg.enable_checkpoints): bars.mb_state_acc_done.wait(state_upd_index.phase) @@ -1594,18 +1477,18 @@ def compute1_warp_group( bars.mb_state_acc_done.wait(state_upd_index.phase) state_upd_index = advance(state_upd_index, 1) - owns_final = wend == num_chunks_b + owns_final = write_end == batch_num_chunks - # ---- final-state drain: state acc TMEM -> GMEM --------------------------- + # ---- final state store: TMEM -> GMEM ----------------------------------------- if cutlass.const_expr(mState_out is not None): - if seqlen_b > 0: + if batch_seqlen > 0: if owns_final: state_vw = 16 // (mState_out.element_type.width // 8) state_dst = (mState_out.iterator + mState_out.layout((batch_idx, head_o, value_dim, 0))).raw_ptr() for key_block_start in cutlass.range_constexpr(0, cfg.d_k, 32): loaded = nvvm.tcgen05_ld( "32x32b", - nvvm.make_tmem_ptr((row_id << 16) + (tmem_col + cfg.tmem_state_acc_offset + key_block_start), cutlass.Float32), + nvvm.make_tmem_ptr(row_lo_addr + (tmem_col + cfg.tmem_state_acc_offset + key_block_start), cutlass.Float32), num=32, ) @@ -1625,7 +1508,7 @@ def compute1_warp_group( mState_out[batch_idx, head_o, value_dim, key_dim] = mState_init[batch_idx, head_o, value_dim, key_dim] else: mState_out[batch_idx, head_o, value_dim, key_dim] = cutlass.Float32(0.0).to(mState_out.element_type) - tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, scheduler_state = scheduler_next_tile(cfg, bars, sScheduler, scheduler_state, tile_idx, num_ctas, elect_one) bars.mb_tmem_done[0].arrive() @@ -1644,7 +1527,7 @@ def host( final_state: cute.Tensor | None, work_items: cute.Tensor | None, work_count: cute.Tensor | None, - sched_ctr: cute.Tensor | None, + scheduler_counter: cute.Tensor | None, tensormap_workspace: cute.Tensor, checkpoint_every_n_tokens: cutlass.Int32, stream, @@ -1666,7 +1549,7 @@ def host( final_state, work_items, work_count, - sched_ctr, + scheduler_counter, checkpoint_every_n_tokens, ).launch( grid=grid_shape, @@ -1692,7 +1575,7 @@ def kernel( mState_out: cute.Tensor | None, mWorkItems: cute.Tensor, mCount: cute.Tensor, - mSched: cute.Tensor | None, + mScheduler: cute.Tensor | None, checkpoint_every_n_tokens: cutlass.Int32, ) -> None: """BT=16 KDA recompute (state/checkpoints-only) persistent kernel body: every warp @@ -1702,11 +1585,11 @@ def kernel( bidx = cute.arch.block_idx()[0] num_ctas = cute.arch.grid_dim()[0] warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - lane = tidx % cfg.threads_per_warp + lane_idx = tidx % cfg.threads_per_warp total_tiles = mCount[0] - if cutlass.const_expr(cfg.dyn_sched): - assert mSched is not None and mSched.element_type == cutlass.Int32 + if cutlass.const_expr(cfg.dynamic_scheduling): + assert mScheduler is not None and mScheduler.element_type == cutlass.Int32 assert mK.element_type == cfg.io_dtype and mV.element_type == cfg.io_dtype assert mGate.element_type == cutlass.Float32 beta_expected = cfg.io_dtype if cutlass.const_expr(cfg.beta_sigmoid) else cutlass.Float32 @@ -1729,11 +1612,10 @@ def kernel( desc_gate_base = desc_base_words + cutlass.Int32(2) * arr_words desc_checkpoint_base = desc_base_words + cutlass.Int32(3) * arr_words - # Buffers are declaration-ordered and intentionally non-aliased. SMEM = cutlass.AddressSpace.smem bars = make_kda_bars(cfg) sTmem_base = cutlass.Array(cutlass.Int32, 1, space=SMEM, alignment=4) - sSched = cutlass.Array(cutlass.Int32, cfg.sched_stages, space=SMEM, alignment=16) + sScheduler = cutlass.Array(cutlass.Int32, cfg.scheduler_stages, space=SMEM, alignment=16) sK_decay_raw = cutlass.Array(cfg.io_dtype, cfg.k_decay_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) sK_restore_raw = cutlass.Array(cfg.io_dtype, cfg.k_restore_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) sIntermediate_raw = cutlass.Array(cfg.io_dtype, cfg.intermediate_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) @@ -1756,7 +1638,7 @@ def kernel( stride_byte_offset=1024, layout=nvvm.Tcgen05SmemSwizzle.SWIZZLE_128B, ) - sK_restore = SmemTile( + sK_restore_trans = SmemTile( base=sK_restore_raw, elems_per_stage=(cfg.d_k * cfg.b_t), stages=cfg.smem_decay_stages, @@ -1798,15 +1680,15 @@ def kernel( bars.mb_state_k_acc_ready.init() bars.mb_u_acc_ready.init() bars.mb_state_acc_done.init() - bars.mb_state_inp_ready.init() + bars.mb_state_input_ready.init() for stage in cutlass.range_constexpr(cfg.smem_state_scale_diag_stages): bars.mb_state_scale_diag_done[stage].init() for stage in cutlass.range_constexpr(cfg.smem_decay_stages): bars.mb_decay_tcgen05_done[stage].init() bars.mb_decay_super_done[stage].init() bars.mb_k_restore_done[stage].init() - bars.mb_y_inp_ready.init() - bars.mb_u_inp_ready.init() + bars.mb_y_input_ready.init() + bars.mb_u_input_ready.init() bars.mb_tmem_done[0].init() elif warp_idx == cfg.super_mma_warp_id: if elect_one: @@ -1819,9 +1701,9 @@ def kernel( bars.mb_k_decay_inv_cg0_ready[stage].init() elif warp_idx == cfg.epilogue_warp_id: if elect_one: - for stage in cutlass.range_constexpr(cfg.sched_stages): - bars.mb_sched_ready[stage].init() - bars.mb_sched_done[stage].init() + for stage in cutlass.range_constexpr(cfg.scheduler_stages): + bars.mb_scheduler_ready[stage].init() + bars.mb_scheduler_done[stage].init() if cutlass.const_expr(cfg.enable_checkpoints): for stage in cutlass.range_constexpr(cfg.smem_checkpoint_stages): bars.mb_checkpoint_tmastg_ready[stage].init() @@ -1840,9 +1722,9 @@ def kernel( num_ctas, cu_seqlens, mWorkItems, - mSched, - sSched, - lane, + mScheduler, + sScheduler, + lane_idx, sK_raw, sV_raw, sGate_raw, @@ -1859,8 +1741,8 @@ def kernel( num_ctas, cu_seqlens, mWorkItems, - sSched, - lane, + sScheduler, + lane_idx, sK_inv_raw, sIntermediate_raw, sBeta_raw, @@ -1875,11 +1757,11 @@ def kernel( num_ctas, cu_seqlens, mWorkItems, - sSched, + sScheduler, sTmem_base, sIntermediate, sK_decay, - sK_restore, + sK_restore_trans, sState_scale_diag, bars, ) @@ -1891,7 +1773,7 @@ def kernel( num_ctas, cu_seqlens, mWorkItems, - sSched, + sScheduler, sCheckpoint_raw, desc_checkpoint_base, checkpoint_every_n_tokens, @@ -1905,8 +1787,8 @@ def kernel( num_ctas, cu_seqlens, mWorkItems, - sSched, - lane, + sScheduler, + lane_idx, warp_idx, mA_log, mDt_bias, @@ -1928,8 +1810,8 @@ def kernel( num_ctas, cu_seqlens, mWorkItems, - sSched, - lane, + sScheduler, + lane_idx, sTmem_base, warp_idx, mState_out, @@ -1962,8 +1844,8 @@ class KdaRecomputeCfg: v_ratio: int n_heads_out: int max_active_clusters: int - dyn_sched: bool = False - sched_stages: int = CFG.SMEM_SCHED_STAGES + dynamic_scheduling: bool = False + scheduler_stages: int = CFG.SMEM_SCHEDULER_STAGES compute_group_0_warp_ids: tuple[int, ...] = CFG.COMPUTE_GROUP_0_WARP_IDS compute_group_1_warp_ids: tuple[int, ...] = CFG.COMPUTE_GROUP_1_WARP_IDS @@ -1980,7 +1862,7 @@ class KdaRecomputeCfg: cg0_group_count: int = 2 cg0_warps_per_group: int = 4 cg0_threads_per_group: int = 0 - cg0_group_sync_barrier_base_id: int = 1 # CG0 group g syncs on nbar id 1 + g + cg0_group_sync_barrier_base_id: int = 1 # CG0 group g syncs on named-barrier id 1 + g cg0_tile_entry_barrier_id: int = 5 # CG0-wide (both groups) work-item entry sync tmem_user_threads: int = 0 tmem_lifecycle_barrier_id: int = 3 @@ -1988,7 +1870,7 @@ class KdaRecomputeCfg: num_regs_compute_group_1: int = CFG.NUM_REGS_COMPUTE_GROUP_1 num_regs_other: int = CFG.NUM_REGS_OTHER - # ---- SMEM / TMEM ring stage counts ------------------------------------------- + # ---- SMEM / TMEM ring stage counts ----------------------------------------------- smem_raw_stages: int = CFG.SMEM_RAW_STAGES smem_checkpoint_stages: int = 1 smem_decay_stages: int = CFG.SMEM_DECAY_STAGES @@ -1996,15 +1878,15 @@ class KdaRecomputeCfg: smem_state_scale_diag_stages: int = CFG.SMEM_STATE_SCALE_DIAG_STAGES qk_scale_ready_stages: int = CFG.QK_SCALE_READY_STAGES - # ---- TMEM column offsets (state doubles as the final_state acc) -------------- + # ---- TMEM column offsets (state doubles as the final state acc) ------------------ tmem_state_acc_offset: int = 0 - tmem_state_inp_offset: int = 0 + tmem_state_input_offset: int = 0 tmem_state_k_acc_offset: int = 0 tmem_u_acc_offset: int = 0 - tmem_y_inp_offset: int = 0 - tmem_u_inp_offset: int = 0 + tmem_y_input_offset: int = 0 + tmem_u_input_offset: int = 0 - # ---- SMEM buffer cosizes ----------------------------------------------------- + # ---- SMEM buffer cosizes --------------------------------------------------------- k_cosize: int = 0 v_cosize: int = 0 gate_cosize: int = 0 @@ -2036,7 +1918,7 @@ def build_cfg( v_ratio: int, n_heads_out: int, max_active_clusters: int, - dyn_sched: bool = False, + dynamic_scheduling: bool = False, ) -> KdaRecomputeCfg: """Build the per-compile ``KdaRecomputeCfg`` (io_dtype in {Float16, BFloat16}); fills the derived TMEM column offsets and SMEM buffer cosizes.""" @@ -2056,7 +1938,7 @@ def build_cfg( v_ratio=v_ratio, n_heads_out=n_heads_out, max_active_clusters=max_active_clusters, - dyn_sched=dyn_sched, + dynamic_scheduling=dynamic_scheduling, ) if enable_checkpoints: cfg.smem_raw_stages = 6 @@ -2069,12 +1951,12 @@ def build_cfg( if cfg.smem_state_scale_diag_stages != cfg.qk_scale_ready_stages: raise ValueError("diag and qk-scale ready rings must share their rolling stage") - cfg.tmem_state_inp_offset = cfg.tmem_state_acc_offset + cfg.d_k - cfg.tmem_state_k_acc_offset = cfg.tmem_state_inp_offset + (cfg.d_k // 2) + cfg.tmem_state_input_offset = cfg.tmem_state_acc_offset + cfg.d_k + cfg.tmem_state_k_acc_offset = cfg.tmem_state_input_offset + (cfg.d_k // 2) cfg.tmem_u_acc_offset = cfg.tmem_state_k_acc_offset + cfg.b_t - cfg.tmem_y_inp_offset = cfg.tmem_u_acc_offset + cfg.b_t - cfg.tmem_u_inp_offset = cfg.tmem_y_inp_offset + (cfg.b_t // 2) - assert (cfg.tmem_u_inp_offset + (cfg.b_t // 2)) <= 512 + cfg.tmem_y_input_offset = cfg.tmem_u_acc_offset + cfg.b_t + cfg.tmem_u_input_offset = cfg.tmem_y_input_offset + (cfg.b_t // 2) + assert (cfg.tmem_u_input_offset + (cfg.b_t // 2)) <= 512 cfg.k_cosize = cfg.smem_raw_stages * cfg.d_k * cfg.b_t cfg.v_cosize = cfg.smem_raw_stages * cfg.d_v * cfg.b_t @@ -2102,7 +1984,7 @@ def build_descs_body( base_v, base_gate, base_checkpoint, - desc_ws: cute.Tensor, + desc_workspace: cute.Tensor, cu_seqlens: cute.Tensor, k: cute.Tensor, v: cute.Tensor, @@ -2119,10 +2001,10 @@ def build_descs_body( prologue kernel after its order pass; warps past the array count fall through the widx guards.""" arr_words = n_batch * cutlass.Int32(TENSOR_MAP_QWORDS) - desc_words_k = cute.make_tensor(desc_ws.iterator, cute.make_layout((arr_words,), stride=(1,))) - desc_words_v = cute.make_tensor(desc_ws.iterator + arr_words, cute.make_layout((arr_words,), stride=(1,))) - desc_words_gate = cute.make_tensor(desc_ws.iterator + 2 * arr_words, cute.make_layout((arr_words,), stride=(1,))) - desc_words_checkpoint = cute.make_tensor(desc_ws.iterator + 3 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + desc_words_k = cute.make_tensor(desc_workspace.iterator, cute.make_layout((arr_words,), stride=(1,))) + desc_words_v = cute.make_tensor(desc_workspace.iterator + arr_words, cute.make_layout((arr_words,), stride=(1,))) + desc_words_gate = cute.make_tensor(desc_workspace.iterator + 2 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + desc_words_checkpoint = cute.make_tensor(desc_workspace.iterator + 3 * arr_words, cute.make_layout((arr_words,), stride=(1,))) if widx == 0: if nvvm.elect_sync(): @@ -2149,13 +2031,13 @@ def build_descs_body( def prologue_kernel( run_order: cutlass.Constexpr[bool], order_gen: cutlass.Constexpr[bool], - has_sched: cutlass.Constexpr[bool], + has_scheduler: cutlass.Constexpr[bool], b_t: cutlass.Constexpr[int], base_k: cutlass.GridConstant[cuda.tensor_map.TensorMap], base_v: cutlass.GridConstant[cuda.tensor_map.TensorMap], base_gate: cutlass.GridConstant[cuda.tensor_map.TensorMap], base_checkpoint: cutlass.GridConstant[cuda.tensor_map.TensorMap], - desc_ws: cute.Tensor, + desc_workspace: cute.Tensor, cu_seqlens: cute.Tensor, k: cute.Tensor, v: cute.Tensor, @@ -2164,7 +2046,7 @@ def prologue_kernel( mStaging: cute.Tensor | None, mCount: cute.Tensor, mWorkItems: cute.Tensor | None, - mSched: cute.Tensor | None, + mScheduler: cute.Tensor | None, n_batch: cutlass.Int32, k_row_stride: cutlass.Int32, v_row_stride: cutlass.Int32, @@ -2174,7 +2056,7 @@ def prologue_kernel( ) -> None: """Single-CTA prologue. Under ``run_order`` this kernel is the first work-item-table consumer, so it LPT-orders the table and zeroes both - consumers' sched rings via :func:`order_body`; it then builds the + consumers' scheduler rings via :func:`order_body`; it then builds the per-batch TMA-descriptor arrays via :func:`build_descs_body`, one warp per array (the extra warps only take part in the order phase).""" tidx, _, _ = cute.arch.thread_idx() @@ -2187,10 +2069,10 @@ def prologue_kernel( n_heads_out = cutlass.Int32(gate.shape[1]) order_body( order_gen, - has_sched, + has_scheduler, b_t, ORDER_THREADS, - ORDER_ELEMS, + ORDER_ELEMENTS, tidx, n_heads_out, n_heads_out * n_batch, @@ -2198,7 +2080,7 @@ def prologue_kernel( mStaging, mCount, mWorkItems, - mSched, + mScheduler, sKey, sIdx, sSpread, @@ -2209,7 +2091,7 @@ def prologue_kernel( base_v, base_gate, base_checkpoint, - desc_ws, + desc_workspace, cu_seqlens, k, v, @@ -2230,7 +2112,7 @@ def prologue( b_t: cutlass.Constexpr[int], run_order: cutlass.Constexpr[bool], order_gen: cutlass.Constexpr[bool], - has_sched: cutlass.Constexpr[bool], + has_scheduler: cutlass.Constexpr[bool], k: cute.Tensor, v: cute.Tensor, gate: cute.Tensor, @@ -2239,7 +2121,7 @@ def prologue( work_item_staging: cute.Tensor | None, work_count: cute.Tensor, work_items: cute.Tensor | None, - sched_all: cute.Tensor | None, + scheduler_all: cute.Tensor | None, tensormap_workspace: cute.Tensor, checkpoint_every_n: cutlass.Int32, stream: cuda_driver.CUstream, @@ -2279,7 +2161,7 @@ def prologue( prologue_kernel( run_order, order_gen, - has_sched, + has_scheduler, b_t, base_k, base_v, @@ -2294,7 +2176,7 @@ def prologue( work_item_staging, work_count, work_items, - sched_all, + scheduler_all, cutlass.Int32(batch_size), cutlass.Int32(k.stride[0]), cutlass.Int32(v.stride[0]), @@ -2304,7 +2186,7 @@ def prologue( ).launch(grid=(1, 1, 1), block=(ORDER_THREADS, 1, 1), stream=stream) -# ---- Torch adapter / host-side compilation --------------------------------------- +# ---- Torch adapter / host-side compilation ------------------------------------------- @lru_cache(maxsize=None) @@ -2322,7 +2204,7 @@ def get_compiled_cache( safe_gate: bool, gate_lower_bound: float, beta_sigmoid: bool, - dyn_sched: bool, + dynamic_scheduling: bool, run_order: bool, order_gen: bool, ): @@ -2343,7 +2225,7 @@ def compile( k_ratio: int, v_ratio: int, n_heads_out: int, - dyn_sched: bool = False, + dynamic_scheduling: bool = False, *, num_sm: int, k_cute, @@ -2357,8 +2239,8 @@ def compile( state_out_cute, work_items_cute=None, work_count_cute=None, - sched_ctr_cute=None, - tensormap_ws_cute, + scheduler_counter_cute=None, + tensormap_workspace_cute, checkpoint_every_n_tokens, stream, ): @@ -2377,7 +2259,7 @@ def compile( v_ratio=v_ratio, n_heads_out=n_heads_out, max_active_clusters=num_sm, - dyn_sched=dyn_sched, + dynamic_scheduling=dynamic_scheduling, ) return cute.compile( @@ -2394,8 +2276,8 @@ def compile( state_out_cute, work_items_cute, work_count_cute, - sched_ctr_cute, - tensormap_ws_cute, + scheduler_counter_cute, + tensormap_workspace_cute, checkpoint_every_n_tokens, stream, options="--enable-tvm-ffi --opt-level 2", @@ -2420,8 +2302,8 @@ def chunk_kda_recompute_sm100( use_beta_sigmoid: bool = False, work_items=None, work_count=None, - sched_ctr=None, - sched_all=None, + scheduler_counter=None, + scheduler_all=None, work_item_scratch=None, order_in_prologue: bool = False, *, @@ -2462,13 +2344,14 @@ def chunk_kda_recompute_sm100( use_beta_sigmoid: ``beta`` holds logits; sigmoid in-kernel work_items: ``(max_items, 8)`` int32 work-item table from ``common/split_k.py`` (REQUIRED; an uncut table row is the whole - (b, h) sequence). Each item computes chunks ``[cstart, wend)`` - and writes checkpoints only for ``[wstart, wend)``. + (b, h) sequence). Each item computes chunks ``[compute_start, write_end)`` + and writes checkpoints only for ``[write_start, write_end)``. work_count: ``(1,)`` int32 device-side item count (REQUIRED) - sched_ctr: ``(2,)`` int32 device scratch ``[ticket, done]`` enabling + scheduler_counter: ``(2,)`` int32 device scratch ``[ticket, done]`` enabling the dynamic (work-stealing) tile scheduler; must be zeroed before - every launch (``build_split_table`` does this when it is passed as - ``sched_ctr``). None keeps the static CTA stride. + every launch (the split-table stage and the order-generating + prologue both zero it when passed as ``scheduler_counter``). + None keeps the static CTA stride. """ HK = k.shape[1] HV = v.shape[1] @@ -2484,12 +2367,14 @@ def chunk_kda_recompute_sm100( f"output_state_checkpoints dtype must match the io dtype (fp32 state belongs to output_state): got {output_state_checkpoints.dtype} with io {k.dtype}" ) if work_items is None or work_count is None: - raise ValueError("work_items/work_count are required (the split-table stage builds them for every launch)") - dyn_sched = sched_ctr is not None + raise ValueError( + "work_items/work_count are required (built by the split-table stage, or by an order-generating prologue when work_item_scratch is None)" + ) + dynamic_scheduling = scheduler_counter is not None run_order = order_in_prologue order_gen = order_in_prologue and work_item_scratch is None - if run_order and sched_all is None: - raise ValueError("order in the prologue requires sched_all (the prologue zeroes both consumers' sched rings)") + if run_order and scheduler_all is None: + raise ValueError("order in the prologue requires scheduler_all (the prologue zeroes both consumers' scheduler rings)") if initial_state is not None: state_dtype_src = initial_state.dtype @@ -2526,7 +2411,7 @@ def chunk_kda_recompute_sm100( safe_gate, gate_lower_bound, use_beta_sigmoid, - dyn_sched, + dynamic_scheduling, run_order, order_gen, ) @@ -2554,11 +2439,11 @@ def chunk_kda_recompute_sm100( work_items_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1), divisibility=1) work_count_cute = from_dlpack(work_count, assumed_align=4).mark_layout_dynamic() - sched_ctr_cute = None - if dyn_sched: - sched_ctr_cute = from_dlpack(sched_ctr, assumed_align=4).mark_layout_dynamic() + scheduler_counter_cute = None + if dynamic_scheduling: + scheduler_counter_cute = from_dlpack(scheduler_counter, assumed_align=4).mark_layout_dynamic() - tensormap_ws_cute = from_dlpack(tensormap_workspace, assumed_align=128).mark_layout_dynamic() + tensormap_workspace_cute = from_dlpack(tensormap_workspace, assumed_align=128).mark_layout_dynamic() cache["compiled"] = compile( io_dtype, @@ -2573,7 +2458,7 @@ def chunk_kda_recompute_sm100( k_ratio, v_ratio, HO, - dyn_sched, + dynamic_scheduling, num_sm=multiprocessor_count(current_device()), k_cute=k_cute, v_cute=v_cute, @@ -2586,8 +2471,8 @@ def chunk_kda_recompute_sm100( state_out_cute=state_out_cute, work_items_cute=work_items_cute, work_count_cute=work_count_cute, - sched_ctr_cute=sched_ctr_cute, - tensormap_ws_cute=tensormap_ws_cute, + scheduler_counter_cute=scheduler_counter_cute, + tensormap_workspace_cute=tensormap_workspace_cute, checkpoint_every_n_tokens=checkpoint_every_n_tokens, stream=cu_stream, ) @@ -2600,7 +2485,7 @@ def chunk_kda_recompute_sm100( v_pl = from_dlpack(v, assumed_align=16).mark_layout_dynamic(leading_dim=2) gate_pl = from_dlpack(gate, assumed_align=16).mark_layout_dynamic(leading_dim=2) cu_pl = from_dlpack(cu_seqlens, assumed_align=8).mark_layout_dynamic() - ws_pl = from_dlpack(tensormap_workspace, assumed_align=128).mark_layout_dynamic() + workspace_pl = from_dlpack(tensormap_workspace, assumed_align=128).mark_layout_dynamic() state_checkpoints_pl = None if state_checkpoints_for_descs is not None: state_checkpoints_pl = from_dlpack(state_checkpoints_for_descs, assumed_align=16).mark_layout_dynamic(leading_dim=3) @@ -2611,9 +2496,9 @@ def chunk_kda_recompute_sm100( work_items_pl = from_dlpack(work_items, assumed_align=16) work_items_pl.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1), divisibility=1) work_count_pl = from_dlpack(work_count, assumed_align=4).mark_layout_dynamic() - sched_pl = None + scheduler_pl = None if run_order: - sched_pl = from_dlpack(sched_all, assumed_align=4).mark_layout_dynamic() + scheduler_pl = from_dlpack(scheduler_all, assumed_align=4).mark_layout_dynamic() cache["prologue"] = cute.compile( prologue, io_dtype, @@ -2629,8 +2514,8 @@ def chunk_kda_recompute_sm100( staging_pl, work_count_pl, work_items_pl, - sched_pl, - ws_pl, + scheduler_pl, + workspace_pl, cutlass.Int32(checkpoint_every_n_tokens), cu_stream, options="--enable-tvm-ffi", @@ -2644,7 +2529,7 @@ def chunk_kda_recompute_sm100( work_item_scratch if run_order else None, work_count, work_items, - sched_all if run_order else None, + scheduler_all if run_order else None, tensormap_workspace, checkpoint_every_n_tokens, cu_stream, @@ -2661,7 +2546,7 @@ def chunk_kda_recompute_sm100( output_state if store_final_state else None, work_items, work_count, - sched_ctr, + scheduler_counter, tensormap_workspace, checkpoint_every_n_tokens, cu_stream, @@ -2683,8 +2568,8 @@ def run_recompute( output_state_checkpoints, work_items, work_count, - sched_ctr, - sched_all, + scheduler_counter, + scheduler_all, work_item_scratch, tensormap_workspace, checkpoint_every_n_tokens, @@ -2703,7 +2588,7 @@ def run_recompute( work_item_scratch, work_count, work_items, - sched_all, + scheduler_all, tensormap_workspace, checkpoint_every_n_tokens, cu_stream, @@ -2720,7 +2605,7 @@ def run_recompute( output_state, work_items, work_count, - sched_ctr, + scheduler_counter, tensormap_workspace, checkpoint_every_n_tokens, cu_stream, diff --git a/python/cudnn/linear_attention/graph_analyzer.py b/python/cudnn/linear_attention/graph_analyzer.py index ef09ee556..9b8817339 100644 --- a/python/cudnn/linear_attention/graph_analyzer.py +++ b/python/cudnn/linear_attention/graph_analyzer.py @@ -136,7 +136,7 @@ def analyze(graph: "cudnn.pygraph") -> Optional[LaGraphFacts]: required_out = (["dQ", "dK", "dV", "dG", "dBeta"] + (["dW"] if op == "GDN2" else [])) if is_bwd else ["O"] safe_gate = bool(params.get("safe_gate", False)) - ckpt = int(params.get("checkpoint_every_n_tokens", 0) or 0) + checkpoint = int(params.get("checkpoint_every_n_tokens", 0) or 0) invalid = None missing_in = [p for p in required_in if p not in ins] missing_out = [p for p in required_out if p not in outs] @@ -158,11 +158,11 @@ def analyze(graph: "cudnn.pygraph") -> Optional[LaGraphFacts]: invalid = "d_a_log/d_dt_bias dims must match a_log/dt_bias" elif params.get("gate_lower_bound") is not None and not safe_gate: invalid = "gate_lower_bound requires safe_gate=True" - elif ckpt < 0: + elif checkpoint < 0: invalid = "checkpoint_every_n_tokens must be non-negative" - elif not is_bwd and ckpt > 0 and "state_checkpoints" not in outs: + elif not is_bwd and checkpoint > 0 and "state_checkpoints" not in outs: invalid = "checkpoint_every_n_tokens > 0 requires the state_checkpoints output" - elif not is_bwd and ckpt == 0 and "state_checkpoints" in outs: + elif not is_bwd and checkpoint == 0 and "state_checkpoints" in outs: invalid = "state_checkpoints output requires checkpoint_every_n_tokens > 0" if invalid is not None: return LaGraphFacts(invalid=invalid, op=op, is_bwd=is_bwd) @@ -181,10 +181,10 @@ def analyze(graph: "cudnn.pygraph") -> Optional[LaGraphFacts]: total_t = int(q.dim[0]) if thd_layout else 0 cu = ins["cu_seqlens"] n_seq = int(cu.dim[0]) - 1 if cu.dim else 0 - ckpt_port = ins.get("state_checkpoints") - if ckpt_port is None: - ckpt_port = outs.get("state_checkpoints") - state_checkpoint_rows = int(ckpt_port.dim[0]) if ckpt_port is not None and ckpt_port.dim else 0 + checkpoint_port = ins.get("state_checkpoints") + if checkpoint_port is None: + checkpoint_port = outs.get("state_checkpoints") + state_checkpoint_rows = int(checkpoint_port.dim[0]) if checkpoint_port is not None and checkpoint_port.dim else 0 gates_at_ho = all(t is None or not t.dim or (len(t.dim) > 1 and int(t.dim[1]) == h_o) for t in (ins["g"], ins["beta"], ins.get("w"))) io_dtypes = {in_dt["q"], in_dt["k"], in_dt["v"]} - {None} state_dtypes = {in_dt.get("initial_state"), out_dt.get("final_state")} - {None} @@ -237,6 +237,6 @@ def analyze(graph: "cudnn.pygraph") -> Optional[LaGraphFacts]: safe_gate=safe_gate, use_beta_sigmoid=bool(params.get("use_beta_sigmoid", False)), gate_lower_bound=float(params["gate_lower_bound"]) if params.get("gate_lower_bound") is not None else None, - checkpoint_every_n_tokens=ckpt, + checkpoint_every_n_tokens=checkpoint, batch_invariant=bool(params.get("batch_invariant", False)), ) diff --git a/python/cudnn/linear_attention/ops/gdn.py b/python/cudnn/linear_attention/ops/gdn.py index ffab63a8a..cc177f4df 100644 --- a/python/cudnn/linear_attention/ops/gdn.py +++ b/python/cudnn/linear_attention/ops/gdn.py @@ -39,7 +39,7 @@ # --------------------------------------------------------------------------- -_TORCH_DTYPE_TO_CUDNN = { +TORCH_DTYPE_TO_CUDNN = { torch.float16: cudnn.data_type.HALF, torch.bfloat16: cudnn.data_type.BFLOAT16, torch.float32: cudnn.data_type.FLOAT, @@ -47,10 +47,9 @@ torch.int64: cudnn.data_type.INT64, } -# one graph per static configuration (shapes, dtypes, scale, flags, device) -_fprop_cache: Dict[tuple, tuple] = {} -_bprop_cache: Dict[tuple, tuple] = {} -_cudnn_handles: Dict[int, int] = {} +fprop_cache: Dict[tuple, tuple] = {} +bprop_cache: Dict[tuple, tuple] = {} +cudnn_handles: Dict[int, int] = {} # --------------------------------------------------------------------------- @@ -72,13 +71,11 @@ def select_plan(graph, plan_name): graph.check_support() -def _graph_workspace(graph, device): +def graph_workspace(graph, device): """Caller-side workspace for a compiled graph (grow-only, held on the graph object itself — same lifetime by construction, no id()-keyed side table).""" if not graph._is_built: - # plan first: a bare build() would lower GDN to the backend, - # which has no lowering if not graph._planning_done: graph.create_execution_plans() if graph.selected_engine is None: @@ -86,36 +83,36 @@ def _graph_workspace(graph, device): else: graph.build_plans() size = graph.get_workspace_size() - ws = getattr(graph, "_la_ops_workspace", None) - if ws is None or ws.numel() < size or ws.device != device: - ws = torch.empty(max(size, 1), dtype=torch.uint8, device=device) - graph._la_ops_workspace = ws - return ws + workspace = getattr(graph, "la_ops_workspace", None) + if workspace is None or workspace.numel() < size or workspace.device != device: + workspace = torch.empty(max(size, 1), dtype=torch.uint8, device=device) + graph.la_ops_workspace = workspace + return workspace -def _get_handle(device): +def get_handle(device): """Per-device cuDNN handle carrying the caller's current stream.""" idx = device.index if device.index is not None else torch.cuda.current_device() - handle = _cudnn_handles.get(idx) + handle = cudnn_handles.get(idx) if handle is None: with torch.cuda.device(idx): handle = cudnn.create_handle() - _cudnn_handles[idx] = handle + cudnn_handles[idx] = handle cudnn.set_stream(handle=handle, stream=torch.cuda.current_stream(device).cuda_stream) return handle -def _torch_dtype_to_cudnn(dtype: torch.dtype): +def torch_dtype_to_cudnn(dtype: torch.dtype): """Map a PyTorch dtype to a cuDNN data_type enum.""" - return _TORCH_DTYPE_TO_CUDNN[dtype] + return TORCH_DTYPE_TO_CUDNN[dtype] -def _check_dtype(name, t, want) -> None: +def check_dtype(name, t, want) -> None: if t.dtype != want: raise TypeError(f"gated_delta_net: {name} must be {want} (kernel-native; callers convert), got {t.dtype}") -def _make_fprop_cache_key( +def make_fprop_cache_key( total, N, H, @@ -136,7 +133,7 @@ def _make_fprop_cache_key( use_beta_sigmoid, safe_gate, has_initial_state, - ckpt, + checkpoint, device, plan_name, ): @@ -162,13 +159,13 @@ def _make_fprop_cache_key( bool(use_beta_sigmoid), bool(safe_gate), bool(has_initial_state), - ckpt, + checkpoint, device, plan_name, ) -def _make_bprop_cache_key( +def make_bprop_cache_key( total, N, H, @@ -185,7 +182,7 @@ def _make_bprop_cache_key( cu_dtype, has_initial_state, has_d_final_state, - ckpt_rows, + checkpoint_rows, scale, use_qk_l2norm, batch_invariant, @@ -212,7 +209,7 @@ def _make_bprop_cache_key( cu_dtype, bool(has_initial_state), bool(has_d_final_state), - ckpt_rows, + checkpoint_rows, float(scale), bool(use_qk_l2norm), bool(batch_invariant), @@ -228,7 +225,7 @@ def _make_bprop_cache_key( # --------------------------------------------------------------------------- -def _build_fprop_graph( +def build_fprop_graph( total, N, H, @@ -245,7 +242,7 @@ def _build_fprop_graph( output_final_state, use_qk_l2norm, batch_invariant, - ckpt, + checkpoint, use_beta_sigmoid=False, safe_gate=False, ): @@ -281,7 +278,7 @@ def _build_fprop_graph( batch_invariant=batch_invariant, use_beta_sigmoid=use_beta_sigmoid or None, safe_gate=safe_gate or None, - checkpoint_every_n_tokens=ckpt, + checkpoint_every_n_tokens=checkpoint, name="gdn", ) return graph, dict( @@ -306,7 +303,7 @@ def _build_fprop_graph( @torch.library.custom_op("cudnn::gated_delta_net_fwd", mutates_args=()) -def _gdn_fwd( +def gdn_fwd( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, @@ -341,23 +338,23 @@ def _gdn_fwd( if cu_seqlens.dtype not in (torch.int32, torch.int64): raise ValueError(f"gated_delta_net: cu_seqlens must be int32 or int64; got {cu_seqlens.dtype}") cu = cu_seqlens - _check_dtype("g", g, torch.float32) + check_dtype("g", g, torch.float32) if use_beta_sigmoid_in_kernel: - _check_dtype("beta", beta, q.dtype) + check_dtype("beta", beta, q.dtype) else: - _check_dtype("beta", beta, torch.float32) + check_dtype("beta", beta, torch.float32) if safe_gate: if a_log is None or dt_bias is None: raise ValueError("gated_delta_net: safe_gate requires a_log and dt_bias") - _check_dtype("a_log", a_log, torch.float32) - _check_dtype("dt_bias", dt_bias, torch.float32) + check_dtype("a_log", a_log, torch.float32) + check_dtype("dt_bias", dt_bias, torch.float32) elif a_log is not None or dt_bias is not None: raise ValueError("gated_delta_net: a_log/dt_bias require safe_gate=True") if initial_state is not None: - _check_dtype("initial_state", initial_state, torch.float32) + check_dtype("initial_state", initial_state, torch.float32) if initial_state.shape[0] != N: raise ValueError(f"initial_state must carry one state per sequence: got {initial_state.shape[0]} for {N} sequences") - for _name, _t in ( + for tensor_name, tensor in ( ("k", k), ("v", v), ("g", g), @@ -367,14 +364,14 @@ def _gdn_fwd( ("a_log", a_log), ("dt_bias", dt_bias), ): - if _t is not None and _t.device != device: - raise ValueError(f"gated_delta_net: {_name} must be on q's device ({device}); got {_t.device}") + if tensor is not None and tensor.device != device: + raise ValueError(f"gated_delta_net: {tensor_name} must be on q's device ({device}); got {tensor.device}") g32 = g beta32 = beta state0 = initial_state if initial_state is not None else None - ckpt = int(checkpoint_every_n_tokens) + checkpoint = int(checkpoint_every_n_tokens) - cache_key = _make_fprop_cache_key( + cache_key = make_fprop_cache_key( total, N, H, @@ -395,12 +392,12 @@ def _gdn_fwd( use_beta_sigmoid_in_kernel, safe_gate, state0 is not None, - ckpt, + checkpoint, device, plan_name, ) - if cache_key not in _fprop_cache: - _fprop_cache[cache_key] = _build_fprop_graph( + if cache_key not in fprop_cache: + fprop_cache[cache_key] = build_fprop_graph( total, N, H, @@ -408,22 +405,22 @@ def _gdn_fwd( HV, K, V, - _torch_dtype_to_cudnn(q.dtype), + torch_dtype_to_cudnn(q.dtype), cudnn.data_type.FLOAT, - _torch_dtype_to_cudnn(beta.dtype), + torch_dtype_to_cudnn(beta.dtype), cudnn.data_type.FLOAT if state0 is not None else None, - _torch_dtype_to_cudnn(cu_seqlens.dtype), + torch_dtype_to_cudnn(cu_seqlens.dtype), float(scale), bool(output_final_state), bool(use_qk_l2norm_in_kernel), bool(batch_invariant), - ckpt, + checkpoint, use_beta_sigmoid=bool(use_beta_sigmoid_in_kernel), safe_gate=bool(safe_gate), ) - select_plan(_fprop_cache[cache_key][0], plan_name) + select_plan(fprop_cache[cache_key][0], plan_name) - graph, t = _fprop_cache[cache_key] + graph, t = fprop_cache[cache_key] HO = max(H, HV) o = torch.empty(total, HO, V, dtype=q.dtype, device=device) @@ -446,16 +443,16 @@ def _gdn_fwd( final_state = torch.empty(N, HO, V, K, dtype=torch.float32, device=device) variant_pack[t["fs"]] = final_state state_checkpoints = torch.empty(0, dtype=q.dtype, device=device) - if ckpt > 0: - total_checkpoints = max(total // ckpt + N, 1) + if checkpoint > 0: + total_checkpoints = max(total // checkpoint + N, 1) state_checkpoints = torch.empty(total_checkpoints, HO, V, K, dtype=q.dtype, device=device) variant_pack[t["state_checkpoints"]] = state_checkpoints - graph.execute(variant_pack, workspace=_graph_workspace(graph, device), handle=_get_handle(device)) + graph.execute(variant_pack, workspace=graph_workspace(graph, device), handle=get_handle(device)) return o, final_state, state_checkpoints -@_gdn_fwd.register_fake -def _gdn_fwd_fake( +@gdn_fwd.register_fake +def gdn_fwd_fake( q, k, v, @@ -500,7 +497,7 @@ def _gdn_fwd_fake( # --------------------------------------------------------------------------- -def _build_bprop_graph( +def build_bprop_graph( total, N, H, @@ -514,7 +511,7 @@ def _build_bprop_graph( state_dtype, dstate_in_dtype, cu_dtype, - ckpt_rows, + checkpoint_rows, scale, use_qk_l2norm, batch_invariant, @@ -536,9 +533,9 @@ def _build_bprop_graph( dfs_t = None if dstate_in_dtype is not None: dfs_t = graph.tensor([N, HO, V, K], data_type=dstate_in_dtype, name="d_final_state") - ckpts_t = None - if ckpt_rows is not None: - ckpts_t = graph.tensor([ckpt_rows, HO, V, K], data_type=io_dtype, name="state_checkpoints") + checkpoints_t = None + if checkpoint_rows is not None: + checkpoints_t = graph.tensor([checkpoint_rows, HO, V, K], data_type=io_dtype, name="state_checkpoints") a_log_t = None dt_bias_t = None if safe_gate: @@ -552,7 +549,7 @@ def _build_bprop_graph( beta=beta_t, cu_seqlens=cu_t, dO=dO_t, - state_checkpoints=ckpts_t, + state_checkpoints=checkpoints_t, initial_state=state0_t, d_final_state=dfs_t, a_log=a_log_t, @@ -584,7 +581,7 @@ def _build_bprop_graph( dstate0=dstate0_t, d_a_log=dA_t, d_dt_bias=dDt_t, - ckpts=ckpts_t, + checkpoints=checkpoints_t, ) @@ -594,7 +591,7 @@ def _build_bprop_graph( @torch.library.custom_op("cudnn::gated_delta_net_bwd", mutates_args=()) -def _gdn_bwd( +def gdn_bwd( dO: torch.Tensor, q: torch.Tensor, k: torch.Tensor, @@ -627,8 +624,6 @@ def _gdn_bwd( logits and ``dbeta`` is the raw-logit gradient. """ total, H, K = q.shape - # autograd materializes reduction grads as broadcast (stride-0) - # views; densify ONLY those (dense callers pass through untouched) if 0 in dO.stride(): dO = dO.contiguous() if d_final_state is not None and 0 in d_final_state.stride(): @@ -642,25 +637,25 @@ def _gdn_bwd( if cu_seqlens.dtype not in (torch.int32, torch.int64): raise ValueError(f"gated_delta_net: cu_seqlens must be int32 or int64; got {cu_seqlens.dtype}") cu = cu_seqlens - _check_dtype("g", g, torch.float32) + check_dtype("g", g, torch.float32) if use_beta_sigmoid_in_kernel: - _check_dtype("beta", beta, q.dtype) + check_dtype("beta", beta, q.dtype) else: - _check_dtype("beta", beta, torch.float32) + check_dtype("beta", beta, torch.float32) if safe_gate: if a_log is None or dt_bias is None: raise ValueError("gated_delta_net: safe_gate requires a_log and dt_bias") - _check_dtype("a_log", a_log, torch.float32) - _check_dtype("dt_bias", dt_bias, torch.float32) + check_dtype("a_log", a_log, torch.float32) + check_dtype("dt_bias", dt_bias, torch.float32) elif a_log is not None or dt_bias is not None: raise ValueError("gated_delta_net: a_log/dt_bias require safe_gate=True") if initial_state is not None: - _check_dtype("initial_state", initial_state, torch.float32) + check_dtype("initial_state", initial_state, torch.float32) if initial_state.shape[0] != N: raise ValueError(f"initial_state must carry one state per sequence: got {initial_state.shape[0]} for {N} sequences") if state_checkpoints is not None: - _check_dtype("state_checkpoints", state_checkpoints, q.dtype) - for _name, _t in ( + check_dtype("state_checkpoints", state_checkpoints, q.dtype) + for tensor_name, tensor in ( ("k", k), ("v", v), ("g", g), @@ -672,16 +667,16 @@ def _gdn_bwd( ("a_log", a_log), ("dt_bias", dt_bias), ): - if _t is not None and _t.device != device: - raise ValueError(f"gated_delta_net: {_name} must be on q's device ({device}); got {_t.device}") + if tensor is not None and tensor.device != device: + raise ValueError(f"gated_delta_net: {tensor_name} must be on q's device ({device}); got {tensor.device}") g32 = g beta32 = beta state0 = initial_state if initial_state is not None else None if d_final_state is not None: - _check_dtype("d_final_state", d_final_state, torch.float32) + check_dtype("d_final_state", d_final_state, torch.float32) dstate_in = d_final_state if d_final_state is not None else None - cache_key = _make_bprop_cache_key( + cache_key = make_bprop_cache_key( total, N, H, @@ -707,8 +702,8 @@ def _gdn_bwd( device, plan_name, ) - if cache_key not in _bprop_cache: - _bprop_cache[cache_key] = _build_bprop_graph( + if cache_key not in bprop_cache: + bprop_cache[cache_key] = build_bprop_graph( total, N, H, @@ -716,12 +711,12 @@ def _gdn_bwd( HV, K, V, - _torch_dtype_to_cudnn(q.dtype), + torch_dtype_to_cudnn(q.dtype), cudnn.data_type.FLOAT, - _torch_dtype_to_cudnn(beta.dtype), + torch_dtype_to_cudnn(beta.dtype), cudnn.data_type.FLOAT if state0 is not None else None, cudnn.data_type.FLOAT if dstate_in is not None else None, - _torch_dtype_to_cudnn(cu_seqlens.dtype), + torch_dtype_to_cudnn(cu_seqlens.dtype), state_checkpoints.shape[0] if state_checkpoints is not None else None, float(scale), bool(use_qk_l2norm_in_kernel), @@ -729,9 +724,9 @@ def _gdn_bwd( use_beta_sigmoid=bool(use_beta_sigmoid_in_kernel), safe_gate=bool(safe_gate), ) - select_plan(_bprop_cache[cache_key][0], plan_name) + select_plan(bprop_cache[cache_key][0], plan_name) - graph, t = _bprop_cache[cache_key] + graph, t = bprop_cache[cache_key] HO = max(H, HV) dq = torch.empty(total, H, K, dtype=q.dtype, device=device) @@ -761,7 +756,7 @@ def _gdn_bwd( if dstate_in is not None: variant_pack[t["dfs"]] = dstate_in if state_checkpoints is not None: - variant_pack[t["ckpts"]] = state_checkpoints + variant_pack[t["checkpoints"]] = state_checkpoints d_a_log = torch.empty(0, dtype=torch.float32, device=device) d_dt_bias = torch.empty(0, dtype=torch.float32, device=device) if safe_gate: @@ -771,14 +766,14 @@ def _gdn_bwd( d_dt_bias = torch.empty(HO, dtype=torch.float32, device=device) variant_pack[t["d_a_log"]] = d_a_log variant_pack[t["d_dt_bias"]] = d_dt_bias - graph.execute(variant_pack, workspace=_graph_workspace(graph, device), handle=_get_handle(device)) + graph.execute(variant_pack, workspace=graph_workspace(graph, device), handle=get_handle(device)) if dstate0 is None: dstate0 = torch.empty(0, dtype=torch.float32, device=device) return dq, dk, dv, dg32, dbeta, dstate0, d_a_log, d_dt_bias -@_gdn_bwd.register_fake -def _gdn_bwd_fake( +@gdn_bwd.register_fake +def gdn_bwd_fake( dO, q, k, @@ -820,7 +815,7 @@ def _gdn_bwd_fake( # --------------------------------------------------------------------------- -def _gdn_setup_context(ctx, inputs, output): +def gdn_setup_context(ctx, inputs, output): ( q, k, @@ -840,11 +835,9 @@ def _gdn_setup_context(ctx, inputs, output): checkpoint_every_n_tokens, plan_name, ) = inputs - # save_for_backward cannot hold None; keep initial_state as an attribute. - # g/beta are saved as passed: raw logits under safe_gate / use_beta_sigmoid. saved = [q, k, v, g, beta, cu_seqlens] - ctx.ckpt_reuse = checkpoint_every_n_tokens == 64 and output[2].numel() > 0 - if ctx.ckpt_reuse: + ctx.checkpoint_reuse = checkpoint_every_n_tokens == 64 and output[2].numel() > 0 + if ctx.checkpoint_reuse: saved.append(output[2]) if safe_gate: saved.extend([a_log, dt_bias]) @@ -860,11 +853,11 @@ def _gdn_setup_context(ctx, inputs, output): ctx.mark_non_differentiable(output[2]) -def _gdn_backward(ctx, dO, dFinal, _dstate_checkpoints): +def gdn_backward(ctx, dO, dFinal, dstate_checkpoints): a_log = dt_bias = None if ctx.safe_gate: a_log, dt_bias = ctx.saved_tensors[-2:] - if ctx.ckpt_reuse: + if ctx.checkpoint_reuse: q, k, v, g, beta, cu_seqlens, state_checkpoints = ctx.saved_tensors[:7] else: q, k, v, g, beta, cu_seqlens = ctx.saved_tensors[:6] @@ -894,9 +887,6 @@ def _gdn_backward(ctx, dO, dFinal, _dstate_checkpoints): dt_bias=dt_bias, plan_name=ctx.plan_name, ) - # q, k, v, g, beta, cu_seqlens, scale, initial_state, output_final_state, - # use_qk_l2norm_in_kernel, batch_invariant, use_beta_sigmoid_in_kernel, - # safe_gate, a_log, dt_bias, checkpoint_every_n_tokens, plan_name return ( dq, dk, @@ -920,8 +910,8 @@ def _gdn_backward(ctx, dO, dFinal, _dstate_checkpoints): torch.library.register_autograd( "cudnn::gated_delta_net_fwd", - _gdn_backward, - setup_context=_gdn_setup_context, + gdn_backward, + setup_context=gdn_setup_context, ) diff --git a/python/cudnn/linear_attention/ops/gdn2.py b/python/cudnn/linear_attention/ops/gdn2.py index 6b36d3c10..0c37e45c5 100644 --- a/python/cudnn/linear_attention/ops/gdn2.py +++ b/python/cudnn/linear_attention/ops/gdn2.py @@ -40,7 +40,7 @@ # --------------------------------------------------------------------------- -_TORCH_DTYPE_TO_CUDNN = { +TORCH_DTYPE_TO_CUDNN = { torch.float16: cudnn.data_type.HALF, torch.bfloat16: cudnn.data_type.BFLOAT16, torch.float32: cudnn.data_type.FLOAT, @@ -48,10 +48,9 @@ torch.int64: cudnn.data_type.INT64, } -# one graph per static configuration (shapes, dtypes, scale, flags, device) -_fprop_cache: Dict[tuple, tuple] = {} -_bprop_cache: Dict[tuple, tuple] = {} -_cudnn_handles: Dict[int, int] = {} +fprop_cache: Dict[tuple, tuple] = {} +bprop_cache: Dict[tuple, tuple] = {} +cudnn_handles: Dict[int, int] = {} # --------------------------------------------------------------------------- @@ -73,11 +72,9 @@ def select_plan(graph, plan_name): graph.check_support() -def _graph_workspace(graph, device): +def graph_workspace(graph, device): """Caller-side workspace for a compiled graph.""" if not graph._is_built: - # mirror execute()'s auto-build: plan first (a bare build() would - # lower GDN2 to the backend, which has no lowering) if not graph._planning_done: graph.create_execution_plans() if graph.selected_engine is None: @@ -85,36 +82,36 @@ def _graph_workspace(graph, device): else: graph.build_plans() size = graph.get_workspace_size() - ws = getattr(graph, "_la_ops_workspace", None) - if ws is None or ws.numel() < size or ws.device != device: - ws = torch.empty(max(size, 1), dtype=torch.uint8, device=device) - graph._la_ops_workspace = ws - return ws + workspace = getattr(graph, "la_ops_workspace", None) + if workspace is None or workspace.numel() < size or workspace.device != device: + workspace = torch.empty(max(size, 1), dtype=torch.uint8, device=device) + graph.la_ops_workspace = workspace + return workspace -def _get_handle(device): +def get_handle(device): """Per-device cuDNN handle carrying the caller's current stream.""" idx = device.index if device.index is not None else torch.cuda.current_device() - handle = _cudnn_handles.get(idx) + handle = cudnn_handles.get(idx) if handle is None: with torch.cuda.device(idx): handle = cudnn.create_handle() - _cudnn_handles[idx] = handle + cudnn_handles[idx] = handle cudnn.set_stream(handle=handle, stream=torch.cuda.current_stream(device).cuda_stream) return handle -def _torch_dtype_to_cudnn(dtype: torch.dtype): +def torch_dtype_to_cudnn(dtype: torch.dtype): """Map a PyTorch dtype to a cuDNN data_type enum.""" - return _TORCH_DTYPE_TO_CUDNN[dtype] + return TORCH_DTYPE_TO_CUDNN[dtype] -def _check_dtype(name, t, want) -> None: +def check_dtype(name, t, want) -> None: if t.dtype != want: raise TypeError(f"gated_delta_net_v2: {name} must be {want} (kernel-native; callers convert), got {t.dtype}") -def _make_fprop_cache_key( +def make_fprop_cache_key( total, N, H, @@ -136,7 +133,7 @@ def _make_fprop_cache_key( safe_gate, gate_lower_bound, has_initial_state, - ckpt, + checkpoint, device, plan_name, ): @@ -163,13 +160,13 @@ def _make_fprop_cache_key( bool(safe_gate), float(gate_lower_bound) if gate_lower_bound is not None else None, bool(has_initial_state), - ckpt, + checkpoint, device, plan_name, ) -def _make_bprop_cache_key( +def make_bprop_cache_key( total, N, H, @@ -186,7 +183,7 @@ def _make_bprop_cache_key( cu_dtype, state_dtype, dstate_in_dtype, - ckpt_rows, + checkpoint_rows, scale, use_qk_l2norm, batch_invariant, @@ -214,7 +211,7 @@ def _make_bprop_cache_key( cu_dtype, state_dtype, dstate_in_dtype, - ckpt_rows, + checkpoint_rows, float(scale), bool(use_qk_l2norm), bool(batch_invariant), @@ -231,7 +228,7 @@ def _make_bprop_cache_key( # --------------------------------------------------------------------------- -def _build_fprop_graph( +def build_fprop_graph( total, N, H, @@ -250,7 +247,7 @@ def _build_fprop_graph( batch_invariant, safe_gate, gate_lower_bound, - ckpt, + checkpoint, use_beta_sigmoid=False, ): graph = cudnn.pygraph() @@ -288,7 +285,7 @@ def _build_fprop_graph( use_beta_sigmoid=use_beta_sigmoid or None, safe_gate=safe_gate, gate_lower_bound=gate_lower_bound, - checkpoint_every_n_tokens=ckpt, + checkpoint_every_n_tokens=checkpoint, name="gdn2", ) return graph, dict( @@ -314,7 +311,7 @@ def _build_fprop_graph( @torch.library.custom_op("cudnn::gated_delta_net_v2_fwd", mutates_args=()) -def _gdn2_fwd( +def gdn2_fwd( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, @@ -352,24 +349,24 @@ def _gdn2_fwd( if cu_seqlens.dtype not in (torch.int32, torch.int64): raise ValueError(f"gated_delta_net_v2: cu_seqlens must be int32 or int64; got {cu_seqlens.dtype}") cu = cu_seqlens - _check_dtype("g", g, torch.float32) - _check_dtype("beta", beta, q.dtype) - _check_dtype("w", w, q.dtype) + check_dtype("g", g, torch.float32) + check_dtype("beta", beta, q.dtype) + check_dtype("w", w, q.dtype) if safe_gate: if a_log is None or dt_bias is None: raise ValueError("gated_delta_net_v2: safe_gate requires a_log and dt_bias") - _check_dtype("a_log", a_log, torch.float32) - _check_dtype("dt_bias", dt_bias, torch.float32) + check_dtype("a_log", a_log, torch.float32) + check_dtype("dt_bias", dt_bias, torch.float32) elif a_log is not None or dt_bias is not None: raise ValueError("gated_delta_net_v2: a_log/dt_bias require safe_gate=True") if initial_state is not None: - _check_dtype("initial_state", initial_state, torch.float32) + check_dtype("initial_state", initial_state, torch.float32) if initial_state.shape[0] != N: raise ValueError(f"initial_state must carry one state per sequence: got {initial_state.shape[0]} for {N} sequences") g32 = g beta_io = beta w_io = w - for _name, _t in ( + for tensor_name, tensor in ( ("k", k), ("v", v), ("g", g), @@ -380,12 +377,12 @@ def _gdn2_fwd( ("a_log", a_log), ("dt_bias", dt_bias), ): - if _t is not None and _t.device != device: - raise ValueError(f"gated_delta_net_v2: {_name} must be on q's device ({device}); got {_t.device}") + if tensor is not None and tensor.device != device: + raise ValueError(f"gated_delta_net_v2: {tensor_name} must be on q's device ({device}); got {tensor.device}") state0 = initial_state if initial_state is not None else None - ckpt = int(checkpoint_every_n_tokens) + checkpoint = int(checkpoint_every_n_tokens) - cache_key = _make_fprop_cache_key( + cache_key = make_fprop_cache_key( total, N, H, @@ -407,12 +404,12 @@ def _gdn2_fwd( safe_gate, gate_lower_bound, state0 is not None, - ckpt, + checkpoint, device, plan_name, ) - if cache_key not in _fprop_cache: - _fprop_cache[cache_key] = _build_fprop_graph( + if cache_key not in fprop_cache: + fprop_cache[cache_key] = build_fprop_graph( total, N, H, @@ -420,23 +417,23 @@ def _gdn2_fwd( HV, K, V, - _torch_dtype_to_cudnn(q.dtype), + torch_dtype_to_cudnn(q.dtype), cudnn.data_type.FLOAT, - _torch_dtype_to_cudnn(q.dtype), + torch_dtype_to_cudnn(q.dtype), cudnn.data_type.FLOAT if state0 is not None else None, - _torch_dtype_to_cudnn(cu_seqlens.dtype), + torch_dtype_to_cudnn(cu_seqlens.dtype), float(scale), bool(output_final_state), bool(use_qk_l2norm_in_kernel), bool(batch_invariant), bool(safe_gate), float(gate_lower_bound) if gate_lower_bound is not None else None, - ckpt, + checkpoint, use_beta_sigmoid=bool(use_beta_sigmoid_in_kernel), ) - select_plan(_fprop_cache[cache_key][0], plan_name) + select_plan(fprop_cache[cache_key][0], plan_name) - graph, t = _fprop_cache[cache_key] + graph, t = fprop_cache[cache_key] o = torch.empty(total, HO, V, dtype=q.dtype, device=device) variant_pack = { @@ -459,16 +456,16 @@ def _gdn2_fwd( final_state = torch.empty(N, HO, V, K, dtype=torch.float32, device=device) variant_pack[t["fs"]] = final_state state_checkpoints = torch.empty(0, dtype=q.dtype, device=device) - if ckpt > 0: - total_checkpoints = max(total // ckpt + N, 1) + if checkpoint > 0: + total_checkpoints = max(total // checkpoint + N, 1) state_checkpoints = torch.empty(total_checkpoints, HO, V, K, dtype=q.dtype, device=device) variant_pack[t["state_checkpoints"]] = state_checkpoints - graph.execute(variant_pack, workspace=_graph_workspace(graph, device), handle=_get_handle(device)) + graph.execute(variant_pack, workspace=graph_workspace(graph, device), handle=get_handle(device)) return o, final_state, state_checkpoints -@_gdn2_fwd.register_fake -def _gdn2_fwd_fake( +@gdn2_fwd.register_fake +def gdn2_fwd_fake( q, k, v, @@ -515,7 +512,7 @@ def _gdn2_fwd_fake( # --------------------------------------------------------------------------- -def _build_bprop_graph( +def build_bprop_graph( total, N, H, @@ -529,7 +526,7 @@ def _build_bprop_graph( state_dtype, dstate_in_dtype, cu_dtype, - ckpt_rows, + checkpoint_rows, scale, use_qk_l2norm, batch_invariant, @@ -553,9 +550,9 @@ def _build_bprop_graph( dfs_t = None if dstate_in_dtype is not None: dfs_t = graph.tensor([N, HO, V, K], data_type=dstate_in_dtype, name="d_final_state") - ckpts_t = None - if ckpt_rows is not None: - ckpts_t = graph.tensor([ckpt_rows, HO, V, K], data_type=io_dtype, name="state_checkpoints") + checkpoints_t = None + if checkpoint_rows is not None: + checkpoints_t = graph.tensor([checkpoint_rows, HO, V, K], data_type=io_dtype, name="state_checkpoints") a_log_t = None dt_bias_t = None if safe_gate: @@ -570,7 +567,7 @@ def _build_bprop_graph( w=w_t, cu_seqlens=cu_t, dO=dO_t, - state_checkpoints=ckpts_t, + state_checkpoints=checkpoints_t, initial_state=state0_t, d_final_state=dfs_t, a_log=a_log_t, @@ -605,7 +602,7 @@ def _build_bprop_graph( dstate0=dstate0_t, d_a_log=dA_t, d_dt_bias=dDt_t, - ckpts=ckpts_t, + checkpoints=checkpoints_t, ) @@ -615,7 +612,7 @@ def _build_bprop_graph( @torch.library.custom_op("cudnn::gated_delta_net_v2_bwd", mutates_args=()) -def _gdn2_bwd( +def gdn2_bwd( dO: torch.Tensor, q: torch.Tensor, k: torch.Tensor, @@ -650,8 +647,6 @@ def _gdn2_bwd( logits and ``dbeta`` is the raw-logit gradient. """ total, H, K = q.shape - # autograd materializes reduction grads as broadcast (stride-0) - # views; densify ONLY those (dense callers pass through untouched) if 0 in dO.stride(): dO = dO.contiguous() if d_final_state is not None and 0 in d_final_state.stride(): @@ -666,25 +661,25 @@ def _gdn2_bwd( if cu_seqlens.dtype not in (torch.int32, torch.int64): raise ValueError(f"gated_delta_net_v2: cu_seqlens must be int32 or int64; got {cu_seqlens.dtype}") cu = cu_seqlens - _check_dtype("g", g, torch.float32) - _check_dtype("beta", beta, q.dtype) - _check_dtype("w", w, q.dtype) + check_dtype("g", g, torch.float32) + check_dtype("beta", beta, q.dtype) + check_dtype("w", w, q.dtype) if safe_gate: if a_log is None or dt_bias is None: raise ValueError("gated_delta_net_v2: safe_gate requires a_log and dt_bias") - _check_dtype("a_log", a_log, torch.float32) - _check_dtype("dt_bias", dt_bias, torch.float32) + check_dtype("a_log", a_log, torch.float32) + check_dtype("dt_bias", dt_bias, torch.float32) elif a_log is not None or dt_bias is not None: raise ValueError("gated_delta_net_v2: a_log/dt_bias require safe_gate=True") if initial_state is not None: - _check_dtype("initial_state", initial_state, torch.float32) + check_dtype("initial_state", initial_state, torch.float32) if initial_state.shape[0] != N: raise ValueError(f"initial_state must carry one state per sequence: got {initial_state.shape[0]} for {N} sequences") if d_final_state is not None: - _check_dtype("d_final_state", d_final_state, torch.float32) + check_dtype("d_final_state", d_final_state, torch.float32) if state_checkpoints is not None: - _check_dtype("state_checkpoints", state_checkpoints, q.dtype) - for _name, _t in ( + check_dtype("state_checkpoints", state_checkpoints, q.dtype) + for tensor_name, tensor in ( ("k", k), ("v", v), ("g", g), @@ -697,12 +692,12 @@ def _gdn2_bwd( ("a_log", a_log), ("dt_bias", dt_bias), ): - if _t is not None and _t.device != device: - raise ValueError(f"gated_delta_net_v2: {_name} must be on q's device ({device}); got {_t.device}") + if tensor is not None and tensor.device != device: + raise ValueError(f"gated_delta_net_v2: {tensor_name} must be on q's device ({device}); got {tensor.device}") state0 = initial_state if initial_state is not None else None dstate_in = d_final_state if d_final_state is not None else None - cache_key = _make_bprop_cache_key( + cache_key = make_bprop_cache_key( total, N, H, @@ -729,8 +724,8 @@ def _gdn2_bwd( device, plan_name, ) - if cache_key not in _bprop_cache: - _bprop_cache[cache_key] = _build_bprop_graph( + if cache_key not in bprop_cache: + bprop_cache[cache_key] = build_bprop_graph( total, N, H, @@ -738,12 +733,12 @@ def _gdn2_bwd( HV, K, V, - _torch_dtype_to_cudnn(q.dtype), + torch_dtype_to_cudnn(q.dtype), cudnn.data_type.FLOAT, - _torch_dtype_to_cudnn(q.dtype), - _torch_dtype_to_cudnn(state0.dtype) if state0 is not None else None, - _torch_dtype_to_cudnn(dstate_in.dtype) if dstate_in is not None else None, - _torch_dtype_to_cudnn(cu_seqlens.dtype), + torch_dtype_to_cudnn(q.dtype), + torch_dtype_to_cudnn(state0.dtype) if state0 is not None else None, + torch_dtype_to_cudnn(dstate_in.dtype) if dstate_in is not None else None, + torch_dtype_to_cudnn(cu_seqlens.dtype), state_checkpoints.shape[0] if state_checkpoints is not None else None, float(scale), bool(use_qk_l2norm_in_kernel), @@ -752,9 +747,9 @@ def _gdn2_bwd( safe_gate=bool(safe_gate), gate_lower_bound=float(gate_lower_bound) if gate_lower_bound is not None else None, ) - select_plan(_bprop_cache[cache_key][0], plan_name) + select_plan(bprop_cache[cache_key][0], plan_name) - graph, t = _bprop_cache[cache_key] + graph, t = bprop_cache[cache_key] dq = torch.empty(total, H, K, dtype=q.dtype, device=device) dk = torch.empty(total, HK, K, dtype=q.dtype, device=device) @@ -786,7 +781,7 @@ def _gdn2_bwd( if dstate_in is not None: variant_pack[t["dfs"]] = dstate_in if state_checkpoints is not None: - variant_pack[t["ckpts"]] = state_checkpoints + variant_pack[t["checkpoints"]] = state_checkpoints d_a_log = torch.empty(0, dtype=torch.float32, device=device) d_dt_bias = torch.empty(0, dtype=torch.float32, device=device) if safe_gate: @@ -796,14 +791,14 @@ def _gdn2_bwd( d_dt_bias = torch.empty(HO, K, dtype=torch.float32, device=device) variant_pack[t["d_a_log"]] = d_a_log variant_pack[t["d_dt_bias"]] = d_dt_bias - graph.execute(variant_pack, workspace=_graph_workspace(graph, device), handle=_get_handle(device)) + graph.execute(variant_pack, workspace=graph_workspace(graph, device), handle=get_handle(device)) if dstate0 is None: dstate0 = torch.empty(0, dtype=torch.float32, device=device) return dq, dk, dv, dg, dbeta, dw, dstate0, d_a_log, d_dt_bias -@_gdn2_bwd.register_fake -def _gdn2_bwd_fake( +@gdn2_bwd.register_fake +def gdn2_bwd_fake( dO, q, k, @@ -848,7 +843,7 @@ def _gdn2_bwd_fake( # --------------------------------------------------------------------------- -def _gdn2_setup_context(ctx, inputs, output): +def gdn2_setup_context(ctx, inputs, output): ( q, k, @@ -870,11 +865,9 @@ def _gdn2_setup_context(ctx, inputs, output): checkpoint_every_n_tokens, plan_name, ) = inputs - # save_for_backward cannot hold None; keep initial_state as an attribute. - # g/beta are saved as passed: raw logits under safe_gate / use_beta_sigmoid. saved = [q, k, v, g, beta, w, cu_seqlens] - ctx.ckpt_reuse = checkpoint_every_n_tokens == 16 and output[2].numel() > 0 - if ctx.ckpt_reuse: + ctx.checkpoint_reuse = checkpoint_every_n_tokens == 16 and output[2].numel() > 0 + if ctx.checkpoint_reuse: saved.append(output[2]) if safe_gate: saved.extend([a_log, dt_bias]) @@ -891,11 +884,11 @@ def _gdn2_setup_context(ctx, inputs, output): ctx.mark_non_differentiable(output[2]) -def _gdn2_backward(ctx, dO, dFinal, _dstate_checkpoints): +def gdn2_backward(ctx, dO, dFinal, dstate_checkpoints): a_log = dt_bias = None if ctx.safe_gate: a_log, dt_bias = ctx.saved_tensors[-2:] - if ctx.ckpt_reuse: + if ctx.checkpoint_reuse: q, k, v, g, beta, w, cu_seqlens, state_checkpoints = ctx.saved_tensors[:8] else: q, k, v, g, beta, w, cu_seqlens = ctx.saved_tensors[:7] @@ -927,10 +920,6 @@ def _gdn2_backward(ctx, dO, dFinal, _dstate_checkpoints): dt_bias=dt_bias, plan_name=ctx.plan_name, ) - # q, k, v, g, beta, w, cu_seqlens, scale, initial_state, - # output_final_state, use_qk_l2norm_in_kernel, batch_invariant, - # use_beta_sigmoid_in_kernel, safe_gate, gate_lower_bound, a_log, - # dt_bias, checkpoint_every_n_tokens, plan_name return ( dq, dk, @@ -956,8 +945,8 @@ def _gdn2_backward(ctx, dO, dFinal, _dstate_checkpoints): torch.library.register_autograd( "cudnn::gated_delta_net_v2_fwd", - _gdn2_backward, - setup_context=_gdn2_setup_context, + gdn2_backward, + setup_context=gdn2_setup_context, ) diff --git a/python/cudnn/linear_attention/ops/kda.py b/python/cudnn/linear_attention/ops/kda.py index 9be778f7b..3ef479c36 100644 --- a/python/cudnn/linear_attention/ops/kda.py +++ b/python/cudnn/linear_attention/ops/kda.py @@ -43,7 +43,7 @@ # --------------------------------------------------------------------------- -_TORCH_DTYPE_TO_CUDNN = { +TORCH_DTYPE_TO_CUDNN = { torch.float16: cudnn.data_type.HALF, torch.bfloat16: cudnn.data_type.BFLOAT16, torch.float32: cudnn.data_type.FLOAT, @@ -51,10 +51,9 @@ torch.int64: cudnn.data_type.INT64, } -# one graph per static configuration (shapes, dtypes, scale, flags, device) -_fprop_cache: Dict[tuple, tuple] = {} -_bprop_cache: Dict[tuple, tuple] = {} -_cudnn_handles: Dict[int, int] = {} +fprop_cache: Dict[tuple, tuple] = {} +bprop_cache: Dict[tuple, tuple] = {} +cudnn_handles: Dict[int, int] = {} # --------------------------------------------------------------------------- @@ -76,11 +75,9 @@ def select_plan(graph, plan_name): graph.check_support() -def _graph_workspace(graph, device): +def graph_workspace(graph, device): """Caller-side workspace for a compiled graph.""" if not graph._is_built: - # plan first (a bare build() would lower KDA to the backend, - # which has no lowering) if not graph._planning_done: graph.create_execution_plans() if graph.selected_engine is None: @@ -88,36 +85,36 @@ def _graph_workspace(graph, device): else: graph.build_plans() size = graph.get_workspace_size() - ws = getattr(graph, "_la_ops_workspace", None) - if ws is None or ws.numel() < size or ws.device != device: - ws = torch.empty(max(size, 1), dtype=torch.uint8, device=device) - graph._la_ops_workspace = ws - return ws + workspace = getattr(graph, "la_ops_workspace", None) + if workspace is None or workspace.numel() < size or workspace.device != device: + workspace = torch.empty(max(size, 1), dtype=torch.uint8, device=device) + graph.la_ops_workspace = workspace + return workspace -def _get_handle(device): +def get_handle(device): """Per-device cuDNN handle carrying the caller's current stream.""" idx = device.index if device.index is not None else torch.cuda.current_device() - handle = _cudnn_handles.get(idx) + handle = cudnn_handles.get(idx) if handle is None: with torch.cuda.device(idx): handle = cudnn.create_handle() - _cudnn_handles[idx] = handle + cudnn_handles[idx] = handle cudnn.set_stream(handle=handle, stream=torch.cuda.current_stream(device).cuda_stream) return handle -def _torch_dtype_to_cudnn(dtype: torch.dtype): +def torch_dtype_to_cudnn(dtype: torch.dtype): """Map a PyTorch dtype to a cuDNN data_type enum.""" - return _TORCH_DTYPE_TO_CUDNN[dtype] + return TORCH_DTYPE_TO_CUDNN[dtype] -def _check_dtype(name, t, want) -> None: +def check_dtype(name, t, want) -> None: if t.dtype != want: raise TypeError(f"kimi_delta_attention: {name} must be {want} (kernel-native; callers convert), got {t.dtype}") -def _make_fprop_cache_key( +def make_fprop_cache_key( total, N, H, @@ -139,7 +136,7 @@ def _make_fprop_cache_key( safe_gate, gate_lower_bound, has_initial_state, - ckpt, + checkpoint, device, plan_name, ): @@ -166,13 +163,13 @@ def _make_fprop_cache_key( bool(safe_gate), float(gate_lower_bound) if gate_lower_bound is not None else None, bool(has_initial_state), - ckpt, + checkpoint, device, plan_name, ) -def _make_bprop_cache_key( +def make_bprop_cache_key( total, N, H, @@ -191,7 +188,7 @@ def _make_bprop_cache_key( beta_dtype, state_dtype, dstate_in_dtype, - ckpt_rows, + checkpoint_rows, scale, use_qk_l2norm, batch_invariant, @@ -221,7 +218,7 @@ def _make_bprop_cache_key( beta_dtype, state_dtype, dstate_in_dtype, - ckpt_rows, + checkpoint_rows, float(scale), bool(use_qk_l2norm), bool(batch_invariant), @@ -238,7 +235,7 @@ def _make_bprop_cache_key( # --------------------------------------------------------------------------- -def _build_fprop_graph( +def build_fprop_graph( total, N, H, @@ -258,7 +255,7 @@ def _build_fprop_graph( use_beta_sigmoid, safe_gate, gate_lower_bound, - ckpt, + checkpoint, ): graph = cudnn.pygraph() HO = max(H, HV) @@ -293,7 +290,7 @@ def _build_fprop_graph( use_beta_sigmoid=use_beta_sigmoid, safe_gate=safe_gate, gate_lower_bound=gate_lower_bound, - checkpoint_every_n_tokens=ckpt, + checkpoint_every_n_tokens=checkpoint, name="kda", ) return graph, dict( @@ -318,7 +315,7 @@ def _build_fprop_graph( @torch.library.custom_op("cudnn::kimi_delta_attention_fwd", mutates_args=()) -def _kda_fwd( +def kda_fwd( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, @@ -355,23 +352,23 @@ def _kda_fwd( if cu_seqlens.dtype not in (torch.int32, torch.int64): raise ValueError(f"kimi_delta_attention: cu_seqlens must be int32 or int64; got {cu_seqlens.dtype}") cu = cu_seqlens - _check_dtype("g", g, torch.float32) + check_dtype("g", g, torch.float32) if use_beta_sigmoid_in_kernel: - _check_dtype("beta", beta, q.dtype) + check_dtype("beta", beta, q.dtype) else: - _check_dtype("beta", beta, torch.float32) + check_dtype("beta", beta, torch.float32) if safe_gate: if a_log is None or dt_bias is None: raise ValueError("kimi_delta_attention: safe_gate requires a_log and dt_bias") - _check_dtype("a_log", a_log, torch.float32) - _check_dtype("dt_bias", dt_bias, torch.float32) + check_dtype("a_log", a_log, torch.float32) + check_dtype("dt_bias", dt_bias, torch.float32) elif a_log is not None or dt_bias is not None: raise ValueError("kimi_delta_attention: a_log/dt_bias require safe_gate=True") if initial_state is not None: - _check_dtype("initial_state", initial_state, torch.float32) + check_dtype("initial_state", initial_state, torch.float32) if initial_state.shape[0] != N: raise ValueError(f"initial_state must carry one state per sequence: got {initial_state.shape[0]} for {N} sequences") - for _name, _t in ( + for tensor_name, tensor in ( ("k", k), ("v", v), ("g", g), @@ -381,12 +378,12 @@ def _kda_fwd( ("a_log", a_log), ("dt_bias", dt_bias), ): - if _t is not None and _t.device != device: - raise ValueError(f"kimi_delta_attention: {_name} must be on q's device ({device}); got {_t.device}") + if tensor is not None and tensor.device != device: + raise ValueError(f"kimi_delta_attention: {tensor_name} must be on q's device ({device}); got {tensor.device}") state0 = initial_state if initial_state is not None else None - ckpt = int(checkpoint_every_n_tokens) + checkpoint = int(checkpoint_every_n_tokens) - cache_key = _make_fprop_cache_key( + cache_key = make_fprop_cache_key( total, N, H, @@ -408,12 +405,12 @@ def _kda_fwd( safe_gate, gate_lower_bound, state0 is not None, - ckpt, + checkpoint, device, plan_name, ) - if cache_key not in _fprop_cache: - _fprop_cache[cache_key] = _build_fprop_graph( + if cache_key not in fprop_cache: + fprop_cache[cache_key] = build_fprop_graph( total, N, H, @@ -421,11 +418,11 @@ def _kda_fwd( HV, K, V, - _torch_dtype_to_cudnn(q.dtype), + torch_dtype_to_cudnn(q.dtype), cudnn.data_type.FLOAT, - _torch_dtype_to_cudnn(beta.dtype), + torch_dtype_to_cudnn(beta.dtype), cudnn.data_type.FLOAT if state0 is not None else None, - _torch_dtype_to_cudnn(cu_seqlens.dtype), + torch_dtype_to_cudnn(cu_seqlens.dtype), float(scale), bool(output_final_state), bool(use_qk_l2norm_in_kernel), @@ -433,11 +430,11 @@ def _kda_fwd( bool(use_beta_sigmoid_in_kernel), bool(safe_gate), float(gate_lower_bound) if gate_lower_bound is not None else None, - ckpt, + checkpoint, ) - select_plan(_fprop_cache[cache_key][0], plan_name) + select_plan(fprop_cache[cache_key][0], plan_name) - graph, t = _fprop_cache[cache_key] + graph, t = fprop_cache[cache_key] o = torch.empty(total, HO, V, dtype=q.dtype, device=device) variant_pack = { @@ -459,16 +456,16 @@ def _kda_fwd( final_state = torch.empty(N, HO, V, K, dtype=torch.float32, device=device) variant_pack[t["fs"]] = final_state state_checkpoints = torch.empty(0, dtype=q.dtype, device=device) - if ckpt > 0: - total_checkpoints = max(total // ckpt + N, 1) + if checkpoint > 0: + total_checkpoints = max(total // checkpoint + N, 1) state_checkpoints = torch.empty(total_checkpoints, HO, V, K, dtype=q.dtype, device=device) variant_pack[t["state_checkpoints"]] = state_checkpoints - graph.execute(variant_pack, workspace=_graph_workspace(graph, device), handle=_get_handle(device)) + graph.execute(variant_pack, workspace=graph_workspace(graph, device), handle=get_handle(device)) return o, final_state, state_checkpoints -@_kda_fwd.register_fake -def _kda_fwd_fake( +@kda_fwd.register_fake +def kda_fwd_fake( q, k, v, @@ -514,7 +511,7 @@ def _kda_fwd_fake( # --------------------------------------------------------------------------- -def _build_bprop_graph( +def build_bprop_graph( total, N, H, @@ -528,7 +525,7 @@ def _build_bprop_graph( state_dtype, dstate_in_dtype, cu_dtype, - ckpt_rows, + checkpoint_rows, scale, use_qk_l2norm, batch_invariant, @@ -551,9 +548,9 @@ def _build_bprop_graph( dfs_t = None if dstate_in_dtype is not None: dfs_t = graph.tensor([N, HO, V, K], data_type=dstate_in_dtype, name="d_final_state") - ckpts_t = None - if ckpt_rows is not None: - ckpts_t = graph.tensor([ckpt_rows, HO, V, K], data_type=io_dtype, name="state_checkpoints") + checkpoints_t = None + if checkpoint_rows is not None: + checkpoints_t = graph.tensor([checkpoint_rows, HO, V, K], data_type=io_dtype, name="state_checkpoints") a_log_t = None dt_bias_t = None if safe_gate: @@ -567,7 +564,7 @@ def _build_bprop_graph( beta=beta_t, cu_seqlens=cu_t, dO=dO_t, - state_checkpoints=ckpts_t, + state_checkpoints=checkpoints_t, initial_state=state0_t, d_final_state=dfs_t, a_log=a_log_t, @@ -600,7 +597,7 @@ def _build_bprop_graph( dstate0=dstate0_t, d_a_log=dA_t, d_dt_bias=dDt_t, - ckpts=ckpts_t, + checkpoints=checkpoints_t, ) @@ -610,7 +607,7 @@ def _build_bprop_graph( @torch.library.custom_op("cudnn::kimi_delta_attention_bwd", mutates_args=()) -def _kda_bwd( +def kda_bwd( dO: torch.Tensor, q: torch.Tensor, k: torch.Tensor, @@ -644,8 +641,6 @@ def _kda_bwd( logits and ``dbeta`` is the raw-logit gradient. """ total, H, K = q.shape - # autograd materializes reduction grads as broadcast (stride-0) - # views; densify ONLY those (dense callers pass through untouched) if 0 in dO.stride(): dO = dO.contiguous() if d_final_state is not None and 0 in d_final_state.stride(): @@ -660,27 +655,27 @@ def _kda_bwd( if cu_seqlens.dtype not in (torch.int32, torch.int64): raise ValueError(f"kimi_delta_attention: cu_seqlens must be int32 or int64; got {cu_seqlens.dtype}") cu = cu_seqlens - _check_dtype("g", g, torch.float32) + check_dtype("g", g, torch.float32) if use_beta_sigmoid_in_kernel: - _check_dtype("beta", beta, q.dtype) + check_dtype("beta", beta, q.dtype) else: - _check_dtype("beta", beta, torch.float32) + check_dtype("beta", beta, torch.float32) if safe_gate: if a_log is None or dt_bias is None: raise ValueError("kimi_delta_attention: safe_gate requires a_log and dt_bias") - _check_dtype("a_log", a_log, torch.float32) - _check_dtype("dt_bias", dt_bias, torch.float32) + check_dtype("a_log", a_log, torch.float32) + check_dtype("dt_bias", dt_bias, torch.float32) elif a_log is not None or dt_bias is not None: raise ValueError("kimi_delta_attention: a_log/dt_bias require safe_gate=True") if initial_state is not None: - _check_dtype("initial_state", initial_state, torch.float32) + check_dtype("initial_state", initial_state, torch.float32) if initial_state.shape[0] != N: raise ValueError(f"initial_state must carry one state per sequence: got {initial_state.shape[0]} for {N} sequences") if d_final_state is not None: - _check_dtype("d_final_state", d_final_state, torch.float32) + check_dtype("d_final_state", d_final_state, torch.float32) if state_checkpoints is not None: - _check_dtype("state_checkpoints", state_checkpoints, q.dtype) - for _name, _t in ( + check_dtype("state_checkpoints", state_checkpoints, q.dtype) + for tensor_name, tensor in ( ("k", k), ("v", v), ("g", g), @@ -692,12 +687,12 @@ def _kda_bwd( ("a_log", a_log), ("dt_bias", dt_bias), ): - if _t is not None and _t.device != device: - raise ValueError(f"kimi_delta_attention: {_name} must be on q's device ({device}); got {_t.device}") + if tensor is not None and tensor.device != device: + raise ValueError(f"kimi_delta_attention: {tensor_name} must be on q's device ({device}); got {tensor.device}") state0 = initial_state if initial_state is not None else None dstate_in = d_final_state if d_final_state is not None else None - cache_key = _make_bprop_cache_key( + cache_key = make_bprop_cache_key( total, N, H, @@ -726,8 +721,8 @@ def _kda_bwd( device, plan_name, ) - if cache_key not in _bprop_cache: - _bprop_cache[cache_key] = _build_bprop_graph( + if cache_key not in bprop_cache: + bprop_cache[cache_key] = build_bprop_graph( total, N, H, @@ -735,12 +730,12 @@ def _kda_bwd( HV, K, V, - _torch_dtype_to_cudnn(q.dtype), - _torch_dtype_to_cudnn(g.dtype), - _torch_dtype_to_cudnn(beta.dtype), - _torch_dtype_to_cudnn(state0.dtype) if state0 is not None else None, - _torch_dtype_to_cudnn(dstate_in.dtype) if dstate_in is not None else None, - _torch_dtype_to_cudnn(cu_seqlens.dtype), + torch_dtype_to_cudnn(q.dtype), + torch_dtype_to_cudnn(g.dtype), + torch_dtype_to_cudnn(beta.dtype), + torch_dtype_to_cudnn(state0.dtype) if state0 is not None else None, + torch_dtype_to_cudnn(dstate_in.dtype) if dstate_in is not None else None, + torch_dtype_to_cudnn(cu_seqlens.dtype), state_checkpoints.shape[0] if state_checkpoints is not None else None, float(scale), bool(use_qk_l2norm_in_kernel), @@ -749,9 +744,9 @@ def _kda_bwd( safe_gate=bool(safe_gate), gate_lower_bound=float(gate_lower_bound) if gate_lower_bound is not None else None, ) - select_plan(_bprop_cache[cache_key][0], plan_name) + select_plan(bprop_cache[cache_key][0], plan_name) - graph, t = _bprop_cache[cache_key] + graph, t = bprop_cache[cache_key] dq = torch.empty(total, H, K, dtype=q.dtype, device=device) dk = torch.empty(total, HK, K, dtype=q.dtype, device=device) @@ -780,7 +775,7 @@ def _kda_bwd( if dstate_in is not None: variant_pack[t["dfs"]] = dstate_in if state_checkpoints is not None: - variant_pack[t["ckpts"]] = state_checkpoints + variant_pack[t["checkpoints"]] = state_checkpoints d_a_log = torch.empty(0, dtype=torch.float32, device=device) d_dt_bias = torch.empty(0, dtype=torch.float32, device=device) if safe_gate: @@ -790,14 +785,14 @@ def _kda_bwd( d_dt_bias = torch.empty(HO, K, dtype=torch.float32, device=device) variant_pack[t["d_a_log"]] = d_a_log variant_pack[t["d_dt_bias"]] = d_dt_bias - graph.execute(variant_pack, workspace=_graph_workspace(graph, device), handle=_get_handle(device)) + graph.execute(variant_pack, workspace=graph_workspace(graph, device), handle=get_handle(device)) if dstate0 is None: dstate0 = torch.empty(0, dtype=torch.float32, device=device) return dq, dk, dv, dg, dbeta, dstate0, d_a_log, d_dt_bias -@_kda_bwd.register_fake -def _kda_bwd_fake( +@kda_bwd.register_fake +def kda_bwd_fake( dO, q, k, @@ -840,7 +835,7 @@ def _kda_bwd_fake( # --------------------------------------------------------------------------- -def _kda_setup_context(ctx, inputs, output): +def kda_setup_context(ctx, inputs, output): ( q, k, @@ -861,11 +856,9 @@ def _kda_setup_context(ctx, inputs, output): checkpoint_every_n_tokens, plan_name, ) = inputs - # save_for_backward cannot hold None; keep initial_state as an attribute. - # g/beta are saved as passed: raw logits under safe_gate / use_beta_sigmoid. saved = [q, k, v, g, beta, cu_seqlens] - ctx.ckpt_reuse = checkpoint_every_n_tokens == 16 and output[2].numel() > 0 - if ctx.ckpt_reuse: + ctx.checkpoint_reuse = checkpoint_every_n_tokens == 16 and output[2].numel() > 0 + if ctx.checkpoint_reuse: saved.append(output[2]) if safe_gate: saved.extend([a_log, dt_bias]) @@ -882,11 +875,11 @@ def _kda_setup_context(ctx, inputs, output): ctx.mark_non_differentiable(output[2]) -def _kda_backward(ctx, dO, dFinal, _dstate_checkpoints): +def kda_backward(ctx, dO, dFinal, dstate_checkpoints): a_log = dt_bias = None if ctx.safe_gate: a_log, dt_bias = ctx.saved_tensors[-2:] - if ctx.ckpt_reuse: + if ctx.checkpoint_reuse: q, k, v, g, beta, cu_seqlens, state_checkpoints = ctx.saved_tensors[:7] else: q, k, v, g, beta, cu_seqlens = ctx.saved_tensors[:6] @@ -917,10 +910,6 @@ def _kda_backward(ctx, dO, dFinal, _dstate_checkpoints): dt_bias=dt_bias, plan_name=ctx.plan_name, ) - # q, k, v, g, beta, cu_seqlens, scale, initial_state, output_final_state, - # use_qk_l2norm_in_kernel, batch_invariant, use_beta_sigmoid_in_kernel, - # safe_gate, gate_lower_bound, a_log, dt_bias, checkpoint_every_n_tokens, - # plan_name return ( dq, dk, @@ -945,8 +934,8 @@ def _kda_backward(ctx, dO, dFinal, _dstate_checkpoints): torch.library.register_autograd( "cudnn::kimi_delta_attention_fwd", - _kda_backward, - setup_context=_kda_setup_context, + kda_backward, + setup_context=kda_setup_context, ) diff --git a/test/python/linear_attention/test_fla_compat.py b/test/python/linear_attention/test_fla_compat.py index d8bda57cc..f1b60aaa0 100644 --- a/test/python/linear_attention/test_fla_compat.py +++ b/test/python/linear_attention/test_fla_compat.py @@ -61,9 +61,9 @@ def _leaves(master, dtype): return lv -def _run(fn, master, dtype): +def _run(fn, master, dtype, **kw): lv = _leaves(master, dtype) - o, _ = fn(lv["q"], lv["k"], lv["v"], lv["g"], lv["beta"], output_final_state=False) + o, _ = fn(lv["q"], lv["k"], lv["v"], lv["g"], lv["beta"], output_final_state=False, **kw) return o, lv @@ -330,10 +330,15 @@ def check(name, a, b, ref, slack): def test_fallback_is_transparent(): - """An unsupported config (K != 128) falls back and returns FLA's exact result.""" - m = _master(2, 256, 4, 4, 64, 128, seed=1) # K=64 -> native declines - o_fla, _ = _run(chunk_gated_delta_rule, m, torch.bfloat16) - o_cud, _ = _run(shim, m, torch.bfloat16) + """A variant the native op does not model falls back and returns FLA's exact result. + + Keyed on ``allow_neg_eigval``, which the shim declines by construction, so this stays + a fallback no matter which shapes the engines grow support for. + """ + m = _master(2, 256, 4, 4, 128, 128, seed=1) + kw = dict(allow_neg_eigval=True, use_beta_sigmoid_in_kernel=True) # FLA requires the pair + o_fla, _ = _run(chunk_gated_delta_rule, m, torch.bfloat16, **kw) + o_cud, _ = _run(shim, m, torch.bfloat16, **kw) assert last_path().startswith("fallback"), f"expected fallback, got {last_path()}" torch.testing.assert_close(o_cud, o_fla, rtol=0, atol=0) diff --git a/test/python/linear_attention/test_la.py b/test/python/linear_attention/test_la.py index 02af57bde..0e7984b1e 100644 --- a/test/python/linear_attention/test_la.py +++ b/test/python/linear_attention/test_la.py @@ -20,7 +20,9 @@ import contextlib import functools import math +import os import threading +import time import pytest torch = pytest.importorskip("torch") @@ -40,6 +42,7 @@ ] VARIANTS = ("gdn", "kda", "gdn2") +SPLIT_T = 4096 # long enough that the work-item table cuts CHUNK = {"gdn": 64, "kda": 16, "gdn2": 16} FWD_TOL = {torch.bfloat16: 2e-2, torch.float16: 1e-2} @@ -103,8 +106,8 @@ def family_engines(backend_name): def clear_op_caches(): for mod in op_modules().values(): - mod._fprop_cache.clear() - mod._bprop_cache.clear() + mod.fprop_cache.clear() + mod.bprop_cache.clear() class Case: @@ -299,7 +302,7 @@ def test_backend_pin_selects_engine(backend, variant): with waive_unsupported(backend, variant): pinned_op(backend, variant)(*op_args(case)) mod = op_modules()[variant] - names = {g.selected_engine.name for g, entry in mod._fprop_cache.values() if g.selected_engine is not None} + names = {g.selected_engine.name for g, entry in mod.fprop_cache.values() if g.selected_engine is not None} assert names == {f"{variant}_{backend.name}"}, f"expected only {variant}_{backend.name} to serve, got {names}" @@ -410,6 +413,22 @@ def test_fwd_initial_state(backend, variant, T): assert_fwd_parity(backend, case, use_initial_state=True) +@pytest.mark.parametrize("variant", VARIANTS) +def test_fwd_split_initial_state(backend, variant): + """An initial state at a split-inducing length: the default schedule and the + uncut batch-invariant schedule must agree with each other and the reference.""" + case = make_case(variant, torch.bfloat16, B=1, T=SPLIT_T) + set_seed(SEED + 1) + state0 = torch.randn(case.N, case.HO, case.V, case.K, device="cuda", dtype=torch.float32) * 0.05 + o_split, fs_split = run_fwd(backend, case, initial_state=state0, output_final_state=True) + o_uncut, fs_uncut = run_fwd(backend, case, initial_state=state0, output_final_state=True, batch_invariant=True) + assert_rms_close("o split-vs-uncut", o_split, o_uncut.float(), FWD_TOL[torch.bfloat16]) + assert_rms_close("final_state split-vs-uncut", fs_split, fs_uncut.float(), STATE_TOL[torch.bfloat16]) + o_ref, fs_ref = reference(case, initial_state=state0) + assert_rms_close("o vs reference", o_split, o_ref, FWD_TOL[torch.bfloat16]) + assert_rms_close("final_state vs reference", fs_split, fs_ref, STATE_TOL[torch.bfloat16]) + + @pytest.mark.parametrize("T1,T2", [(128, 128), (64, 192), (192, 121)]) @pytest.mark.parametrize("variant", VARIANTS) def test_fwd_chunked_prefill(backend, variant, T1, T2): @@ -598,6 +617,29 @@ def test_bwd_initial_state(backend, variant): assert_bwd_parity(backend, make_case(variant, torch.bfloat16, T=128), use_initial_state=True) +@pytest.mark.parametrize("variant", VARIANTS) +def test_bwd_split_initial_state(backend, variant): + """Backward over a cut work-item table with an initial state: the gradients + must match the uncut (batch-invariant) table.""" + case = make_case(variant, torch.bfloat16, B=1, T=SPLIT_T) + set_seed(SEED + 1) + state0 = torch.randn(case.N, case.HO, case.V, case.K, device="cuda", dtype=torch.float32) * 0.05 + tensors = {"q": case.q, "k": case.k, "v": case.v, "g": case.gates["g"], "beta": case.gates["beta"]} + if variant == "gdn2": + tensors["w"] = case.gates["w"] + dO, grads = None, {} + for tag, kw in (("split", {}), ("uncut", {"batch_invariant": True})): + leaves = [to_thd(t).detach().clone().requires_grad_(True) for t in tensors.values()] + s0 = state0.detach().clone().requires_grad_(True) + with waive_unsupported(backend, variant): + o, _ = pinned_op(backend, variant)(*leaves, case.cu, initial_state=s0, output_final_state=True, **kw) + if dO is None: + dO = torch.randn_like(o) + grads[tag] = torch.autograd.grad([o], leaves + [s0], [dO]) + for name, got, want in zip(list(tensors) + ["initial_state"], grads["split"], grads["uncut"]): + assert_rms_close(f"d{name} split-vs-uncut", got, want.float(), BWD_TOL[torch.bfloat16]) + + @pytest.mark.parametrize("variant", VARIANTS) def test_bwd_d_final_state(backend, variant): assert_bwd_parity(backend, make_case(variant, torch.bfloat16, T=128), use_initial_state=True, use_dfs=True) @@ -645,6 +687,96 @@ def test_bwd_with_checkpoints(backend, variant): assert t.grad is not None and torch.isfinite(t.grad).all(), f"bad grad for {name}" +# --------------------------------------------------------------------------- +# Layout (innermost-contiguous inputs; outer strides pass straight to the kernels) +# --------------------------------------------------------------------------- + + +def strided_copy(t): + """A non-contiguous copy of ``t``: the values land in the leading columns of + a buffer with a doubled innermost extent, so every outer stride changes + while stride(-1) stays 1 (the one layout fact the FROST engines gate on).""" + wide = torch.empty(*t.shape[:-1], 2 * t.shape[-1], device=t.device, dtype=t.dtype) + view = wide[..., : t.shape[-1]] + view.copy_(t) + assert view.stride(-1) == 1 and not view.is_contiguous() + return view + + +def fused_qkv_views(case): + """q/k/v as slices of one fused projection buffer — the layout a fused QKV + matmul hands the op: innermost-contiguous, never whole-tensor contiguous. + A doc-level whole-tensor rule forces callers to copy exactly these.""" + tensors = [to_thd(case.q), to_thd(case.k), to_thd(case.v)] + widths = [t.shape[1] * t.shape[2] for t in tensors] + fused = torch.empty(tensors[0].shape[0], sum(widths), device="cuda", dtype=case.dtype) + views, base = [], 0 + for t, width in zip(tensors, widths): + view = fused[:, base : base + width].unflatten(-1, t.shape[1:]) + view.copy_(t) + assert view.stride(-1) == 1 and not view.is_contiguous() + views.append(view) + base += width + return views + + +@pytest.mark.parametrize("backend", ["frost"], indirect=True) +@pytest.mark.parametrize("variant", VARIANTS) +def test_fwd_innermost_contiguous_inputs(backend, variant): + """Fused-projection q/k/v slices, strided gates and a strided initial state + match the contiguous run bitwise: the kernels take every outer stride as a + runtime argument (the TMA descriptors are rebuilt from the live strides on + each launch) and the engine gates only on a stride-1 innermost dim.""" + case = make_case(variant, torch.bfloat16, seq_lens=[192, 251]) + set_seed(SEED + 1) + state0 = torch.randn(case.N, case.HO, case.V, case.K, device="cuda", dtype=torch.float32) * 0.05 + o_ref, fs_ref = run_fwd(backend, case, initial_state=state0, output_final_state=True) + args = fused_qkv_views(case) + [strided_copy(to_thd(case.gates["g"])), strided_copy(to_thd(case.gates["beta"]))] + if variant == "gdn2": + args.append(strided_copy(to_thd(case.gates["w"]))) + with waive_unsupported(backend, variant): + o, fs = pinned_op(backend, variant)(*args, case.cu, initial_state=strided_copy(state0), output_final_state=True) + assert torch.equal(bits(o), bits(o_ref)), "strided inputs changed o" + assert torch.equal(bits(fs), bits(fs_ref)), "strided inputs changed final_state" + + +@pytest.mark.parametrize("backend", ["frost"], indirect=True) +@pytest.mark.parametrize("variant", VARIANTS) +def test_bwd_innermost_contiguous_inputs(backend, variant): + """Backward from strided leaves and a strided incoming dO: every gradient + matches the contiguous run bitwise (the op densifies only stride-0 + broadcast grads; the recompute and bprop kernels take the strides).""" + case = make_case(variant, torch.bfloat16, seq_lens=[192, 251]) + gate_names = ["g", "beta"] + (["w"] if variant == "gdn2" else []) + + def grads_from(leaves, dO): + with waive_unsupported(backend, variant): + o, _ = pinned_op(backend, variant)(*leaves, case.cu) + return torch.autograd.grad([o], leaves, [dO]) + + set_seed(SEED + 3) + dO = torch.randn(case.T, case.HO, case.V, device="cuda", dtype=case.dtype) + contiguous = [to_thd(t).detach().clone() for t in (case.q, case.k, case.v)] + [to_thd(case.gates[n]).detach().clone() for n in gate_names] + strided = fused_qkv_views(case) + [strided_copy(to_thd(case.gates[n])) for n in gate_names] + grads_c = grads_from([t.requires_grad_(True) for t in contiguous], dO) + grads_s = grads_from([t.requires_grad_(True) for t in strided], strided_copy(dO)) + for name, gc, gs in zip(["q", "k", "v", *gate_names], grads_c, grads_s): + assert torch.equal(bits(gc), bits(gs)), f"d{name} differs between contiguous and strided inputs" + + +@pytest.mark.parametrize("backend", ["cutile"], indirect=True) +@pytest.mark.parametrize("variant", ["gdn", "kda"]) +def test_cutile_rejects_strided_inputs(backend, variant): + """The cuTile backend stages rank-merged views, so it cannot take outer + strides: a strided buffer must raise its contract error, never read the + padding or silently copy.""" + case = make_case(variant, torch.bfloat16, T=64) + args = [strided_copy(to_thd(case.q)), to_thd(case.k), to_thd(case.v), to_thd(case.gates["g"]), to_thd(case.gates["beta"])] + with waive_unsupported(backend, variant): + with pytest.raises(ValueError, match="must be contiguous"): + pinned_op(backend, variant)(*args, case.cu) + + # --------------------------------------------------------------------------- # Checkpoints (per-chunk state series) # --------------------------------------------------------------------------- @@ -1358,3 +1490,134 @@ def launch(): torch.cuda.synchronize() for i, (a, b) in enumerate(zip(eager, captured)): assert torch.equal(bits(a), bits(b)), f"replayed output {i} differs from eager" + + +# --------------------------------------------------------------------------- +# Hang regression (mbarrier parity-aperture class) +# --------------------------------------------------------------------------- +# A single-slot mbarrier parity wait wedges forever if the barrier completes +# twice between one waiter's polls, so any arrive whose issue is not gated on +# every waiter having observed the previous completion is a latent deadlock. +# The windows open at work-item boundaries (tile-last chunks, zero-length +# items, initial-state seeds), so these tests run boundary-dense shapes many +# times with per-iteration syncs and NaN-poisoned buffer reuse (the host-side +# traffic that makes the windows hittable), under a watchdog that turns a +# wedge into a loud abort instead of a silent suite hang. + +HANG_STRESS_ITERS = int(os.environ.get("CUDNN_LA_HANG_STRESS_ITERS", "400")) +HANG_STRESS_TIMEOUT = float(os.environ.get("CUDNN_LA_HANG_STRESS_TIMEOUT", "120")) +HANG_STRESS_COMPILE_TIMEOUT = float(os.environ.get("CUDNN_LA_HANG_STRESS_COMPILE_TIMEOUT", "900")) + + +@contextlib.contextmanager +def wedge_watchdog(label, heartbeat, capfd=None, timeout=None): + timeout = HANG_STRESS_TIMEOUT if timeout is None else timeout + done = threading.Event() + + def watch(): + while not done.wait(5.0): + if time.monotonic() - heartbeat[0] > timeout: + message = f"\nHANG: {label} made no progress for {timeout:.0f}s; " "GPU kernel wedge (mbarrier parity-aperture class); aborting process" + with contextlib.suppress(Exception): + with capfd.disabled() if capfd is not None else contextlib.nullcontext(): + print(message, flush=True) + os._exit(70) + + thread = threading.Thread(target=watch, daemon=True) + thread.start() + try: + yield + finally: + done.set() + thread.join() + + +def run_hang_stress(backend, case, *, use_initial_state=False, fwd_each_iter=False, iters=None, label="", capfd=None): + """Repeated backward over a retained graph of a boundary-dense case; a + wedge aborts loudly. The previous grads are NaN-filled and freed BEFORE + the next launch so the allocator hands poisoned blocks straight to it + (the host-side traffic pattern the wedges need), and the final + iteration's finiteness check doubles as a partial-store canary.""" + iters = HANG_STRESS_ITERS if iters is None else iters + + def build_inputs(c): + tensors = [to_thd(c.q), to_thd(c.k), to_thd(c.v), to_thd(c.gates["g"]), to_thd(c.gates["beta"])] + if c.variant == "gdn2": + tensors.append(to_thd(c.gates["w"])) + leaves = [t.detach().clone().requires_grad_(True) for t in tensors] + state0 = None + if use_initial_state: + state0 = (torch.randn(c.N, c.HO, c.V, c.K, device="cuda", dtype=torch.float32) * 0.05).requires_grad_(True) + return leaves, state0 + + heartbeat = [time.monotonic()] + with waive_unsupported(backend, case.variant): + fn = pinned_op(backend, case.variant) + # compile on a single-wave case (one tile per CTA: no cross-tile + # handshakes, so it cannot wedge) so the real case runs entirely + # under the tight watchdog + tiny = make_case(case.variant, case.dtype, T=64, H=16) + tiny_leaves, tiny_state0 = build_inputs(tiny) + with wedge_watchdog(f"{label} (compile warmup)", heartbeat, capfd=capfd, timeout=HANG_STRESS_COMPILE_TIMEOUT): + o, fs = fn(*tiny_leaves, tiny.cu, initial_state=tiny_state0, output_final_state=True) + torch.autograd.grad([o], tiny_leaves + ([tiny_state0] if use_initial_state else []), [torch.randn_like(o)]) + torch.cuda.synchronize() + del o, fs, tiny_leaves, tiny_state0 + leaves, state0 = build_inputs(case) + grad_inputs = leaves + ([state0] if use_initial_state else []) + heartbeat[0] = time.monotonic() + with wedge_watchdog(label, heartbeat, capfd=capfd): + o, fs = fn(*leaves, case.cu, initial_state=state0, output_final_state=True) + dO = torch.randn_like(o) + grads = torch.autograd.grad([o], grad_inputs, [dO], retain_graph=True) + torch.cuda.synchronize() + heartbeat[0] = time.monotonic() + for _ in range(iters): + for g in grads: + g.fill_(float("nan")) + del grads + if fwd_each_iter: + o, fs = fn(*leaves, case.cu, initial_state=state0, output_final_state=True) + grads = torch.autograd.grad([o], grad_inputs, [dO], retain_graph=not fwd_each_iter) + torch.cuda.synchronize() + heartbeat[0] = time.monotonic() + assert torch.isfinite(o.float()).all(), "non-finite forward output after stress" + names = ["q", "k", "v", "g", "beta"] + (["w"] if case.variant == "gdn2" else []) + (["initial_state"] if use_initial_state else []) + for name, g in zip(names, grads): + assert torch.isfinite(g.float()).all(), f"non-finite d{name} after stress" + + +@pytest.mark.gpu_exclusive +@pytest.mark.xdist_group(name="gpu_exclusive") +@pytest.mark.parametrize("backend", ["frost"], indirect=True) +@pytest.mark.parametrize("variant", VARIANTS) +def test_hang_stress_tile_boundary_pipeline(backend, variant, capfd): + """Many short tiles per CTA: every tile-last chunk opens the parity window + of any arrive not gated on its waiter (the gdn2 bwd sDy release wedge).""" + case = make_case(variant, torch.bfloat16, seq_lens=[128] * 148, H=16) + run_hang_stress(backend, case, label=f"tile_boundary_pipeline[{variant}]", capfd=capfd) + + +@pytest.mark.gpu_exclusive +@pytest.mark.xdist_group(name="gpu_exclusive") +@pytest.mark.parametrize("backend", ["frost"], indirect=True) +@pytest.mark.parametrize("variant", VARIANTS) +def test_hang_stress_zero_length_tiles(backend, variant, capfd): + """Chunked tiles alternating with zero-length work items: an empty item's + body has no cross-warp waits, so per-tile handshakes whose arrive is + unconditional can complete twice between one waiter's polls (the bwd + dstate0 wedge).""" + case = make_case(variant, torch.bfloat16, seq_lens=[64, 0] * 96, H=16) + run_hang_stress(backend, case, label=f"zero_length_tiles[{variant}]", capfd=capfd) + + +@pytest.mark.gpu_exclusive +@pytest.mark.xdist_group(name="gpu_exclusive") +@pytest.mark.parametrize("backend", ["frost"], indirect=True) +@pytest.mark.parametrize("variant", VARIANTS) +def test_hang_stress_initial_state_boundaries(backend, variant, capfd): + """Initial-state builds with boundary-dense varlen: a seed-slot arrive + issued at the top of a tile races the previous tile's final state wait + (the gdn prefill/recompute seed-credit wedge).""" + case = make_case(variant, torch.bfloat16, seq_lens=[64, 0, 128, 0, 64] * 24, H=16) + run_hang_stress(backend, case, use_initial_state=True, fwd_each_iter=True, label=f"initial_state_boundaries[{variant}]", capfd=capfd)