From b1fe18acbbb06e7a02d0e0090b3f7f1e978acbaf Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 23 May 2026 03:05:08 +0000 Subject: [PATCH] cleanup: remove archive/ (240 stale files), stale example9/10, fix test table, add Stage D plan --- README.md | 54 +- tests/archive/debug_output.py | 70 -- tests/archive/debug_stages.py | 59 -- tests/archive/diag_layouts.py | 85 --- tests/archive/diag_tma_layout.py | 116 ---- tests/archive/diag_tma_shapes.py | 120 ---- tests/archive/diag_tma_shapes2.py | 122 ---- tests/archive/diag_tma_shapes3.py | 133 ---- tests/archive/diag_tmem.py | 91 --- tests/archive/fmha_v3_identity_diag.py | 313 --------- tests/archive/fmha_v3_real_softmax.py | 344 ---------- tests/archive/fmha_v3_stage_c_example1.py | 479 ------------- tests/archive/fmha_v3_stage_c_example2.py | 451 ------------- tests/archive/fmha_v3_stage_c_example3.py | 456 ------------- tests/archive/fmha_v3_stage_c_example4.py | 469 ------------- tests/archive/fmha_v3_stage_c_example5.py | 469 ------------- tests/archive/fmha_v3_stage_c_example6.py | 472 ------------- tests/archive/fmha_v3_stage_c_example8.py | 486 -------------- tests/archive/native_stage_c_patch.py | 351 ---------- tests/archive/quick_v3_multitile.py | 53 -- tests/archive/stage_b_debug5.py | 187 ------ tests/archive/test_128_128_fmha_v.py | 236 ------- tests/archive/test_128_16_bigP.py | 383 ----------- tests/archive/test_128_16_debug.py | 384 ----------- tests/archive/test_128_16_debug2.py | 389 ----------- tests/archive/test_128_16_debug3.py | 384 ----------- tests/archive/test_128_16_fp16.py | 383 ----------- tests/archive/test_128_16_full.py | 385 ----------- tests/archive/test_128_16_minimal.py | 385 ----------- tests/archive/test_128_16_nogC.py | 383 ----------- tests/archive/test_128_16_nopack.py | 383 ----------- tests/archive/test_128_16_nosoftmax.py | 383 ----------- tests/archive/test_128_16_pAtS.py | 367 ---------- tests/archive/test_128_16_pvlayout.py | 387 ----------- tests/archive/test_128_16_pvpack.py | 383 ----------- tests/archive/test_128_16_pvwrite.py | 383 ----------- tests/archive/test_128_16_qkread.py | 383 ----------- tests/archive/test_128_16_smem.py | 382 ----------- tests/archive/test_128_16_stepA.py | 383 ----------- tests/archive/test_128_16_stepB.py | 383 ----------- tests/archive/test_128_16_stepC.py | 383 ----------- tests/archive/test_128_16_stepD.py | 386 ----------- tests/archive/test_128_16_stepE.py | 385 ----------- tests/archive/test_128_16_tiler.py | 383 ----------- tests/archive/test_128_16_v8.py | 385 ----------- tests/archive/test_128_16_zeropad.py | 385 ----------- tests/archive/test_128_32_ctafix.py | 386 ----------- tests/archive/test_128_32_ctafix2.py | 384 ----------- tests/archive/test_128_32_native.py | 382 ----------- tests/archive/test_128_32_vdiag.py | 385 ----------- tests/archive/test_128_32_zeropad.py | 385 ----------- tests/archive/test_afrag_roundtrip.py | 175 ----- tests/archive/test_attention.py | 173 ----- tests/archive/test_attention_path_b200.py | 267 -------- tests/archive/test_b_afrag2.py | 216 ------ tests/archive/test_b_layout.py | 49 -- tests/archive/test_bf16_elemwise.py | 237 ------- tests/archive/test_bf16_pack.py | 237 ------- tests/archive/test_bf16_recast_full.py | 242 ------- tests/archive/test_bf16_recast_simple.py | 239 ------- tests/archive/test_blackwell_attn_b200.py | 318 --------- tests/archive/test_cache.py | 252 ------- tests/archive/test_compile_custom_op.py | 189 ------ tests/archive/test_csa_attention_b200.py | 251 ------- tests/archive/test_csa_sparse_attn_b200.py | 399 ----------- tests/archive/test_custom_op.py | 138 ---- tests/archive/test_decode_attention_b200.py | 460 ------------- tests/archive/test_decode_pipeline.py | 140 ---- tests/archive/test_decode_vs_prefill_b200.py | 274 -------- tests/archive/test_dense_router.py | 27 - tests/archive/test_diag_layout.py | 373 ----------- tests/archive/test_diag_multitile.py | 30 - tests/archive/test_diag_permute.py | 80 --- tests/archive/test_diag_smem_layout.py | 72 -- tests/archive/test_diag_v_mma128.py | 374 ----------- tests/archive/test_diag_v_ones.py | 62 -- tests/archive/test_diag_v_truncid.py | 369 ---------- tests/archive/test_e2e_decode_b200.py | 425 ------------ tests/archive/test_error_pattern.py | 85 --- tests/archive/test_fmha_pipeline.py | 354 ---------- tests/archive/test_fmha_v1.py | 253 ------- tests/archive/test_fmha_v2.py | 245 ------- tests/archive/test_fmha_v2_fixed.py | 277 -------- tests/archive/test_fmha_v3_12w.py | 327 --------- tests/archive/test_fmha_v3_debug.py | 284 -------- tests/archive/test_fmha_v3_diag.py | 317 --------- tests/archive/test_fmha_v3_diag_fixed.py | 307 --------- tests/archive/test_full_layer_b200.py | 258 ------- tests/archive/test_full_layer_nan_b200.py | 348 ---------- tests/archive/test_full_model_b200.py | 314 --------- tests/archive/test_fused_step1.py | 92 --- tests/archive/test_inspect_types.py | 73 -- tests/archive/test_interleave.py | 144 ---- tests/archive/test_interleave_gemm.py | 137 ---- tests/archive/test_inv_rope.py | 126 ---- tests/archive/test_kv_cache_b200.py | 358 ---------- tests/archive/test_layer_schedule.py | 85 --- tests/archive/test_layout_compare.py | 95 --- tests/archive/test_mma_si_only.py | 247 ------- tests/archive/test_mma_si_pv.py | 355 ---------- tests/archive/test_model_forward_b200.py | 238 ------- tests/archive/test_moe_nan_b200.py | 225 ------- tests/archive/test_moe_runner_nan_b200.py | 190 ------ tests/archive/test_multilayer.py | 161 ----- tests/archive/test_nvfp4_attention_b200.py | 255 ------- tests/archive/test_nvfp4_attn_gemm_b200.py | 373 ----------- tests/archive/test_nvfp4_mapper.py | 186 ------ tests/archive/test_o_projection.py | 159 ----- tests/archive/test_o_projection_b200.py | 306 --------- tests/archive/test_packing_diag.py | 133 ---- tests/archive/test_pair_swap.py | 119 ---- tests/archive/test_pair_swap2.py | 95 --- tests/archive/test_pipeline_real_weights.py | 164 ----- tests/archive/test_pv64.py | 244 ------- tests/archive/test_pv64_fmha_v.py | 258 ------- tests/archive/test_pv64_kmajor_v.py | 238 ------- tests/archive/test_pv64_no_softmax.py | 222 ------ tests/archive/test_pv64_nosoftmax_fmha_v.py | 207 ------ tests/archive/test_pv_diag.py | 383 ----------- tests/archive/test_pv_mma_mn_major.py | 303 --------- tests/archive/test_quick_rand.py | 36 - tests/archive/test_recast_minimal.py | 237 ------- tests/archive/test_ref_minimal.py | 41 -- tests/archive/test_reference_fmha.py | 80 --- tests/archive/test_rope_kv_b200.py | 152 ----- tests/archive/test_router.py | 217 ------ tests/archive/test_runner_vs_pipeline.py | 210 ------ tests/archive/test_scale_assembly.py | 116 ---- tests/archive/test_scale_debug.py | 77 --- tests/archive/test_shared_expert.py | 163 ----- tests/archive/test_silu_step1.py | 97 --- tests/archive/test_softmax_only.py | 288 -------- tests/archive/test_softmax_store_debug.py | 247 ------- tests/archive/test_sparse_attn_b200.py | 364 ---------- tests/archive/test_sparse_decode.py | 71 -- tests/archive/test_stage_a_copy.py | 372 ----------- tests/archive/test_stage_a_minimal.py | 395 ----------- tests/archive/test_stage_a_pv_created.py | 376 ----------- tests/archive/test_stage_a_pv_param.py | 374 ----------- tests/archive/test_stage_a_qk.py | 632 ------------------ tests/archive/test_stage_a_v2.py | 372 ----------- tests/archive/test_stage_a_with_pv_mma.py | 374 ----------- tests/archive/test_stage_b_afrag.py | 210 ------ tests/archive/test_stage_b_afrag2.py | 217 ------ tests/archive/test_stage_b_debug.py | 252 ------- tests/archive/test_stage_b_debug2.py | 224 ------- tests/archive/test_stage_b_debug3.py | 198 ------ tests/archive/test_stage_b_debug4.py | 205 ------ tests/archive/test_stage_b_diag.py | 384 ----------- tests/archive/test_stage_b_final.py | 207 ------ tests/archive/test_stage_b_identity.py | 487 -------------- tests/archive/test_stage_b_minimal.py | 271 -------- tests/archive/test_stage_b_ntile_v1.py | 399 ----------- tests/archive/test_stage_b_ntile_v3.py | 357 ---------- tests/archive/test_stage_b_ntile_v5.py | 415 ------------ tests/archive/test_stage_b_ntile_v6.py | 400 ----------- tests/archive/test_stage_b_ntile_v7.py | 376 ----------- tests/archive/test_stage_b_ntile_v8.py | 387 ----------- tests/archive/test_stage_b_pipeline_only.py | 281 -------- tests/archive/test_stage_b_v1.py | 420 ------------ tests/archive/test_stage_b_v10.py | 321 --------- tests/archive/test_stage_b_v11.py | 210 ------ tests/archive/test_stage_b_v11b.py | 398 ----------- tests/archive/test_stage_b_v12.py | 408 ----------- tests/archive/test_stage_b_v13.py | 401 ----------- tests/archive/test_stage_b_v14.py | 352 ---------- tests/archive/test_stage_b_v16.py | 453 ------------- tests/archive/test_stage_b_v17.py | 450 ------------- tests/archive/test_stage_b_v18.py | 452 ------------- tests/archive/test_stage_b_v19.py | 450 ------------- tests/archive/test_stage_b_v2.py | 407 ----------- tests/archive/test_stage_b_v20.py | 362 ---------- tests/archive/test_stage_b_v22.py | 314 --------- tests/archive/test_stage_b_v22_bug1fix.py | 364 ---------- tests/archive/test_stage_b_v23.py | 364 ---------- tests/archive/test_stage_b_v24.py | 377 ----------- tests/archive/test_stage_b_v25.py | 380 ----------- tests/archive/test_stage_b_v26.py | 367 ---------- tests/archive/test_stage_b_v27.py | 375 ----------- tests/archive/test_stage_b_v28.py | 377 ----------- tests/archive/test_stage_b_v29.py | 365 ---------- tests/archive/test_stage_b_v3.py | 375 ----------- tests/archive/test_stage_b_v30.py | 382 ----------- tests/archive/test_stage_b_v4.py | 259 ------- tests/archive/test_stage_b_v5.py | 341 ---------- tests/archive/test_stage_b_v6.py | 342 ---------- tests/archive/test_stage_b_v7.py | 450 ------------- tests/archive/test_stage_b_v7_rep128.py | 445 ------------ tests/archive/test_stage_b_v7_rep16.py | 445 ------------ tests/archive/test_stage_b_v7_rep64.py | 445 ------------ tests/archive/test_stage_b_v7_rep8.py | 445 ------------ tests/archive/test_stage_b_v8.py | 403 ----------- tests/archive/test_stage_b_v8b.py | 354 ---------- tests/archive/test_stage_b_v9.py | 348 ---------- tests/archive/test_step2_subtile.py | 98 --- tests/archive/test_step2_subtile_v2.py | 114 ---- tests/archive/test_store_verify.py | 285 -------- tests/archive/test_store_verify2.py | 272 -------- tests/archive/test_tma_coord.py | 102 --- tests/archive/test_tmem_addressing.py | 263 -------- tests/archive/test_tmem_col2.py | 383 ----------- tests/archive/test_tmem_col3.py | 392 ----------- tests/archive/test_tmem_col4.py | 388 ----------- tests/archive/test_tmem_col5.py | 392 ----------- tests/archive/test_tmem_col5_16.py | 397 ----------- tests/archive/test_tmem_col5_32.py | 390 ----------- tests/archive/test_tmem_col_offset.py | 388 ----------- tests/archive/test_tmem_copy_roundtrip.py | 271 -------- tests/archive/test_tmem_debug.py | 397 ----------- tests/archive/test_tmem_debug2.py | 392 ----------- tests/archive/test_tmem_fp32_roundtrip.py | 255 ------- tests/archive/test_tmem_layout_diag.py | 81 --- tests/archive/test_tmem_pure_fp32.py | 237 ------- tests/archive/test_uniform_fp4.py | 30 - tests/archive/test_v28c_noepi.py | 315 --------- tests/archive/test_v4_attention_b200.py | 287 -------- tests/archive/test_v_mode_fix.py | 284 -------- tests/archive/test_vllm_codepaths_b200.py | 252 ------- tests/archive/test_vsmem_diag.py | 385 ----------- tests/archive/test_warmup_gs.py | 206 ------ tests/archive/test_wo_a.py | 171 ----- tests/archive/test_wo_a_bmm.py | 79 --- tests/archive/unit_test_128_128_vdiag.py | 385 ----------- tests/archive/unit_test_fmha_v3_debug.py | 469 ------------- tests/archive/unit_test_fmha_v3_fixed_v.py | 512 -------------- tests/archive/unit_test_fmha_v3_noop_c9.py | 469 ------------- tests/archive/unit_test_fmha_v3_per_row.py | 587 ---------------- .../archive/unit_test_fmha_v3_per_row_min.py | 465 ------------- tests/archive/unit_test_fmha_v3_proper.py | 416 ------------ tests/archive/unit_test_fmha_v3_pva_c9.py | 484 -------------- tests/archive/unit_test_fmha_v3_scalar.py | 493 -------------- tests/archive/unit_test_fmha_v3_shapes.py | 518 -------------- tests/archive/unit_test_fmha_v3_softmax.py | 511 -------------- tests/archive/unit_test_fmha_v3_stage_c.py | 327 --------- tests/archive/unit_test_fmha_v3_stage_c2.py | 466 ------------- .../archive/unit_test_fmha_v3_stage_c_full.py | 454 ------------- .../archive/unit_test_fmha_v3_stage_c_min.py | 305 --------- tests/archive/unit_test_fmha_v3_vec_c9.py | 488 -------------- tests/archive/unit_test_pv64_with_softmax.py | 257 ------- tests/archive/unit_test_qk_softmax.py | 252 ------- tests/archive/unit_test_qkonly.py | 269 -------- tests/fmha_v3_stage_c_example10.py | 510 -------------- tests/fmha_v3_stage_c_example9.py | 495 -------------- 243 files changed, 48 insertions(+), 72898 deletions(-) delete mode 100644 tests/archive/debug_output.py delete mode 100644 tests/archive/debug_stages.py delete mode 100644 tests/archive/diag_layouts.py delete mode 100644 tests/archive/diag_tma_layout.py delete mode 100644 tests/archive/diag_tma_shapes.py delete mode 100644 tests/archive/diag_tma_shapes2.py delete mode 100644 tests/archive/diag_tma_shapes3.py delete mode 100644 tests/archive/diag_tmem.py delete mode 100644 tests/archive/fmha_v3_identity_diag.py delete mode 100644 tests/archive/fmha_v3_real_softmax.py delete mode 100644 tests/archive/fmha_v3_stage_c_example1.py delete mode 100644 tests/archive/fmha_v3_stage_c_example2.py delete mode 100644 tests/archive/fmha_v3_stage_c_example3.py delete mode 100644 tests/archive/fmha_v3_stage_c_example4.py delete mode 100644 tests/archive/fmha_v3_stage_c_example5.py delete mode 100644 tests/archive/fmha_v3_stage_c_example6.py delete mode 100644 tests/archive/fmha_v3_stage_c_example8.py delete mode 100644 tests/archive/native_stage_c_patch.py delete mode 100644 tests/archive/quick_v3_multitile.py delete mode 100644 tests/archive/stage_b_debug5.py delete mode 100644 tests/archive/test_128_128_fmha_v.py delete mode 100644 tests/archive/test_128_16_bigP.py delete mode 100644 tests/archive/test_128_16_debug.py delete mode 100644 tests/archive/test_128_16_debug2.py delete mode 100644 tests/archive/test_128_16_debug3.py delete mode 100644 tests/archive/test_128_16_fp16.py delete mode 100644 tests/archive/test_128_16_full.py delete mode 100644 tests/archive/test_128_16_minimal.py delete mode 100644 tests/archive/test_128_16_nogC.py delete mode 100644 tests/archive/test_128_16_nopack.py delete mode 100644 tests/archive/test_128_16_nosoftmax.py delete mode 100644 tests/archive/test_128_16_pAtS.py delete mode 100644 tests/archive/test_128_16_pvlayout.py delete mode 100644 tests/archive/test_128_16_pvpack.py delete mode 100644 tests/archive/test_128_16_pvwrite.py delete mode 100644 tests/archive/test_128_16_qkread.py delete mode 100644 tests/archive/test_128_16_smem.py delete mode 100644 tests/archive/test_128_16_stepA.py delete mode 100644 tests/archive/test_128_16_stepB.py delete mode 100644 tests/archive/test_128_16_stepC.py delete mode 100644 tests/archive/test_128_16_stepD.py delete mode 100644 tests/archive/test_128_16_stepE.py delete mode 100644 tests/archive/test_128_16_tiler.py delete mode 100644 tests/archive/test_128_16_v8.py delete mode 100644 tests/archive/test_128_16_zeropad.py delete mode 100644 tests/archive/test_128_32_ctafix.py delete mode 100644 tests/archive/test_128_32_ctafix2.py delete mode 100644 tests/archive/test_128_32_native.py delete mode 100644 tests/archive/test_128_32_vdiag.py delete mode 100644 tests/archive/test_128_32_zeropad.py delete mode 100644 tests/archive/test_afrag_roundtrip.py delete mode 100644 tests/archive/test_attention.py delete mode 100644 tests/archive/test_attention_path_b200.py delete mode 100644 tests/archive/test_b_afrag2.py delete mode 100644 tests/archive/test_b_layout.py delete mode 100644 tests/archive/test_bf16_elemwise.py delete mode 100644 tests/archive/test_bf16_pack.py delete mode 100644 tests/archive/test_bf16_recast_full.py delete mode 100644 tests/archive/test_bf16_recast_simple.py delete mode 100644 tests/archive/test_blackwell_attn_b200.py delete mode 100644 tests/archive/test_cache.py delete mode 100644 tests/archive/test_compile_custom_op.py delete mode 100644 tests/archive/test_csa_attention_b200.py delete mode 100644 tests/archive/test_csa_sparse_attn_b200.py delete mode 100644 tests/archive/test_custom_op.py delete mode 100644 tests/archive/test_decode_attention_b200.py delete mode 100644 tests/archive/test_decode_pipeline.py delete mode 100644 tests/archive/test_decode_vs_prefill_b200.py delete mode 100644 tests/archive/test_dense_router.py delete mode 100644 tests/archive/test_diag_layout.py delete mode 100644 tests/archive/test_diag_multitile.py delete mode 100644 tests/archive/test_diag_permute.py delete mode 100644 tests/archive/test_diag_smem_layout.py delete mode 100644 tests/archive/test_diag_v_mma128.py delete mode 100644 tests/archive/test_diag_v_ones.py delete mode 100644 tests/archive/test_diag_v_truncid.py delete mode 100644 tests/archive/test_e2e_decode_b200.py delete mode 100644 tests/archive/test_error_pattern.py delete mode 100644 tests/archive/test_fmha_pipeline.py delete mode 100644 tests/archive/test_fmha_v1.py delete mode 100644 tests/archive/test_fmha_v2.py delete mode 100644 tests/archive/test_fmha_v2_fixed.py delete mode 100644 tests/archive/test_fmha_v3_12w.py delete mode 100644 tests/archive/test_fmha_v3_debug.py delete mode 100644 tests/archive/test_fmha_v3_diag.py delete mode 100644 tests/archive/test_fmha_v3_diag_fixed.py delete mode 100644 tests/archive/test_full_layer_b200.py delete mode 100644 tests/archive/test_full_layer_nan_b200.py delete mode 100644 tests/archive/test_full_model_b200.py delete mode 100644 tests/archive/test_fused_step1.py delete mode 100644 tests/archive/test_inspect_types.py delete mode 100644 tests/archive/test_interleave.py delete mode 100644 tests/archive/test_interleave_gemm.py delete mode 100644 tests/archive/test_inv_rope.py delete mode 100644 tests/archive/test_kv_cache_b200.py delete mode 100644 tests/archive/test_layer_schedule.py delete mode 100644 tests/archive/test_layout_compare.py delete mode 100644 tests/archive/test_mma_si_only.py delete mode 100644 tests/archive/test_mma_si_pv.py delete mode 100644 tests/archive/test_model_forward_b200.py delete mode 100644 tests/archive/test_moe_nan_b200.py delete mode 100644 tests/archive/test_moe_runner_nan_b200.py delete mode 100644 tests/archive/test_multilayer.py delete mode 100644 tests/archive/test_nvfp4_attention_b200.py delete mode 100644 tests/archive/test_nvfp4_attn_gemm_b200.py delete mode 100644 tests/archive/test_nvfp4_mapper.py delete mode 100644 tests/archive/test_o_projection.py delete mode 100644 tests/archive/test_o_projection_b200.py delete mode 100644 tests/archive/test_packing_diag.py delete mode 100644 tests/archive/test_pair_swap.py delete mode 100644 tests/archive/test_pair_swap2.py delete mode 100644 tests/archive/test_pipeline_real_weights.py delete mode 100644 tests/archive/test_pv64.py delete mode 100644 tests/archive/test_pv64_fmha_v.py delete mode 100644 tests/archive/test_pv64_kmajor_v.py delete mode 100644 tests/archive/test_pv64_no_softmax.py delete mode 100644 tests/archive/test_pv64_nosoftmax_fmha_v.py delete mode 100644 tests/archive/test_pv_diag.py delete mode 100644 tests/archive/test_pv_mma_mn_major.py delete mode 100644 tests/archive/test_quick_rand.py delete mode 100644 tests/archive/test_recast_minimal.py delete mode 100644 tests/archive/test_ref_minimal.py delete mode 100644 tests/archive/test_reference_fmha.py delete mode 100644 tests/archive/test_rope_kv_b200.py delete mode 100644 tests/archive/test_router.py delete mode 100644 tests/archive/test_runner_vs_pipeline.py delete mode 100644 tests/archive/test_scale_assembly.py delete mode 100644 tests/archive/test_scale_debug.py delete mode 100644 tests/archive/test_shared_expert.py delete mode 100644 tests/archive/test_silu_step1.py delete mode 100644 tests/archive/test_softmax_only.py delete mode 100644 tests/archive/test_softmax_store_debug.py delete mode 100644 tests/archive/test_sparse_attn_b200.py delete mode 100644 tests/archive/test_sparse_decode.py delete mode 100644 tests/archive/test_stage_a_copy.py delete mode 100644 tests/archive/test_stage_a_minimal.py delete mode 100644 tests/archive/test_stage_a_pv_created.py delete mode 100644 tests/archive/test_stage_a_pv_param.py delete mode 100644 tests/archive/test_stage_a_qk.py delete mode 100644 tests/archive/test_stage_a_v2.py delete mode 100644 tests/archive/test_stage_a_with_pv_mma.py delete mode 100644 tests/archive/test_stage_b_afrag.py delete mode 100644 tests/archive/test_stage_b_afrag2.py delete mode 100644 tests/archive/test_stage_b_debug.py delete mode 100644 tests/archive/test_stage_b_debug2.py delete mode 100644 tests/archive/test_stage_b_debug3.py delete mode 100644 tests/archive/test_stage_b_debug4.py delete mode 100644 tests/archive/test_stage_b_diag.py delete mode 100644 tests/archive/test_stage_b_final.py delete mode 100644 tests/archive/test_stage_b_identity.py delete mode 100644 tests/archive/test_stage_b_minimal.py delete mode 100644 tests/archive/test_stage_b_ntile_v1.py delete mode 100644 tests/archive/test_stage_b_ntile_v3.py delete mode 100644 tests/archive/test_stage_b_ntile_v5.py delete mode 100644 tests/archive/test_stage_b_ntile_v6.py delete mode 100644 tests/archive/test_stage_b_ntile_v7.py delete mode 100644 tests/archive/test_stage_b_ntile_v8.py delete mode 100644 tests/archive/test_stage_b_pipeline_only.py delete mode 100644 tests/archive/test_stage_b_v1.py delete mode 100644 tests/archive/test_stage_b_v10.py delete mode 100644 tests/archive/test_stage_b_v11.py delete mode 100644 tests/archive/test_stage_b_v11b.py delete mode 100644 tests/archive/test_stage_b_v12.py delete mode 100644 tests/archive/test_stage_b_v13.py delete mode 100644 tests/archive/test_stage_b_v14.py delete mode 100644 tests/archive/test_stage_b_v16.py delete mode 100644 tests/archive/test_stage_b_v17.py delete mode 100644 tests/archive/test_stage_b_v18.py delete mode 100644 tests/archive/test_stage_b_v19.py delete mode 100644 tests/archive/test_stage_b_v2.py delete mode 100644 tests/archive/test_stage_b_v20.py delete mode 100644 tests/archive/test_stage_b_v22.py delete mode 100644 tests/archive/test_stage_b_v22_bug1fix.py delete mode 100644 tests/archive/test_stage_b_v23.py delete mode 100644 tests/archive/test_stage_b_v24.py delete mode 100644 tests/archive/test_stage_b_v25.py delete mode 100644 tests/archive/test_stage_b_v26.py delete mode 100644 tests/archive/test_stage_b_v27.py delete mode 100644 tests/archive/test_stage_b_v28.py delete mode 100644 tests/archive/test_stage_b_v29.py delete mode 100644 tests/archive/test_stage_b_v3.py delete mode 100644 tests/archive/test_stage_b_v30.py delete mode 100644 tests/archive/test_stage_b_v4.py delete mode 100644 tests/archive/test_stage_b_v5.py delete mode 100644 tests/archive/test_stage_b_v6.py delete mode 100644 tests/archive/test_stage_b_v7.py delete mode 100644 tests/archive/test_stage_b_v7_rep128.py delete mode 100644 tests/archive/test_stage_b_v7_rep16.py delete mode 100644 tests/archive/test_stage_b_v7_rep64.py delete mode 100644 tests/archive/test_stage_b_v7_rep8.py delete mode 100644 tests/archive/test_stage_b_v8.py delete mode 100644 tests/archive/test_stage_b_v8b.py delete mode 100644 tests/archive/test_stage_b_v9.py delete mode 100644 tests/archive/test_step2_subtile.py delete mode 100644 tests/archive/test_step2_subtile_v2.py delete mode 100644 tests/archive/test_store_verify.py delete mode 100644 tests/archive/test_store_verify2.py delete mode 100644 tests/archive/test_tma_coord.py delete mode 100644 tests/archive/test_tmem_addressing.py delete mode 100644 tests/archive/test_tmem_col2.py delete mode 100644 tests/archive/test_tmem_col3.py delete mode 100644 tests/archive/test_tmem_col4.py delete mode 100644 tests/archive/test_tmem_col5.py delete mode 100644 tests/archive/test_tmem_col5_16.py delete mode 100644 tests/archive/test_tmem_col5_32.py delete mode 100644 tests/archive/test_tmem_col_offset.py delete mode 100644 tests/archive/test_tmem_copy_roundtrip.py delete mode 100644 tests/archive/test_tmem_debug.py delete mode 100644 tests/archive/test_tmem_debug2.py delete mode 100644 tests/archive/test_tmem_fp32_roundtrip.py delete mode 100644 tests/archive/test_tmem_layout_diag.py delete mode 100644 tests/archive/test_tmem_pure_fp32.py delete mode 100644 tests/archive/test_uniform_fp4.py delete mode 100644 tests/archive/test_v28c_noepi.py delete mode 100644 tests/archive/test_v4_attention_b200.py delete mode 100644 tests/archive/test_v_mode_fix.py delete mode 100644 tests/archive/test_vllm_codepaths_b200.py delete mode 100644 tests/archive/test_vsmem_diag.py delete mode 100644 tests/archive/test_warmup_gs.py delete mode 100644 tests/archive/test_wo_a.py delete mode 100644 tests/archive/test_wo_a_bmm.py delete mode 100644 tests/archive/unit_test_128_128_vdiag.py delete mode 100644 tests/archive/unit_test_fmha_v3_debug.py delete mode 100644 tests/archive/unit_test_fmha_v3_fixed_v.py delete mode 100644 tests/archive/unit_test_fmha_v3_noop_c9.py delete mode 100644 tests/archive/unit_test_fmha_v3_per_row.py delete mode 100644 tests/archive/unit_test_fmha_v3_per_row_min.py delete mode 100644 tests/archive/unit_test_fmha_v3_proper.py delete mode 100644 tests/archive/unit_test_fmha_v3_pva_c9.py delete mode 100644 tests/archive/unit_test_fmha_v3_scalar.py delete mode 100644 tests/archive/unit_test_fmha_v3_shapes.py delete mode 100644 tests/archive/unit_test_fmha_v3_softmax.py delete mode 100644 tests/archive/unit_test_fmha_v3_stage_c.py delete mode 100644 tests/archive/unit_test_fmha_v3_stage_c2.py delete mode 100644 tests/archive/unit_test_fmha_v3_stage_c_full.py delete mode 100644 tests/archive/unit_test_fmha_v3_stage_c_min.py delete mode 100644 tests/archive/unit_test_fmha_v3_vec_c9.py delete mode 100644 tests/archive/unit_test_pv64_with_softmax.py delete mode 100644 tests/archive/unit_test_qk_softmax.py delete mode 100644 tests/archive/unit_test_qkonly.py delete mode 100644 tests/fmha_v3_stage_c_example10.py delete mode 100644 tests/fmha_v3_stage_c_example9.py diff --git a/README.md b/README.md index 04069365..0ddbe710 100644 --- a/README.md +++ b/README.md @@ -199,8 +199,7 @@ dsv4/ |------|-------|--------| | `test_fmha_v3.py` | A+B | ✅ Full QK→identity softmax→PV, cosine 0.999999 | | `test_fmha_v3_12w.py` | A+B | ✅ 12-warp QK→PV, cosine 0.999999 | -| `test_fmha_v3_stage_c_full.py` | C | ✅ Real online softmax + O normalization, cosine 0.993-0.996 | -| `test_fmha_v3_stage_c_min.py` | C | 🔨 Early 12-warp pipeline (broken pipeline state) | +| `test_fmha_v3_stage_c.py` | C | ⚠️ Real online softmax + normalize, n=128 cos 0.973 (TMEM layout mismatch), n=256 cos 0.793 | | `test_pv64_with_softmax.py` | B | ✅ (128,64) PV, single AB pipeline | | `test_128_128_vdiag.py` | A+B | ✅ (128,128) PV baseline | | `test_qkonly.py` | A | ✅ QK with split Q/KV pipelines | @@ -220,10 +219,6 @@ dsv4/ | `test_interleave_gemm.py` | Interleaved GEMM correctness | | `test_fused_step1.py` | Fused SwiGLU GEMM | -### Archived Tests - -`tests/archive/` contains ~190 debug files from Stages A/B. Not maintained. Can be deleted. - --- ## Test Harness @@ -394,6 +389,53 @@ Col 128+: O (PV acc, 64 FP32, rescale via Ld32x32bOp Repetition(16)) --- +## Stage D: Full Decode Attention (TODO) + +### What Stage D Needs + +The FMHA pipeline (Stages A-C) processes a single query row against K/V. Stage D makes it a **real attention kernel** by adding: + +1. **Paged KV cache** — read K/V from a paged cache (not contiguous GMEM). Each page is a fixed-size block (e.g., 128 tokens). Page tables map logical → physical addresses. This is how vLLM handles variable sequence lengths without reallocation. + +2. **Multi-query attention (MQA) / Grouped-query attention (GQA)** — DSV4 has 128 heads sharing a (T, 512) KV latent. All heads share the same K/V, so Q is (128, 512) but K/V are (T, 512). One K/V load, 128 Q heads. The current single-head kernel needs a Q-head loop. + +3. **Causal mask** — for autoregressive decode, each query position can only attend to K/V positions ≤ its own. This is a simple upper-triangular mask on the QK scores (or just limiting the K/V range loaded per query row). + +4. **CSA/HCA sparse attention** — DSV4 uses compressed sparse attention (m=4) and heavily compressed attention (m=128). The indexer selects top-k blocks, and attention only runs on those blocks. This changes the KV tile iteration from sequential to sparse. + +5. **SWA branch + sink merge** — Sliding window attention (window=128) runs alongside sparse attention. Sink weights combine the two branches: `O = sink * O_sparse + (1-sink) * O_swa`. + +### Architecture Plan + +``` +Input: Q (num_heads, 512, 1), K/V from paged cache + │ + ├─ Q-head loop (128 heads, shared K/V) + │ │ + │ ├─ Causal mask: only load K/V up to current position + │ │ + │ ├─ CSA path: load top-k compressed KV blocks → sparse FMHA + │ │ + │ ├─ SWA path: load last 128 tokens → windowed FMHA + │ │ + │ └─ Sink merge: O = sink * O_csa + (1-sink) * O_swa + │ + └─ Output: O (num_heads, 512, 1) → inv RoPE → o_a BMM → o_b projection +``` + +### Reference Files + +- KV cache spec: `dsv4/cache/` (stubs) +- CSA/HCA attention: `dsv4/reference/csa_attention.py` +- vLLM Blackwell attention backend: `/root/dsv4-nvfp4-workspace/vllm/vllm/attention/backends/` + +### Dependencies + +- Stage C precision fix (correction_epilog) should be resolved before production, but doesn't block Stage D development +- Paged KV cache read/write kernels (Stage D.1) are the first concrete deliverable + +--- + ## Environment - Server: root@45.76.247.107 (B200, 180 GiB HBM3e per GPU) diff --git a/tests/archive/debug_output.py b/tests/archive/debug_output.py deleted file mode 100644 index f7bef38b..00000000 --- a/tests/archive/debug_output.py +++ /dev/null @@ -1,70 +0,0 @@ -"""Debug: Compare runner output vs reference pipeline output. -Focus on whether the scale assembly + GEMM produces correct values.""" -import torch -import sys -sys.path.insert(0, '/root/nvfp4-megamoe-kernel/cutedsl') -sys.path.insert(0, '/root/nvfp4-megamoe-kernel/vllm') - -from cutedsl.reference.moe_pipeline import moe_pipeline -from vllm.nvfp4_cutedsl import Nvfp4MoE - -torch.cuda.set_device(0) - -# Load real model weights for layer 0 -from cutedsl.weight_loader import load_layer_weights -weights = load_layer_weights(layer_idx=0, num_experts=3) - -# Run reference pipeline with dynamic gs -ref_out = moe_pipeline( - hidden_states=torch.randn(4, 256, dtype=torch.bfloat16, device='cuda'), - topk_weights=torch.ones(4, 2, dtype=torch.float32, device='cuda') / 2, - topk_ids=torch.tensor([[0,1],[0,1],[0,1],[0,1]], dtype=torch.int64, device='cuda'), - l1_fp4=weights['l1_fp4'], - l1_sf=weights['l1_sf'], - l1_gs=weights['l1_gs'], - l2_fp4=weights['l2_fp4'], - l2_sf=weights['l2_sf'], - l2_gs=weights['l2_gs'], - num_experts=3, - hidden_size=256, - intermediate_size=512, -) - -print(f"Reference output: amax={ref_out.amax().item():.4f} mean={ref_out.mean().item():.4f}") - -# Run runner with warmup gs -runner = Nvfp4MoE( - num_experts=3, hidden_size=256, intermediate_size=512, - max_num_tokens=4, top_k=2, device='cuda' -) -# Set weights directly -runner.l1_fp4 = weights['l1_fp4'] -runner.l1_sf = weights['l1_sf'] -runner.l1_gs = weights['l1_gs'] -runner.l2_fp4 = weights['l2_fp4'] -runner.l2_sf = weights['l2_sf'] -runner.l2_gs = weights['l2_gs'] - -# Compute warmup gs -hs = torch.randn(4, 256, dtype=torch.bfloat16, device='cuda') -tw = torch.ones(4, 2, dtype=torch.float32, device='cuda') / 2 -ti = torch.tensor([[0,1],[0,1],[0,1],[0,1]], dtype=torch.int64, device='cuda') -runner.compute_activation_global_scales(hs, tw, ti) -print(f"Warmup gs: L1={runner._l1_activation_global_scale} L2={runner._l2_activation_global_scale}") - -# Run with same input as reference -runner_out = runner.run(hs, tw, ti) -print(f"Runner output: amax={runner_out.amax().item():.4f} mean={runner_out.mean().item():.4f}") - -# Cosine similarity -cos = torch.nn.functional.cosine_similarity(ref_out.flatten().unsqueeze(0), runner_out.flatten().unsqueeze(0)).item() -print(f"Cosine similarity: {cos:.6f}") - -# Check for NaN/Inf -print(f"Runner NaN: {torch.isnan(runner_out).any().item()} Inf: {torch.isinf(runner_out).any().item()}") -print(f"Ref NaN: {torch.isnan(ref_out).any().item()} Inf: {torch.isinf(ref_out).any().item()}") - -# Per-token comparison -for i in range(4): - cos_i = torch.nn.functional.cosine_similarity(ref_out[i].unsqueeze(0), runner_out[i].unsqueeze(0)).item() - print(f" Token {i}: cosine={cos_i:.6f} ref_max={ref_out[i].amax().item():.4f} run_max={runner_out[i].amax().item():.4f}") diff --git a/tests/archive/debug_stages.py b/tests/archive/debug_stages.py deleted file mode 100644 index 2055f994..00000000 --- a/tests/archive/debug_stages.py +++ /dev/null @@ -1,59 +0,0 @@ -""" -Debug: test each stage independently. -1. Run Stage A (Q @ K^T only) — should give cosine 0.999 -2. Run Stage B minimal (two MMAs, no softmax) — should give NaN or garbage -3. Run Stage B pipeline-only (pipeline but no ld/st) — should give NaN or garbage -4. Run Stage B full (identity softmax) — should give correct (Q@K^T)@V -""" -import torch -import cutlass.cute as cute -import cutlass.torch as ct -import cuda.bindings.driver as cuda - -torch.manual_seed(42) -m, n, k = 128, 128, 128 -q = torch.randn(m, k, 1, dtype=torch.bfloat16, device='cuda') -kv = torch.randn(n, k, 1, dtype=torch.bfloat16, device='cuda') - -qf = q[:,:,0].float(); kvf = kv[:,:,0].float() -ref_qkt = qf @ kvf.T -ref_qktv = ref_qkt @ kvf - -print(f"Q shape: {q.shape}, KV shape: {kv.shape}") -print(f"Q@K^T shape: {ref_qkt.shape}, (Q@K^T)@V shape: {ref_qktv.shape}") -print(f"Q@K^T range: [{ref_qkt.min():.2f}, {ref_qkt.max():.2f}]") -print(f"(Q@K^T)@V range: [{ref_qktv.min():.2f}, {ref_qktv.max():.2f}]") - -# Test Stage A first -from test_stage_a_v2 import StageAQKTKernel -c_a = torch.zeros(m, n, 1, dtype=torch.bfloat16, device='cuda') -mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) -mK = ct.from_dlpack(kv).mark_layout_dynamic(leading_dim=ct.get_leading_dim(kv)) -mC = ct.from_dlpack(c_a).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_a)) -stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - -kernel_a = StageAQKTKernel(mma_tiler_mn=(128, 128)) -compiled_a = cute.compile(kernel_a, mQ, mK, mC, stream) -compiled_a(mQ, mK, mC, stream) -torch.cuda.synchronize() -out_a = c_a[:,:,0].float() -cos_a = torch.nn.functional.cosine_similarity(out_a.flatten().unsqueeze(0), ref_qkt.flatten().unsqueeze(0)).item() -print(f"\nStage A (Q@K^T): cosine = {cos_a:.6f} {'✅' if cos_a > 0.99 else '❌'}") - -# Test Stage B v7 (identity softmax) -from test_stage_b_v7 import StageBIdentitySoftmax -c_b = torch.zeros(m, n, 1, dtype=torch.bfloat16, device='cuda') -mC2 = ct.from_dlpack(c_b).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_b)) -kernel_b = StageBIdentitySoftmax(mma_tiler_mn=(128, 128)) -compiled_b = cute.compile(kernel_b, mQ, mK, mC2, stream) -compiled_b(mQ, mK, mC2, stream) -torch.cuda.synchronize() -out_b = c_b[:,:,0].float() -cos_b = torch.nn.functional.cosine_similarity(out_b.flatten().unsqueeze(0), ref_qktv.flatten().unsqueeze(0)).item() -has_nan = torch.isnan(out_b).any().item() -print(f"Stage B (identity softmax): cosine = {cos_b:.6f}, has_nan = {has_nan} {'✅' if cos_b > 0.99 else '❌'}") - -# Check: is the output close to Q@K^T (not Q@K^T@V)? -cos_b_qkt = torch.nn.functional.cosine_similarity(out_b.flatten().unsqueeze(0), ref_qkt.flatten().unsqueeze(0)).item() -print(f" vs Q@K^T: cosine = {cos_b_qkt:.6f} (should be ~0 if it's Q@K^T@V)") -print(f" Output range: [{out_b.nan_to_num().min():.2f}, {out_b.nan_to_num().max():.2f}]") diff --git a/tests/archive/diag_layouts.py b/tests/archive/diag_layouts.py deleted file mode 100644 index 1f7eaa1b..00000000 --- a/tests/archive/diag_layouts.py +++ /dev/null @@ -1,85 +0,0 @@ -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils -from cutlass.cute.nvgpu import tcgen05 -from cutlass import Float32, BFloat16 -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset - -a_dtype = BFloat16; b_dtype = BFloat16 -from cutlass.cute.nvgpu import OperandMajorMode -a_major = OperandMajorMode.K -b_major = OperandMajorMode.K -mma_tiler_mn = (128, 128) - -qk_mma = utils.sm100.make_trivial_tiled_mma( - a_dtype, b_dtype, a_major, b_major, Float32, tcgen05.CtaGroup.ONE, mma_tiler_mn, tcgen05.OperandSource.SMEM) -pv_mma = utils.sm100.make_trivial_tiled_mma( - a_dtype, b_dtype, OperandMajorMode.K, b_major, Float32, tcgen05.CtaGroup.ONE, mma_tiler_mn, tcgen05.OperandSource.TMEM) - -qk_thr = qk_mma.get_slice(0) -qk_acc_shape = qk_thr.partition_shape_C(mma_tiler_mn) -tStS = qk_thr.make_fragment_C(qk_acc_shape) - -pv_thr = pv_mma.get_slice(0) -pv_acc_shape = pv_thr.partition_shape_C(mma_tiler_mn) -tOtO = pv_thr.make_fragment_C(pv_acc_shape) - -qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) -qk_mma_tiler = (*mma_tiler_mn, qk_inst_k * 4) -pv_inst_k = cute.size(pv_mma.shape_mnk, mode=[2]) -pv_mma_tiler = (*mma_tiler_mn, pv_inst_k * 4) - -p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, pv_mma_tiler, BFloat16, 1) -tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) -tOrP_base = pv_thr.make_fragment_A(tP) -tOrP = tOrP_base[(None, None, None, 0)] -tOrP0 = cute.make_tensor( - tOrP.iterator + Float32.width // BFloat16.width * 32, - tOrP.layout) - -print('=== Layout diagnostics ===') -print('tStS.layout:', tStS.layout) -print('tStS.size:', cute.size(tStS)) -print('tStS s_cols:', find_tmem_tensor_col_offset(tStS)) -print() -print('tOtO.layout:', tOtO.layout) -print('tOtO.size:', cute.size(tOtO)) -print('tOtO o_cols:', find_tmem_tensor_col_offset(tOtO)) -print() -print('tOrP.layout:', tOrP.layout) -print('tOrP.size:', cute.size(tOrP)) -print('tOrP0.layout:', tOrP0.layout) -print('tOrP0.size:', cute.size(tOrP0)) -print() - -tilePlikeFP32 = 128 * 16 // 32 -tStS_P_layout = cute.composition(tStS.layout, cute.make_layout((128, tilePlikeFP32))) -print('tStS_P_layout:', tStS_P_layout) -print() - -# LOAD -tmem_load_atom = cute.make_copy_atom( - tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), Float32) -tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS) -thr_load = tiled_tmem_load.get_slice(0) -cS = cute.make_identity_tensor((128, 128)) -tScS = qk_thr.partition_C(cS) -tTMEM_LOADcS = thr_load.partition_D(tScS) -print('LOAD tTMEM_LOADcS.shape:', tTMEM_LOADcS.shape) -print('LOAD per-thread elements:', cute.size(tTMEM_LOADcS)) - -# STORE (composition) -tStS_P = cute.make_tensor(tStS.iterator + 32, tStS_P_layout) -tmem_store_atom = cute.make_copy_atom( - tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), Float32) -tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStS_P) -thr_store = tiled_tmem_store.get_slice(0) -tTMEM_STOREcS = thr_store.partition_S(cute.make_identity_tensor(tStS_P.shape)) -print('STORE tTMEM_STOREcS.shape:', tTMEM_STOREcS.shape) -print('STORE per-thread elements:', cute.size(tTMEM_STOREcS)) - -# What about the tOrP0 shape for store? -print() -print('tOrP0.shape:', tOrP0.shape if hasattr(tOrP0, 'shape') else 'N/A') -# tOrP0 is BF16 so we'd need a BF16 store atom - but cute.copy requires equal bit widths -# The F32 store to a BF16 target doesn't work either -# This is the fundamental tension diff --git a/tests/archive/diag_tma_layout.py b/tests/archive/diag_tma_layout.py deleted file mode 100644 index bd781763..00000000 --- a/tests/archive/diag_tma_layout.py +++ /dev/null @@ -1,116 +0,0 @@ -"""Diagnostic: Check tBgK layout strides to see if the GMEM tile dim is degenerate.""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32 -from cutlass.utils import LayoutEnum -import cutlass.torch as ct -import math - -HEAD_DIM = 64 -n = 256 - -q = torch.randn(128, HEAD_DIM, 1, dtype=torch.bfloat16, device='cuda') -k = torch.randn(n, HEAD_DIM, 1, dtype=torch.bfloat16, device='cuda') -v = torch.randn(n, HEAD_DIM, dtype=torch.bfloat16, device='cuda') -v_kernel = v.unsqueeze(-1) - -mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) -mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) -mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel)) - -qk_mma = utils.sm100.make_trivial_tiled_mma(BFloat16, BFloat16, cute.nvgpu.OperandMajorMode.K, cute.nvgpu.OperandMajorMode.K, Float32, tcgen05.CtaGroup.ONE, (128,128), tcgen05.OperandSource.SMEM) -pv_mma = utils.sm100.make_trivial_tiled_mma(BFloat16, BFloat16, cute.nvgpu.OperandMajorMode.K, cute.nvgpu.OperandMajorMode.MN, Float32, tcgen05.CtaGroup.ONE, (128,HEAD_DIM), tcgen05.OperandSource.TMEM) - -qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) -qk_mma_tiler = (128, 128, qk_ik * 4) -pv_ik = cute.size(pv_mma.shape_mnk, mode=[2]) -pv_mma_tiler = (128, HEAD_DIM, pv_ik * (128 // pv_ik)) -cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - -kv_stage = 2; q_stage = 1 -k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, qk_mma_tiler, BFloat16, kv_stage) -v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, pv_mma_tiler, BFloat16, kv_stage) - -k_s = cute.slice_(k_smem_s,(None,None,None,0)) -v_s = cute.slice_(v_smem_s,(None,None,None,0)) - -tma_k, mK_tma = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(cluster_layout_vmnk.shape, qk_mma.thr_id), - mK, k_s, qk_mma_tiler, qk_mma, cluster_layout_vmnk.shape -) -tma_v, mV_tma = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(cluster_layout_vmnk.shape, pv_mma.thr_id), - mV, v_s, pv_mma_tiler, pv_mma, cluster_layout_vmnk.shape -) - -gK = cute.local_tile(mK_tma, cute.slice_(qk_mma_tiler,(0,None,None)),(None,None,None)) -gV = cute.local_tile(mV_tma, cute.slice_(pv_mma_tiler,(0,None,None)),(None,None,None)) - -print(f'gK shape: {cute.shape(gK)}') -print(f'gK layout: {gK.layout}') -print(f'gK stride: {gK.layout.stride}') -print(f'gK size per mode: {[cute.size(gK, mode=[i]) for i in range(len(cute.shape(gK)))]}') -print() -print(f'gV shape: {cute.shape(gV)}') -print(f'gV layout: {gV.layout}') -print(f'gV stride: {gV.layout.stride}') -print() - -qk_thr = qk_mma.get_slice(0) -pv_thr = pv_mma.get_slice(0) -tCgK = qk_thr.partition_B(gK) -tCgV = pv_thr.partition_B(gV) - -print(f'tCgK shape: {cute.shape(tCgK)}') -print(f'tCgK layout: {tCgK.layout}') -print(f'tCgK stride: {tCgK.layout.stride}') -print() -print(f'tCgV shape: {cute.shape(tCgV)}') -print(f'tCgV layout: {tCgV.layout}') -print(f'tCgV stride: {tCgV.layout.stride}') -print() - -sK = cute.make_tensor(BFloat16, k_s) -sV = cute.make_tensor(BFloat16, v_s) -b_lay = cute.make_layout(cute.slice_(cluster_layout_vmnk,(0,None,0,0)).shape) - -tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) -tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - -print(f'=== tBgK (K TMA partition) ===') -print(f'shape: {cute.shape(tBgK)}') -print(f'layout: {tBgK.layout}') -print(f'stride: {tBgK.layout.stride}') -print(f'size per mode: {[cute.size(tBgK, mode=[i]) for i in range(len(cute.shape(tBgK)))]}') -print() -print(f'=== tVgV (V TMA partition) ===') -print(f'shape: {cute.shape(tVgV)}') -print(f'layout: {tVgV.layout}') -print(f'stride: {tVgV.layout.stride}') -print(f'size per mode: {[cute.size(tVgV, mode=[i]) for i in range(len(cute.shape(tVgV)))]}') -print() - -# Now check the slices -print(f'=== tBgK after (None,None,0,0) ===') -tBgK_nn = tBgK[(None,None,0,0)] -print(f'shape: {cute.shape(tBgK_nn)}') -print(f'layout: {tBgK_nn.layout}') -print(f'stride: {tBgK_nn.layout.stride}') -print() -print(f'=== tBgK after (None,0,None,0) ===') -tBgK_n0 = tBgK[(None,0,None,0)] -print(f'shape: {cute.shape(tBgK_n0)}') -print(f'layout: {tBgK_n0.layout}') -print(f'stride: {tBgK_n0.layout.stride}') -print() - -print(f'=== tVgV after (None,0,None,0) ===') -tVgV_n0 = tVgV[(None,0,None,0)] -print(f'shape: {cute.shape(tVgV_n0)}') -print(f'layout: {tVgV_n0.layout}') -print(f'stride: {tVgV_n0.layout.stride}') -print(f'=== tVgV after (None,None,0,0) ===') -tVgV_nn = tVgV[(None,None,0,0)] -print(f'shape: {cute.shape(tVgV_nn)}') -print(f'layout: {tVgV_nn.layout}') -print(f'stride: {tVgV_nn.layout.stride}') diff --git a/tests/archive/diag_tma_shapes.py b/tests/archive/diag_tma_shapes.py deleted file mode 100644 index 0737338e..00000000 --- a/tests/archive/diag_tma_shapes.py +++ /dev/null @@ -1,120 +0,0 @@ -"""Diagnostic: print tma_partition output shapes for Q, K, V tensors.""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -import cuda.bindings.driver as cuda -import cutlass.torch as ct -import math - -HEAD_DIM = 64 - -class DiagShapes: - def __init__(self, s_k=256): - self.s_k = s_k - self.n_kv_tiles = s_k // 128 - self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0,1,2,3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.kv_stage = 2; self.q_stage = 1; self.num_c_stage = 2 - self.scale_softmax = 1.0 / math.sqrt(HEAD_DIM) - self.scale_softmax_log2 = self.scale_softmax * math.log2(math.e) - - def _setup(self, qk_mma, pv_mma): - qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (128, 128, qk_ik * 4) - pv_ik = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (128, HEAD_DIM, pv_ik * (128 // pv_ik)) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0]//cute.size(qk_mma.thr_id.shape), HEAD_DIM, self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - self.q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, self.q_dtype, self.q_stage) - self.k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, self.q_dtype, self.kv_stage) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, self.kv_stage) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - v_fmha = cute.make_tensor( - v.iterator, - cute.make_layout((HEAD_DIM, self.s_k, 1), stride=(1, HEAD_DIM, HEAD_DIM * self.s_k)), - ) - self.v_major = LayoutEnum.from_tensor(v_fmha).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, self.a_major, self.b_major, self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, self.qk_acc_dtype, self.cta_group, (128,HEAD_DIM), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)) - k_s = cute.slice_(self.k_smem_s,(None,None,None,0)) - v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id),q,q_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id),k,k_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v_fmha,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape) - - gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None)) - gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None)) - gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None)) - - print(f"gQ shape: {cute.shape(gQ)}") - print(f"gK shape: {cute.shape(gK)}") - print(f"gV shape: {cute.shape(gV)}") - - qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - tCgV = pv_thr.partition_B(gV) - - print(f"tCgQ shape: {cute.shape(tCgQ)}") - print(f"tCgK shape: {cute.shape(tCgK)}") - print(f"tCgV shape: {cute.shape(tCgV)}") - - a_lay = cute.make_layout(cute.slice_(self.cluster_layout_vmnk,(0,0,None,0)).shape) - b_lay = cute.make_layout(cute.slice_(self.cluster_layout_vmnk,(0,None,0,0)).shape) - sQ = cute.slice_(self.q_smem_s,(None,None,None,0)) - sK = cute.slice_(self.k_smem_s,(None,None,None,0)) - sV = cute.slice_(self.v_smem_s,(None,None,None,0)) - tAsQ,tAgQ = cpasync.tma_partition(tma_q,0,a_lay,cute.group_modes(sQ,0,3),cute.group_modes(tCgQ,0,3)) - tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3)) - tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3)) - - print(f"tAgQ shape: {cute.shape(tAgQ)} stride: {tAgQ.layout.stride}") - print(f"tBgK shape: {cute.shape(tBgK)} stride: {tBgK.layout.stride}") - print(f"tVgV shape: {cute.shape(tVgV)} stride: {tVgV.layout.stride}") - print(f"tAsQ shape: {cute.shape(tAsQ)} stride: {tAsQ.layout.stride}") - print(f"tBsK shape: {cute.shape(tBsK)} stride: {tBsK.layout.stride}") - print(f"tVsV shape: {cute.shape(tVsV)} stride: {tVsV.layout.stride}") - - -def test(): - torch.manual_seed(42) - n = 256 - m, hd = 128, HEAD_DIM - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(n, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(n, hd, dtype=torch.bfloat16, device='cuda') - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - v_kernel = v.unsqueeze(-1) - - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - kernel = DiagShapes(s_k=n) - # Just run __call__ to trigger the prints at trace time - try: - kernel(mQ, mK, mV, mC, stream) - except Exception as e: - print(f"Error (expected — just needed the prints): {e}") - -if __name__ == '__main__': - test() diff --git a/tests/archive/diag_tma_shapes2.py b/tests/archive/diag_tma_shapes2.py deleted file mode 100644 index faa33d0f..00000000 --- a/tests/archive/diag_tma_shapes2.py +++ /dev/null @@ -1,122 +0,0 @@ -"""Diagnostic: print tma_partition output shapes for K, V tensors only.""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct -import math - -HEAD_DIM = 64 - -class DiagShapes: - def __init__(self, s_k=256): - self.s_k = s_k - self.n_kv_tiles = s_k // 128 - self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0,1,2,3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.kv_stage = 2; self.q_stage = 1; self.num_c_stage = 2 - self.scale_softmax = 1.0 / math.sqrt(HEAD_DIM) - self.scale_softmax_log2 = self.scale_softmax * math.log2(math.e) - - def _setup(self, qk_mma, pv_mma): - qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (128, 128, qk_ik * 4) - pv_ik = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (128, HEAD_DIM, pv_ik * (128 // pv_ik)) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0]//cute.size(qk_mma.thr_id.shape), HEAD_DIM, self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - self.q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, self.q_dtype, self.q_stage) - self.k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, self.q_dtype, self.kv_stage) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, self.kv_stage) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - v_fmha = cute.make_tensor( - v.iterator, - cute.make_layout((HEAD_DIM, self.s_k, 1), stride=(1, HEAD_DIM, HEAD_DIM * self.s_k)), - ) - self.v_major = LayoutEnum.from_tensor(v_fmha).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, self.a_major, self.b_major, self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, self.qk_acc_dtype, self.cta_group, (128,HEAD_DIM), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)) - k_s = cute.slice_(self.k_smem_s,(None,None,None,0)) - v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id),q,q_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id),k,k_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v_fmha,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape) - - gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None)) - gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None)) - gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None)) - - print(f"gQ shape: {cute.shape(gQ)}") - print(f"gK shape: {cute.shape(gK)}") - print(f"gV shape: {cute.shape(gV)}") - - qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - tCgV = pv_thr.partition_B(gV) - - print(f"tCgQ shape: {cute.shape(tCgQ)}") - print(f"tCgK shape: {cute.shape(tCgK)}") - print(f"tCgV shape: {cute.shape(tCgV)}") - - a_lay = cute.make_layout(cute.slice_(self.cluster_layout_vmnk,(0,0,None,0)).shape) - b_lay = cute.make_layout(cute.slice_(self.cluster_layout_vmnk,(0,None,0,0)).shape) - - # K and V only (Q tma_partition fails in this context) - tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(k_s,0,3),cute.group_modes(tCgK,0,3)) - tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(v_s,0,3),cute.group_modes(tCgV,0,3)) - - print(f"tBsK shape: {cute.shape(tBsK)} stride: {tBsK.layout.stride}") - print(f"tBgK shape: {cute.shape(tBgK)} stride: {tBgK.layout.stride}") - print(f"tVsV shape: {cute.shape(tVsV)} stride: {tVsV.layout.stride}") - print(f"tVgV shape: {cute.shape(tVgV)} stride: {tVgV.layout.stride}") - - # Try the (None,0,None,0) pre-slice - tBgK_sliced = tBgK[(None,0,None,0)] - tVgV_sliced = tVgV[(None,0,None,0)] - print(f"tBgK after (None,0,None,0) shape: {cute.shape(tBgK_sliced)} stride: {tBgK_sliced.layout.stride}") - print(f"tVgV after (None,0,None,0) shape: {cute.shape(tVgV_sliced)} stride: {tVgV_sliced.layout.stride}") - - -def test(): - torch.manual_seed(42) - n = 256 - m, hd = 128, HEAD_DIM - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(n, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(n, hd, dtype=torch.bfloat16, device='cuda') - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - v_kernel = v.unsqueeze(-1) - - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - kernel = DiagShapes(s_k=n) - try: - kernel(mQ, mK, mV, mC, stream) - except Exception as e: - print(f"Error (expected — just needed the prints): {e}") - -if __name__ == '__main__': - test() diff --git a/tests/archive/diag_tma_shapes3.py b/tests/archive/diag_tma_shapes3.py deleted file mode 100644 index 771acf85..00000000 --- a/tests/archive/diag_tma_shapes3.py +++ /dev/null @@ -1,133 +0,0 @@ -"""Diagnostic: print tma_partition output shapes inside @cute.kernel context.""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct -import math - -HEAD_DIM = 64 - -class DiagShapes: - def __init__(self, s_k=256): - self.s_k = s_k - self.n_kv_tiles = s_k // 128 - self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0,1,2,3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.kv_stage = 2; self.q_stage = 1; self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (128, 128, qk_ik * 4) - pv_ik = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (128, HEAD_DIM, pv_ik * (128 // pv_ik)) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0]//cute.size(qk_mma.thr_id.shape), HEAD_DIM, self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - self.q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, self.q_dtype, self.q_stage) - self.k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, self.q_dtype, self.kv_stage) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, self.kv_stage) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - v_fmha = cute.make_tensor(v.iterator, cute.make_layout((HEAD_DIM, self.s_k, 1), stride=(1, HEAD_DIM, HEAD_DIM * self.s_k))) - self.v_major = LayoutEnum.from_tensor(v_fmha).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, self.a_major, self.b_major, self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, self.qk_acc_dtype, self.cta_group, (128,HEAD_DIM), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)) - k_s = cute.slice_(self.k_smem_s,(None,None,None,0)) - v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id),q,q_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id),k,k_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v_fmha,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape) - self._kernel(qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, self.cluster_layout_vmnk, self.q_smem_s, self.k_smem_s, self.v_smem_s).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, cl_vmnk, q_smem_s, k_smem_s, v_smem_s): - tidx,_,_ = cute.arch.thread_idx() - - sQ = cute.make_tensor(element_type=self.q_dtype, layout=q_smem_s.outer, byte_alignment=128, swizzle=q_smem_s.inner) - sK = cute.make_tensor(element_type=self.q_dtype, layout=k_smem_s.outer, byte_alignment=128, swizzle=k_smem_s.inner) - sV = cute.make_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - - q_s = cute.slice_(q_smem_s,(None,None,None,0)) - k_s = cute.slice_(k_smem_s,(None,None,None,0)) - v_s = cute.slice_(v_smem_s,(None,None,None,0)) - - gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None)) - gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None)) - gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None)) - - qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - tCgV = pv_thr.partition_B(gV) - - a_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,0,None,0)).shape) - b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape) - tAsQ,tAgQ = cpasync.tma_partition(tma_q,0,a_lay,cute.group_modes(q_s,0,3),cute.group_modes(tCgQ,0,3)) - tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(k_s,0,3),cute.group_modes(tCgK,0,3)) - tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(v_s,0,3),cute.group_modes(tCgV,0,3)) - - # Print shapes at trace time - print(f"DIAG tAgQ: shape={cute.shape(tAgQ)} stride={tAgQ.layout.stride}") - print(f"DIAG tBgK: shape={cute.shape(tBgK)} stride={tBgK.layout.stride}") - print(f"DIAG tVgV: shape={cute.shape(tVgV)} stride={tVgV.layout.stride}") - print(f"DIAG tAsQ: shape={cute.shape(tAsQ)} stride={tAsQ.layout.stride}") - print(f"DIAG tBsK: shape={cute.shape(tBsK)} stride={tBsK.layout.stride}") - print(f"DIAG tVsV: shape={cute.shape(tVsV)} stride={tVsV.layout.stride}") - - # Try pre-slices and print - tAgQ_s = tAgQ[(None,0,None,0)] - print(f"DIAG tAgQ after (None,0,None,0): shape={cute.shape(tAgQ_s)} stride={tAgQ_s.layout.stride}") - - tBgK_s = tBgK[(None,0,None,0)] - print(f"DIAG tBgK after (None,0,None,0): shape={cute.shape(tBgK_s)} stride={tBgK_s.layout.stride}") - - tVgV_s = tVgV[(None,0,None,0)] - print(f"DIAG tVgV after (None,0,None,0): shape={cute.shape(tVgV_s)} stride={tVgV_s.layout.stride}") - - # Also try old pre-slice - tBgK_old = tBgK[(None,None,0,0)] - print(f"DIAG tBgK after (None,None,0,0): shape={cute.shape(tBgK_old)} stride={tBgK_old.layout.stride}") - - tVgV_old = tVgV[(None,0,None,0)] - print(f"DIAG tVgV after (None,0,None,0): shape={cute.shape(tVgV_old)} stride={tVgV_old.layout.stride}") - - -def test(): - torch.manual_seed(42) - n = 256 - m, hd = 128, HEAD_DIM - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(n, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(n, hd, dtype=torch.bfloat16, device='cuda') - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - v_kernel = v.unsqueeze(-1) - - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - kernel = DiagShapes(s_k=n) - print('Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)), stream) - print('Compiled. Running...', flush=True) - -if __name__ == '__main__': - test() diff --git a/tests/archive/diag_tmem.py b/tests/archive/diag_tmem.py deleted file mode 100644 index 84786eb8..00000000 --- a/tests/archive/diag_tmem.py +++ /dev/null @@ -1,91 +0,0 @@ -"""Diagnostic: Q1, Q2 for Stage B TMEM debugging. -Uses cute.compile with a dummy kernel that prints layout info at JIT time.""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils -from cutlass.cute.nvgpu import tcgen05 -from cutlass import Float32, BFloat16 -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - - -@cute.jit -def diag_tmem(stream: cuda.CUstream): - a_dtype = BFloat16; b_dtype = BFloat16 - a_major = cute.nvgpu.OperandMajorMode.K - b_major = cute.nvgpu.OperandMajorMode.K - - qk_mma = utils.sm100.make_trivial_tiled_mma( - a_dtype, b_dtype, a_major, b_major, - Float32, tcgen05.CtaGroup.ONE, (128, 128), tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma( - a_dtype, b_dtype, cute.nvgpu.OperandMajorMode.K, b_major, - Float32, tcgen05.CtaGroup.ONE, (128, 128), tcgen05.OperandSource.TMEM) - - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - pv_inst_k = cute.size(pv_mma.shape_mnk, mode=[2]) - mma_tiler = (128, 128, qk_inst_k * 4) - pv_mma_tiler = (128, 128, pv_inst_k * 4) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - - # Q1: QK accumulator C fragment - qk_acc_shape = qk_thr.partition_shape_C(mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - print(f"=== Q1: QK accumulator C fragment ===") - print(f" tStS.layout = {tStS.layout}") - print(f" cute.size(tStS.layout) = {cute.size(tStS.layout)}") - print(f" cute.cosize(tStS.layout) = {cute.cosize(tStS.layout)}") - print(f" cute.size(mode=[0]) = {cute.size(tStS.layout, mode=[0])}") - print(f" cute.size(mode=[1]) = {cute.size(tStS.layout, mode=[1])}") - s_tmem_cols = find_tmem_tensor_col_offset(tStS) - print(f" find_tmem_tensor_col_offset(tStS) = {s_tmem_cols}") - - # PV accumulator O fragment - pv_acc_shape = pv_thr.partition_shape_C(mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - print(f"=== PV accumulator O fragment ===") - print(f" tOtO.layout = {tOtO.layout}") - print(f" cute.size(tOtO.layout) = {cute.size(tOtO.layout)}") - print(f" cute.cosize(tOtO.layout) = {cute.cosize(tOtO.layout)}") - print(f" cute.size(mode=[0]) = {cute.size(tOtO.layout, mode=[0])}") - print(f" cute.size(mode=[1]) = {cute.size(tOtO.layout, mode=[1])}") - o_tmem_cols = find_tmem_tensor_col_offset(tOtO) - print(f" find_tmem_tensor_col_offset(tOtO) = {o_tmem_cols}") - - # Q2: PV A-fragment (P operand from TMEM) - p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, pv_mma_tiler, BFloat16, 1) - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP_sliced = tOrP_base[(None, None, None, 0)] - print(f"=== Q2: PV A-fragment (P operand) ===") - print(f" tP.layout = {tP.layout}") - print(f" cute.size(tP.layout) = {cute.size(tP.layout)}") - print(f" cute.cosize(tP.layout) = {cute.cosize(tP.layout)}") - print(f" tOrP_sliced.layout = {tOrP_sliced.layout}") - print(f" cute.size(tOrP_sliced.layout) = {cute.size(tOrP_sliced.layout)}") - print(f" cute.cosize(tOrP_sliced.layout) = {cute.cosize(tOrP_sliced.layout)}") - p_tmem_cols = find_tmem_tensor_col_offset(tOrP_sliced) - print(f" find_tmem_tensor_col_offset(tOrP_sliced) = {p_tmem_cols}") - # Decompose 32800 - print(f" 32800 in hex = 0x{32800:04x}") - print(f" 32800 - 0x8000 = {32800 - 0x8000}") - print(f" 32800 & 0x0000FFFF = {32800 & 0x0000FFFF}") - print(f" p_tmem_cols in hex = 0x{p_tmem_cols:04x}") - if isinstance(p_tmem_cols, int): - print(f" p_tmem_cols & 0x0000FFFF = {p_tmem_cols & 0x0000FFFF}") - print(f" p_tmem_cols >> 16 = {p_tmem_cols >> 16}") - - # Staged fragments - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, 1)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, 1)) - print(f"=== Staged fragments ===") - print(f" find_tmem_tensor_col_offset(tCtS_fake) = {find_tmem_tensor_col_offset(tCtS_fake)}") - print(f" find_tmem_tensor_col_offset(tCtO_fake) = {find_tmem_tensor_col_offset(tCtO_fake)}") - print(f" get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake]) = {utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch='sm_100')}") - - -stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) -print("Compiling diagnostics...", flush=True) -compiled = cute.compile(diag_tmem, stream) -print("Done. Results above.", flush=True) diff --git a/tests/archive/fmha_v3_identity_diag.py b/tests/archive/fmha_v3_identity_diag.py deleted file mode 100644 index c6ed5c81..00000000 --- a/tests/archive/fmha_v3_identity_diag.py +++ /dev/null @@ -1,313 +0,0 @@ -""" -FMHA v3 Stage-C Multi-Tile — Identity Softmax Diagnostic. -Strips out real softmax and O normalize to isolate whether the pipeline works. -If this passes, the bug is in the softmax/normalize math. -If this fails, the bug is in the pipeline (sync, TMA, MMA). -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct -import math - -HEAD_DIM = 64 - - -class FmhaV3IdentityDiag: - def __init__(self, s_k=128): - self.s_k = s_k - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0,1,2,3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.num_c_stage = 2 - self.kv_stage = 2; self.q_stage = 1; self.num_c_stage = 2 - self.scale_softmax = 1.0 / math.sqrt(HEAD_DIM) - self.scale_softmax_log2 = self.scale_softmax * math.log2(math.e) - - def _setup(self, qk_mma, pv_mma): - qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (128, 128, qk_ik * 4) - pv_ik = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (128, HEAD_DIM, pv_ik * (128 // pv_ik)) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0]//cute.size(qk_mma.thr_id.shape), HEAD_DIM, self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - self.q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, self.q_dtype, self.q_stage) - self.k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, self.q_dtype, self.kv_stage) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, self.kv_stage) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - qk_thr = qk_mma.get_slice(0); qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - pv_thr = pv_mma.get_slice(0); pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - self.tmem_s0_offset = 0; self.tmem_p0_offset = 32 - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - p_end = self.tmem_p0_offset + p_cols_fp32 - s_cols = self.qk_mma_tiler[1] - o_after = max(s_cols, p_end) - self.tmem_o0_offset = ((o_after + 31) // 32) * 32 - o_cols = find_tmem_tensor_col_offset(tOtO) - total = self.tmem_o0_offset + o_cols - self.num_tmem_alloc_cols = 1 - while self.num_tmem_alloc_cols < total: - self.num_tmem_alloc_cols *= 2 - cta = cute.size(qk_mma.thr_id.shape) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)) - k_s = cute.slice_(self.k_smem_s,(None,None,None,0)) - v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - self.q_tx_bytes = cute.size_in_bytes(self.q_dtype, q_s) * cta - self.kv_tx_bytes = (cute.size_in_bytes(self.q_dtype, k_s) + - cute.size_in_bytes(self.q_dtype, v_s)) * cta - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - v_fmha = cute.make_tensor( - v.iterator, - cute.make_layout( - (HEAD_DIM, self.s_k, 1), - stride=(1, HEAD_DIM, HEAD_DIM * self.s_k), - ), - ) - self.v_major = LayoutEnum.from_tensor(v_fmha).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, self.a_major, self.b_major, self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, self.qk_acc_dtype, self.cta_group, (128,HEAD_DIM), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)); v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id),q,q_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id),k,k_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v_fmha,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape) - epi_s = cute.select(self.c_smem_s,mode=[0,1]) - tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile) - self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.v_smem_s,self.p_tmem_s,self.c_smem_s,self.epi_tile).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, q_smem_s, k_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx,_,_ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k); cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - q_bar: cute.struct.MemRange[cutlass.Int64, self.q_stage*2] - kv_bar: cute.struct.MemRange[cutlass.Int64, self.kv_stage*2] - s_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage*2] - tmem_dealloc: cutlass.Int64; holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - qp,qc = pipeline.PipelineTmaUmma.create(barrier_storage=st.q_bar.data_ptr(),num_stages=self.q_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.q_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - kvp,kvc = pipeline.PipelineTmaUmma.create(barrier_storage=st.kv_bar.data_ptr(),num_stages=self.kv_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.kv_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - s_prod,s_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.s_bar.data_ptr(),num_stages=1,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.epilogue_warp_id))).make_participants() - softmax_done_bar = pipeline.NamedBarrier(barrier_id=3, num_threads=32 + 32*len(self.epilogue_warp_id)) - acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(),num_stages=self.num_acc_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,len(self.epilogue_warp_id)),cta_layout_vmnk=cl_vmnk,defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2,num_threads=32*len((self.mma_warp_id,*self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr,barrier_for_retrieve=tmem_bar,allocator_warp_id=self.epilogue_warp_id[0],is_two_cta=cute.size(qk_mma.thr_id.shape)==2,two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk,is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype,layout=q_smem_s.outer,byte_alignment=128,swizzle=q_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype,layout=k_smem_s.outer,byte_alignment=128,swizzle=k_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype,layout=v_smem_s.outer,byte_alignment=128,swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype,layout=c_smem_s.outer,byte_alignment=128,swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None)) - gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None)) - gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None)) - gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None)) - n_kv_tiles = cute.size(gK, mode=[3]) - - qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - tCgV = pv_thr.partition_B(gV); tCgC = pv_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,0,None,0)).shape) - tAsQ,tAgQ = cpasync.tma_partition(tma_q,0,a_lay,cute.group_modes(sQ,0,3),cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape) - tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3)) - tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3)) - # Use the SAME slice as the working diag test - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,None,0,0)]; tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None,None,None,0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ===== TMA LOAD warp ===== - # Use the SAME pattern as the working diag test - if warp_idx == self.tma_warp_id: - qp.reset(); qh = qp.acquire_and_advance() - cute.copy(tma_q, tAgQ[(None, Int32(0))], tAsQ[(None, qh.index)], tma_bar_ptr=qh.barrier) - qp.tail() - kvp.reset(); pk = kvp.try_acquire() - kv_coord = Int32(0 + 0) - for kt in cutlass.range(self.s_k // 128, unroll=1): - kvh = kvp.acquire_and_advance(pk) - cute.copy(tma_k, tBgK[(None, kv_coord)], tBsK[(None, kvh.index)], tma_bar_ptr=kvh.barrier) - cute.copy(tma_v, tVgV[(None, kv_coord)], tVsV[(None, kvh.index)], tma_bar_ptr=kvh.barrier) - kv_coord += 1 - pk = cutlass.Boolean(1) - kvp.tail() - - # ===== MMA warp ===== - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - qc.reset(); qh = qc.wait_and_advance(); qh.release() - kvc.reset(); pk = kvc.try_wait() - acc_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_stage) - acc_pipe.producer_acquire(acc_st) - n_tiles = self.s_k // 128 - for kt in range(n_tiles): - kvh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - sh = s_prod.acquire_and_advance() - qk_mma.set(tcgen05.Field.ACCUMULATE, False) - for kb in cutlass.range(cute.size(tCrQ, mode=[2]), unroll_full=True): - cute.gemm(qk_mma, tStS0, tCrQ[(None,None,kb,0)], tCrK[(None,None,kb,kvh.index)], tStS0) - qk_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - sh.commit() - softmax_done_bar.arrive_and_wait() - pv_mma.set(tcgen05.Field.ACCUMULATE, kt != 0) - for kb in cutlass.range(cute.size(tOrP0, mode=[2]), unroll_full=True): - cute.gemm(pv_mma, tOtO0, tOrP0[(None,None,kb)], tCrV[(None,None,kb,kvh.index)], tOtO0) - pv_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - kvh.release() - acc_pipe.producer_commit(acc_st); acc_st.advance() - acc_pipe.producer_tail(acc_st) - - # ===== SOFTMAX warps — IDENTITY (no real softmax, just S→P copy) ===== - if warp_idx < self.mma_warp_id: - tmem.allocate(self.num_tmem_alloc_cols) - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) - sfw_idx = tidx % (32 * len(self.epilogue_warp_id)) - - # S load - tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS0) - thr_load = tiled_tmem_load.get_slice(sfw_idx) - tTMEM_LOADtS = thr_load.partition_S(tStS0) - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tTMEM_LOADcS = thr_load.partition_D(tScS) - - # P store - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - tStP_layout = cute.composition(tStS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tStP0 = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStP_layout) - tmem_store_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStP0) - thr_store = tiled_tmem_store.get_slice(sfw_idx) - tTMEM_STOREtP = thr_store.partition_D(tStP0) - tScP_layout = cute.composition(tScS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tScP = cute.make_tensor(tScS.iterator, tScP_layout) - tTMEM_STOREcP = thr_store.partition_S(tScP) - - n_tiles = self.s_k // 128 - for kt in range(n_tiles): - si_handle = s_cons.wait_and_advance() - # Load S - tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS) - cute.arch.fence_view_async_tmem_load() - # Identity: S → BF16 → P - rP_words = cute.make_rmem_tensor(tTMEM_STOREcP.shape, self.qk_acc_dtype) - rP_bf16 = cute.make_tensor(cute.recast_ptr(rP_words.iterator, dtype=self.q_dtype), tTMEM_LOADrS.layout) - frg_cnt = 4 - frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt - tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile)) - rP_bf16_frg = cute.logical_divide(rP_bf16, cute.make_layout(frg_tile)) - for j in range(frg_cnt): - s_vec = tTMEM_LOADrS_frg[None, j].load() - rP_bf16_frg[None, j].store(s_vec.to(self.q_dtype)) - cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP) - cute.arch.fence_view_async_tmem_store() - si_handle.release() - softmax_done_bar.arrive() - - # Epilogue: TMEM -> SMEM -> GMEM via TMA store - tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout) - acc_cons_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage) - c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) - c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp) - acc_cons_st = utils.gemm.sm100.epilogue_tma_store(self, tidx, warp_idx, tma_c, tCtO_base, sC, tCgC, epi_tile, 0, const_expr(lambda x: x), (0,0,0), acc_cons_st, acc_pipe, c_pipe) - c_pipe.producer_tail() - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def test(): - for n in [128, 256]: - torch.manual_seed(42) - m, hd = 128, HEAD_DIM - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(n, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(n, hd, dtype=torch.bfloat16, device='cuda') - v_kernel = v.unsqueeze(-1) - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - - # Reference: identity softmax means P = S, O = P @ V (unnormalized attention) - qf = q[:, :, 0].float() - kf = k[:, :, 0].float() - scale = 1.0 / math.sqrt(hd) - attn = qf @ kf.T * scale - # Identity: no softmax, just raw scaled QK^T @ V - ref = (attn @ v.float()) - - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - kernel = FmhaV3IdentityDiag(s_k=n) - print(f'n={n}: Compiling... [IDENTITY_DIAG]', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - - out = c[:, :, 0].float() - cos = torch.nn.functional.cosine_similarity( - out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0) - ).item() - max_abs = (out - ref).abs().max().item() - n_tiles = n // 128 - print(f'Identity FMHA n={n} ({n_tiles} tiles): ' - f'cos {cos:.6f} max_abs {max_abs:.4f} ' - f'{"PASS" if cos >= 0.99 else "FAIL"}') - if cos < 0.99: - print(f' out[0,:4]={out[0,:4].tolist()}') - print(f' ref[0,:4]={ref[0,:4].tolist()}') - - -if __name__ == '__main__': - test() diff --git a/tests/archive/fmha_v3_real_softmax.py b/tests/archive/fmha_v3_real_softmax.py deleted file mode 100644 index 05f6eb5a..00000000 --- a/tests/archive/fmha_v3_real_softmax.py +++ /dev/null @@ -1,344 +0,0 @@ -""" -FMHA v3 Stage-C — Real Softmax, NO O rescale/normalize in TMEM. - -Strategy: Skip O rescale and TMEM-based normalize (TMEM copy of O corrupts data). -For single-tile (n=128), this gives correct unnormalized output (cos 0.999998). -For multi-tile, the O is not rescaled (missing exp2(old_max - new_max) and 1/row_sum). - -The CUTLASS reference applies O rescale via correction_rescale (TMEM read-modify-write) -and the final 1/row_sum via correction_epilog (applied during GMEM write, NOT TMEM modify). -Our TMEM copy of O doesn't work — likely a CuTeDSL version issue or layout mismatch. -Next step: implement correction_epilog that applies 1/row_sum during GMEM write. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct -import math - -HEAD_DIM = 64 - - -class FmhaV3RealSoftmax: - def __init__(self, s_k=128): - self.s_k = s_k - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0,1,2,3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.num_c_stage = 2 - self.kv_stage = 2; self.q_stage = 1; self.num_c_stage = 2 - self.scale_softmax = 1.0 / math.sqrt(HEAD_DIM) - self.scale_softmax_log2 = self.scale_softmax * math.log2(math.e) - - def _setup(self, qk_mma, pv_mma): - qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (128, 128, qk_ik * 4) - pv_ik = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (128, HEAD_DIM, pv_ik * (128 // pv_ik)) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0]//cute.size(qk_mma.thr_id.shape), HEAD_DIM, self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - self.q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, self.q_dtype, self.q_stage) - self.k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, self.q_dtype, self.kv_stage) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, self.kv_stage) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - qk_thr = qk_mma.get_slice(0); qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - pv_thr = pv_mma.get_slice(0); pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - self.tmem_s0_offset = 0; self.tmem_p0_offset = 32 - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - p_end = self.tmem_p0_offset + p_cols_fp32 - s_cols = self.qk_mma_tiler[1] - o_after = max(s_cols, p_end) - self.tmem_o0_offset = ((o_after + 31) // 32) * 32 - o_cols = find_tmem_tensor_col_offset(pv_thr.make_fragment_C(pv_as)) - total = self.tmem_o0_offset + o_cols - self.num_tmem_alloc_cols = 1 - while self.num_tmem_alloc_cols < total: - self.num_tmem_alloc_cols *= 2 - cta = cute.size(qk_mma.thr_id.shape) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)) - k_s = cute.slice_(self.k_smem_s,(None,None,None,0)) - v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - self.q_tx_bytes = cute.size_in_bytes(self.q_dtype, q_s) * cta - self.kv_tx_bytes = (cute.size_in_bytes(self.q_dtype, k_s) + - cute.size_in_bytes(self.q_dtype, v_s)) * cta - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - v_fmha = cute.make_tensor( - v.iterator, - cute.make_layout( - (HEAD_DIM, self.s_k, 1), - stride=(1, HEAD_DIM, HEAD_DIM * self.s_k), - ), - ) - self.v_major = LayoutEnum.from_tensor(v_fmha).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, self.a_major, self.b_major, self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, self.qk_acc_dtype, self.cta_group, (128,HEAD_DIM), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)); v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id),q,q_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id),k,k_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v_fmha,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape) - epi_s = cute.select(self.c_smem_s,mode=[0,1]) - tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile) - self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.v_smem_s,self.p_tmem_s,self.c_smem_s,self.epi_tile).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, q_smem_s, k_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx,_,_ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k); cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - q_bar: cute.struct.MemRange[cutlass.Int64, self.q_stage*2] - kv_bar: cute.struct.MemRange[cutlass.Int64, self.kv_stage*2] - s_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage*2] - tmem_dealloc: cutlass.Int64; holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - qp,qc = pipeline.PipelineTmaUmma.create(barrier_storage=st.q_bar.data_ptr(),num_stages=self.q_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.q_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - kvp,kvc = pipeline.PipelineTmaUmma.create(barrier_storage=st.kv_bar.data_ptr(),num_stages=self.kv_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.kv_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - s_prod,s_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.s_bar.data_ptr(),num_stages=1,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.epilogue_warp_id))).make_participants() - softmax_done_bar = pipeline.NamedBarrier(barrier_id=3, num_threads=32 + 32*len(self.epilogue_warp_id)) - acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(),num_stages=self.num_acc_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,len(self.epilogue_warp_id)),cta_layout_vmnk=cl_vmnk,defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2,num_threads=32*len((self.mma_warp_id,*self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr,barrier_for_retrieve=tmem_bar,allocator_warp_id=self.epilogue_warp_id[0],is_two_cta=cute.size(qk_mma.thr_id.shape)==2,two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk,is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype,layout=q_smem_s.outer,byte_alignment=128,swizzle=q_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype,layout=k_smem_s.outer,byte_alignment=128,swizzle=k_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype,layout=v_smem_s.outer,byte_alignment=128,swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype,layout=c_smem_s.outer,byte_alignment=128,swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None)) - gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None)) - gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None)) - gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None)) - n_kv_tiles = cute.size(gK, mode=[3]) - - qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - tCgV = pv_thr.partition_B(gV); tCgC = pv_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,0,None,0)).shape) - tAsQ,tAgQ = cpasync.tma_partition(tma_q,0,a_lay,cute.group_modes(sQ,0,3),cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape) - tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3)) - tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,None,0,0)]; tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None,None,None,0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ===== TMA LOAD warp ===== - if warp_idx == self.tma_warp_id: - qp.reset(); qh = qp.acquire_and_advance() - cute.copy(tma_q, tAgQ[(None, Int32(0))], tAsQ[(None, qh.index)], tma_bar_ptr=qh.barrier) - qp.tail() - kvp.reset(); pk = kvp.try_acquire() - kv_coord = Int32(0 + 0) - for kt in cutlass.range(self.s_k // 128, unroll=1): - kvh = kvp.acquire_and_advance(pk) - cute.copy(tma_k, tBgK[(None, kv_coord)], tBsK[(None, kvh.index)], tma_bar_ptr=kvh.barrier) - cute.copy(tma_v, tVgV[(None, kv_coord)], tVsV[(None, kvh.index)], tma_bar_ptr=kvh.barrier) - kv_coord += 1 - pk = cutlass.Boolean(1) - kvp.tail() - - # ===== MMA warp ===== - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - qc.reset(); qh = qc.wait_and_advance(); qh.release() - kvc.reset(); pk = kvc.try_wait() - acc_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_stage) - acc_pipe.producer_acquire(acc_st) - n_tiles = self.s_k // 128 - for kt in range(n_tiles): - kvh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - sh = s_prod.acquire_and_advance() - qk_mma.set(tcgen05.Field.ACCUMULATE, False) - for kb in cutlass.range(cute.size(tCrQ, mode=[2]), unroll_full=True): - cute.gemm(qk_mma, tStS0, tCrQ[(None,None,kb,0)], tCrK[(None,None,kb,kvh.index)], tStS0) - qk_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - sh.commit() - softmax_done_bar.arrive_and_wait() - pv_mma.set(tcgen05.Field.ACCUMULATE, kt != 0) - for kb in cutlass.range(cute.size(tOrP0, mode=[2]), unroll_full=True): - cute.gemm(pv_mma, tOtO0, tOrP0[(None,None,kb)], tCrV[(None,None,kb,kvh.index)], tOtO0) - pv_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_shared() - kvh.release() - acc_pipe.producer_commit(acc_st); acc_st.advance() - acc_pipe.producer_tail(acc_st) - - # ===== SOFTMAX warps — REAL SOFTMAX (P only, no O normalize in TMEM) ===== - if warp_idx < self.mma_warp_id: - tmem.allocate(self.num_tmem_alloc_cols) - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) - sfw_idx = tidx % (32 * len(self.epilogue_warp_id)) - - # S load - tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS0) - thr_load = tiled_tmem_load.get_slice(sfw_idx) - tTMEM_LOADtS = thr_load.partition_S(tStS0) - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tTMEM_LOADcS = thr_load.partition_D(tScS) - - # P store - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - tStP_layout = cute.composition(tStS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tStP0 = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStP_layout) - tmem_store_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStP0) - thr_store = tiled_tmem_store.get_slice(sfw_idx) - tTMEM_STOREtP = thr_store.partition_D(tStP0) - tScP_layout = cute.composition(tScS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tScP = cute.make_tensor(tScS.iterator, tScP_layout) - tTMEM_STOREcP = thr_store.partition_S(tScP) - - row_max = -Float32.inf - row_sum = Float32(0.0) - scale_log2 = Float32(self.scale_softmax_log2) - - n_tiles = self.s_k // 128 - for kt in range(n_tiles): - si_handle = s_cons.wait_and_advance() - # Load S[kt] - tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS) - cute.arch.fence_view_async_tmem_load() - - # Pass 1: update row_max - old_row_max = row_max - frg_cnt = 4 - frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt - tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile)) - for j in range(frg_cnt): - for k in range(cute.size(tTMEM_LOADrS_frg, mode=[0])): - row_max = cute.arch.fmax(row_max, tTMEM_LOADrS_frg[k, j] * scale_log2) - - row_max_safe = row_max - if row_max == -cutlass.Float32.inf: - row_max_safe = Float32(0.0) - - # acc_scale: exp2(old_max - new_max) for O rescale - acc_scale_ = old_row_max - row_max_safe - acc_scale = cute.math.exp2(acc_scale_, fastmath=True) - if old_row_max == -cutlass.Float32.inf: - acc_scale = Float32(0.0) - row_sum *= acc_scale - - # Pass 2: P = exp2(S * scale_log2 - row_max), accumulate row_sum - rP_words = cute.make_rmem_tensor(tTMEM_STOREcP.shape, self.qk_acc_dtype) - rP_bf16 = cute.make_tensor(cute.recast_ptr(rP_words.iterator, dtype=self.q_dtype), tTMEM_LOADrS.layout) - minus_row_max = Float32(0.0) - row_max_safe - rP_bf16_frg = cute.logical_divide(rP_bf16, cute.make_layout(frg_tile)) - for j in range(frg_cnt): - for k in range(cute.size(tTMEM_LOADrS_frg, mode=[0])): - tTMEM_LOADrS_frg[k, j] = tTMEM_LOADrS_frg[k, j] * scale_log2 + minus_row_max - tTMEM_LOADrS_frg[k, j] = cute.math.exp2(tTMEM_LOADrS_frg[k, j], fastmath=True) - row_sum = row_sum + tTMEM_LOADrS_frg[k, j] - s_vec = tTMEM_LOADrS_frg[None, j].load() - rP_bf16_frg[None, j].store(s_vec.to(self.q_dtype)) - - cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP) - cute.arch.fence_view_async_tmem_store() - si_handle.release() - softmax_done_bar.arrive() - - # Epilogue: standard epilogue_tma_store - # TODO: replace with custom epilogue that applies 1/row_sum - tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout) - acc_cons_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage) - c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) - c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp) - acc_cons_st = utils.gemm.sm100.epilogue_tma_store(self, tidx, warp_idx, tma_c, tCtO_base, sC, tCgC, epi_tile, 0, const_expr(lambda x: x), (0,0,0), acc_cons_st, acc_pipe, c_pipe) - c_pipe.producer_tail() - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def test(): - for n in [128, 256, 384, 512, 1024]: - torch.manual_seed(42) - m, hd = 128, HEAD_DIM - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(n, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(n, hd, dtype=torch.bfloat16, device='cuda') - v_kernel = v.unsqueeze(-1) - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - - # Reference: proper softmax - qf = q[:, :, 0].float() - kf = k[:, :, 0].float() - scale = 1.0 / math.sqrt(hd) - attn = qf @ kf.T * scale - attn = torch.softmax(attn, dim=-1) - ref = attn @ v.float() - - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - kernel = FmhaV3RealSoftmax(s_k=n) - print(f'n={n}: Compiling... [REAL_SOFTMAX_v2_EPILOGUE]', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - - out = c[:, :, 0].float() - cos = torch.nn.functional.cosine_similarity( - out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0) - ).item() - max_abs = (out - ref).abs().max().item() - n_tiles = n // 128 - print(f'FMHA n={n} ({n_tiles} tiles): ' - f'cos {cos:.6f} max_abs {max_abs:.4f} ' - f'{"PASS" if cos >= 0.99 else "FAIL"}') - if cos < 0.99: - print(f' out[0,:4]={out[0,:4].tolist()}') - print(f' ref[0,:4]={ref[0,:4].tolist()}') - - -if __name__ == '__main__': - test() diff --git a/tests/archive/fmha_v3_stage_c_example1.py b/tests/archive/fmha_v3_stage_c_example1.py deleted file mode 100644 index 185238aa..00000000 --- a/tests/archive/fmha_v3_stage_c_example1.py +++ /dev/null @@ -1,479 +0,0 @@ -""" -FMHA v3 Stage-C Multi-Tile: real online softmax + O normalization across N KV tiles. - -Builds on the working single-tile Stage-C (cos 0.993 at n=128) and fixes the three -multi-tile bugs. - -Fix 1 — caller passes s_k = actual sequence length n. - The single-tile test hardcoded s_k=128 for both n=128 and n=256. v_fmha is - reconstructed as (HEAD_DIM, s_k, 1) MN-major, so when s_k < n the layout - only describes the first s_k V tokens and TMA reads OOB on later tiles. - The kernel takes s_k from __init__; tests must pass s_k=n. - -Fix 2 — TMA producer indexes K and V by kt, not by pipeline count. - The kv pipeline interleaves K and V acquires, so the pipeline `count` - returned by acquire_and_advance goes 0,1,2,3,... The original code did: - - cute.copy(tma_k, tBgK[(None, kh.count)], tBsK[(None, kh.index)], ...) - cute.copy(tma_v, tVgV[(None, vh.count)], tVsV[(None, vh.index)], ...) - - so for kt=0 it loaded K[0] and V[1], for kt=1 it loaded K[2] and V[3]. - Single tile got away with V[1] because the V iteration dim has size 1 - (stride 0 in CuTe → indexing any value yields V[0]); multi-tile loses - that degeneracy and the bug becomes real. - - The SMEM ring buffer index (kh.index, vh.index) was already correct — - consumers read with the same indices — so only the GMEM indexer changes. - -Fix 3 — re-enable O rescale between KV tiles. - Online softmax requires O *= exp2(old_max − new_max) when the row max - grows. The single-tile file disables it during debug; that path is - untaken for n_kv_tiles=1 so single tile still passes. For multi-tile we - have to apply the rescale at the start of every kt > 0, **before** the - softmax_done_bar arrival (so MMA's next PV reads the rescaled O). - -Synchronization (why the per-tile O rescale is safe): - softmax kt's body reads O written by MMA's PV[kt-1]. That PV happened - earlier in MMA's iteration kt-1, then MMA did fence_view_async_tmem_store, - then issued QK[kt], then sh.commit on the S pipe, then arrived on - softmax_done_bar. Our s_cons.wait_and_advance at the top of softmax kt - waits on that S commit, which acts as a memory ordering point — all - MMA writes prior to the commit (including PV[kt-1]) are visible to us. - The fence we add after the tmem_load_o gives thread-local ordering for - the in-register data. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct -import math - -HEAD_DIM = 64 - - -class FmhaV3StageCMulti: - def __init__(self, s_k=128, scale_softmax=None): - # s_k MUST equal the actual sequence length n. See Fix 1. - self.s_k = s_k - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0,1,2,3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.num_c_stage = 2 - self.kv_stage = 2; self.q_stage = 1; self.num_c_stage = 2 - self.scale_softmax = scale_softmax if scale_softmax is not None else 1.0 / math.sqrt(HEAD_DIM) - self.scale_softmax_log2 = self.scale_softmax * math.log2(math.e) - - def _setup(self, qk_mma, pv_mma): - qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (128, 128, qk_ik * 4) - pv_ik = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (128, HEAD_DIM, pv_ik * (128 // pv_ik)) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0]//cute.size(qk_mma.thr_id.shape), HEAD_DIM, self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - self.q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, self.q_dtype, self.q_stage) - self.k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, self.q_dtype, self.kv_stage) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, self.kv_stage) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - qk_thr = qk_mma.get_slice(0); qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - pv_thr = pv_mma.get_slice(0); pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - # TMEM layout: S at 0..128, P at 32..96 (overlaps S, intentional FMHA pattern), - # O after max(S_end, P_end) aligned to 32. - self.tmem_s0_offset = 0; self.tmem_p0_offset = 32 - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - p_end = self.tmem_p0_offset + p_cols_fp32 - s_cols = self.qk_mma_tiler[1] - o_after = max(s_cols, p_end) - self.tmem_o0_offset = ((o_after + 31) // 32) * 32 - o_cols = find_tmem_tensor_col_offset(tOtO) - total = self.tmem_o0_offset + o_cols - self.num_tmem_alloc_cols = 1 - while self.num_tmem_alloc_cols < total: - self.num_tmem_alloc_cols *= 2 - cta = cute.size(qk_mma.thr_id.shape) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)) - self.q_tx_bytes = cute.size_in_bytes(self.q_dtype, q_s) * cta - self.kv_tx_bytes = cute.size_in_bytes(self.q_dtype, k_s) * cta - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - # Fix 1: s_k must equal actual n. With s_k < n, this layout only spans - # the first s_k V tokens and TMA reads past the end on later kv tiles. - v_fmha = cute.make_tensor( - v.iterator, - cute.make_layout( - (HEAD_DIM, self.s_k, 1), - stride=(1, HEAD_DIM, HEAD_DIM * self.s_k), - ), - ) - self.v_major = LayoutEnum.from_tensor(v_fmha).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, self.a_major, self.b_major, self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, self.qk_acc_dtype, self.cta_group, (128,HEAD_DIM), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)); v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id),q,q_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id),k,k_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v_fmha,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape) - epi_s = cute.select(self.c_smem_s,mode=[0,1]) - tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile) - self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.v_smem_s,self.p_tmem_s,self.c_smem_s,self.epi_tile).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, q_smem_s, k_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx,_,_ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k); cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - q_bar: cute.struct.MemRange[cutlass.Int64, self.q_stage*2] - kv_bar: cute.struct.MemRange[cutlass.Int64, self.kv_stage*2] - s_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage*2] - tmem_dealloc: cutlass.Int64; holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - qp,qc = pipeline.PipelineTmaUmma.create(barrier_storage=st.q_bar.data_ptr(),num_stages=self.q_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.q_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - kvp,kvc = pipeline.PipelineTmaUmma.create(barrier_storage=st.kv_bar.data_ptr(),num_stages=self.kv_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.kv_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - s_prod,s_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.s_bar.data_ptr(),num_stages=1,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.epilogue_warp_id))).make_participants() - softmax_done_bar = pipeline.NamedBarrier(barrier_id=3, num_threads=32 + 32*len(self.epilogue_warp_id)) - acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(),num_stages=self.num_acc_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,len(self.epilogue_warp_id)),cta_layout_vmnk=cl_vmnk,defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2,num_threads=32*len((self.mma_warp_id,*self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr,barrier_for_retrieve=tmem_bar,allocator_warp_id=self.epilogue_warp_id[0],is_two_cta=cute.size(qk_mma.thr_id.shape)==2,two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk,is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype,layout=q_smem_s.outer,byte_alignment=128,swizzle=q_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype,layout=k_smem_s.outer,byte_alignment=128,swizzle=k_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype,layout=v_smem_s.outer,byte_alignment=128,swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype,layout=c_smem_s.outer,byte_alignment=128,swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None)) - gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None)) - gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None)) - gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None)) - n_kv_tiles = cute.size(gK, mode=[3]) - - qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - tCgV = pv_thr.partition_B(gV); tCgC = pv_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,0,None,0)).shape) - tAsQ,tAgQ = cpasync.tma_partition(tma_q,0,a_lay,cute.group_modes(sQ,0,3),cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape) - tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3)) - tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)]; tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - # PV reads P from TMEM (where softmax stores it). Same TMEM region as S but - # at offset tmem_p0_offset; values are BF16 packed into FP32 words. - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None,None,None,0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ===== TMA LOAD warp ===== - # Fix 2: index K and V GMEM tensors by kt (the loop variable), NOT by the - # pipeline's interleaved count. SMEM ring buffer index (kh.index/vh.index) - # is still correct — the consumer uses the same indices. - if warp_idx == self.tma_warp_id: - qp.reset(); qh = qp.acquire_and_advance() - cute.copy(tma_q, tAgQ[(None, qh.count)], tAsQ[(None, qh.index)], tma_bar_ptr=qh.barrier) - qp.tail() - kvp.reset(); pk = kvp.try_acquire() - for kt in cutlass.range(n_kv_tiles, unroll=1): - kh = kvp.acquire_and_advance(pk) - # GMEM tile: kt (correct K[kt]). SMEM slot: kh.index (ring buffer). - # TODO: tBgK[(None, kt)] indexes SMEM dim, NOT GMEM tile (the slice - # (None,0,None,0) fixed gmem_iter to 0). CuTeDSL doesn't support - # dynamic gmem tile indexing via subscript. Need pipeline count mapping - # or separate K/V TMA loops. For now, kh.count works for single tile. - cute.copy(tma_k, tBgK[(None, kh.count)], tBsK[(None, kh.index)], tma_bar_ptr=kh.barrier) - pk = cutlass.Boolean(1) - vh = kvp.acquire_and_advance(pk) - # GMEM tile: kt (correct V[kt]). SMEM slot: vh.index (ring buffer). - cute.copy(tma_v, tVgV[(None, vh.count)], tVsV[(None, vh.index)], tma_bar_ptr=vh.barrier) - pk = cutlass.Boolean(1) - kvp.tail() - - # ===== MMA warp ===== - # Unchanged from single-tile. ACCUMULATE flag for PV is True for kt > 0 - # so O = sum_kt(P_kt @ V_kt). The softmax warps rescale O between tiles - # before signaling softmax_done_bar, so by the time PV[kt] runs O has - # already been multiplied by exp2(old_max - new_max). - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - qc.reset(); qh = qc.wait_and_advance(); qh.release() - kvc.reset(); pk = kvc.try_wait() - acc_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_stage) - acc_pipe.producer_acquire(acc_st) - for kt in range(n_kv_tiles): - kh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - sh = s_prod.acquire_and_advance() - qk_mma.set(tcgen05.Field.ACCUMULATE, False) - for kb in cutlass.range(cute.size(tCrQ, mode=[2]), unroll_full=True): - cute.gemm(qk_mma, tStS0, tCrQ[(None,None,kb,0)], tCrK[(None,None,kb,kh.index)], tStS0) - qk_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - sh.commit(); kh.release() - softmax_done_bar.arrive_and_wait() - vh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - pv_mma.set(tcgen05.Field.ACCUMULATE, kt != 0) - for kb in cutlass.range(cute.size(tOrP0, mode=[2]), unroll_full=True): - cute.gemm(pv_mma, tOtO0, tOrP0[(None,None,kb)], tCrV[(None,None,kb,vh.index)], tOtO0) - pv_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - vh.release() - acc_pipe.producer_commit(acc_st); acc_st.advance() - acc_pipe.producer_tail(acc_st) - - # ===== SOFTMAX + EPILOGUE warps (0..3) ===== - if warp_idx < self.mma_warp_id: - tmem.allocate(self.num_tmem_alloc_cols) - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) - sfw_idx = tidx % (32 * len(self.epilogue_warp_id)) - - # --- S load: QK C-fragment layout, FP32 --- - tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS0) - thr_load = tiled_tmem_load.get_slice(sfw_idx) - tTMEM_LOADtS = thr_load.partition_S(tStS0) - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tTMEM_LOADcS = thr_load.partition_D(tScS) - - # --- P store: QK C-fragment composition with P sub-tile, FP32 backing - # of BF16-packed pairs (the FMHA register-bridge pattern). --- - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - tStP_layout = cute.composition(tStS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tStP0 = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStP_layout) - tmem_store_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStP0) - thr_store = tiled_tmem_store.get_slice(sfw_idx) - tTMEM_STOREtP = thr_store.partition_D(tStP0) - tScP_layout = cute.composition(tScS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tScP = cute.make_tensor(tScS.iterator, tScP_layout) - tTMEM_STOREcP = thr_store.partition_S(tScP) - - # --- O load/store path for rescale + final normalize --- - # Tile O column-wise so the load/store fits in registers per pass. - cO = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1])) - tOcO = pv_thr.partition_C(cO) - corr_tile_size = 16 - tOtO_i_layout = cute.composition(tOtO.layout, cute.make_layout((128, corr_tile_size))) - tOcO_i_layout = cute.composition(tOcO.layout, cute.make_layout((128, corr_tile_size))) - tOtO_i = cute.make_tensor(tOtO.iterator, tOtO_i_layout) - tOcO_i = cute.make_tensor(tOcO.iterator, tOcO_i_layout) - tmem_load_o_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), self.acc_dtype) - tmem_store_o_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), self.acc_dtype) - tiled_tmem_load_o = tcgen05.make_tmem_copy(tmem_load_o_atom, tOtO_i) - tiled_tmem_store_o = tcgen05.make_tmem_copy(tmem_store_o_atom, tOtO_i) - thr_load_o = tiled_tmem_load_o.get_slice(sfw_idx) - thr_store_o = tiled_tmem_store_o.get_slice(sfw_idx) - tTMEM_LOAD_OtO = thr_load_o.partition_S(tOtO_i) - tTMEM_LOAD_OcO = thr_load_o.partition_D(tOcO_i) - tTMEM_STORE_OtO = thr_store_o.partition_D(tOtO_i) - - o_col_tiles = self.pv_mma_tiler[1] // corr_tile_size - - row_max = -Float32.inf - row_sum = Float32(0.0) - scale_log2 = Float32(self.scale_softmax_log2) - - for kt in range(n_kv_tiles): - si_handle = s_cons.wait_and_advance() - - # --- Load S[kt] from TMEM into registers --- - tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS) - cute.arch.fence_view_async_tmem_load() - - # --- Pass 1: update row_max with S[kt] (in log2-domain) --- - old_row_max = row_max - frg_cnt = 4 - frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt - tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile)) - for j in range(frg_cnt): - for k in range(cute.size(tTMEM_LOADrS_frg, mode=[0])): - row_max = cute.arch.fmax(row_max, tTMEM_LOADrS_frg[k, j] * scale_log2) - - row_max_safe = row_max - if row_max == -cutlass.Float32.inf: - row_max_safe = Float32(0.0) - - # acc_scale = exp2(old_max - new_max). On the first tile old_max - # is -inf; force acc_scale to 0 so initial row_sum stays 0 and - # we skip the O rescale path entirely. - acc_scale_ = scale_log2 * (old_row_max - row_max_safe) - acc_scale = cute.math.exp2(acc_scale_, fastmath=True) - if old_row_max == -cutlass.Float32.inf: - acc_scale = Float32(0.0) - row_sum *= acc_scale - - # --- Pass 2: P[kt] = exp2((S - new_max) * log2), accumulate row_sum, - # cast to BF16 through the FP32-backed register bridge. --- - rP_words = cute.make_rmem_tensor(tTMEM_STOREcP.shape, self.qk_acc_dtype) - rP_bf16 = cute.make_tensor(cute.recast_ptr(rP_words.iterator, dtype=self.q_dtype), tTMEM_LOADrS.layout) - minus_row_max_scale = (Float32(0.0) - row_max_safe) * scale_log2 - - rP_bf16_frg = cute.logical_divide(rP_bf16, cute.make_layout(frg_tile)) - for j in range(frg_cnt): - for k in range(cute.size(tTMEM_LOADrS_frg, mode=[0])): - tTMEM_LOADrS_frg[k, j] = tTMEM_LOADrS_frg[k, j] * scale_log2 + minus_row_max_scale - tTMEM_LOADrS_frg[k, j] = cute.math.exp2(tTMEM_LOADrS_frg[k, j], fastmath=True) - row_sum = row_sum + tTMEM_LOADrS_frg[k, j] - s_vec = tTMEM_LOADrS_frg[None, j].load() - rP_bf16_frg[None, j].store(s_vec.to(self.q_dtype)) - - cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP) - cute.arch.fence_view_async_tmem_store() - - # --- Fix 3: O rescale for kt > 0 --- - # MMA's previous PV wrote O at offset tmem_o0_offset. We multiply - # O in place by acc_scale, then arrive on softmax_done_bar so MMA - # can start PV[kt] which accumulates the new P @ V on top of the - # rescaled O. - # IMPORTANT: this must happen before softmax_done_bar.arrive(). - if kt > 0: - for i in range(o_col_tiles): - tTMEM_LOAD_O_i = cute.make_tensor( - tTMEM_LOAD_OtO.iterator + i * corr_tile_size, - tTMEM_LOAD_OtO.layout, - ) - tTMEM_STORE_O_i = cute.make_tensor( - tTMEM_STORE_OtO.iterator + i * corr_tile_size, - tTMEM_STORE_OtO.layout, - ) - tTMrO = cute.make_rmem_tensor(tTMEM_LOAD_OcO.shape, self.acc_dtype) - cute.copy(tiled_tmem_load_o, tTMEM_LOAD_O_i, tTMrO) - cute.arch.fence_view_async_tmem_load() - for k in cutlass.range(cute.size(tTMrO), vectorize=True): - tTMrO[k] = tTMrO[k] * acc_scale - cute.copy(tiled_tmem_store_o, tTMrO, tTMEM_STORE_O_i) - cute.arch.fence_view_async_tmem_store() - - si_handle.release() - softmax_done_bar.arrive() - - # --- Wait for MMA to finish PV[N-1] before reading O --- - # Without this, the final normalize reads O while MMA may still be - # writing PV[N-1]. Single-tile wins the race by luck; multi-tile drifts. - acc_cons_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage) - acc_pipe.consumer_wait(acc_cons_st) - - # --- Final O normalization by row_sum --- - # After the last softmax iteration, MMA still needs to finish the final - # PV[N-1]. The acc_pipe consumer wait above ensures O is fully accumulated - # before we divide by row_sum here. - inv_row_sum = Float32(1.0) / row_sum - for i in range(o_col_tiles): - tTMEM_LOAD_O_i = cute.make_tensor( - tTMEM_LOAD_OtO.iterator + i * corr_tile_size, - tTMEM_LOAD_OtO.layout, - ) - tTMEM_STORE_O_i = cute.make_tensor( - tTMEM_STORE_OtO.iterator + i * corr_tile_size, - tTMEM_STORE_OtO.layout, - ) - tTMrO = cute.make_rmem_tensor(tTMEM_LOAD_OcO.shape, self.acc_dtype) - cute.copy(tiled_tmem_load_o, tTMEM_LOAD_O_i, tTMrO) - cute.arch.fence_view_async_tmem_load() - for k in cutlass.range(cute.size(tTMrO), vectorize=True): - tTMrO[k] = tTMrO[k] * inv_row_sum - cute.copy(tiled_tmem_store_o, tTMrO, tTMEM_STORE_O_i) - cute.arch.fence_view_async_tmem_store() - - # --- Epilogue: TMEM -> SMEM -> GMEM via TMA store --- - tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout) - c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) - c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp) - acc_cons_st = utils.gemm.sm100.epilogue_tma_store(self, tidx, warp_idx, tma_c, tCtO_base, sC, tCgC, epi_tile, 0, const_expr(lambda x: x), (0,0,0), acc_cons_st, acc_pipe, c_pipe) - c_pipe.producer_tail() - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def test(): - torch.manual_seed(42) - # Exercise single tile, 2 tiles, 4 tiles, 8 tiles. - for n in [128, 256, 512, 1024]: - torch.manual_seed(42) - m, hd = 128, HEAD_DIM - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(n, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(n, hd, dtype=torch.bfloat16, device='cuda') - v_kernel = v.unsqueeze(-1) - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - - # Reference: full softmax over the actual sequence. - qf = q[:, :, 0].float() - kf = k[:, :, 0].float() - scale = 1.0 / math.sqrt(hd) - attn = qf @ kf.T * scale - attn = torch.softmax(attn, dim=-1) - ref = attn @ v.float() - - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - # Fix 1: s_k MUST equal n. Each n requires its own compiled kernel - # because s_k controls the v_fmha layout (a compile-time constant). - kernel = FmhaV3StageCMulti(s_k=n) - print(f'n={n}: Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - n_tiles = n // 128 - print(f'n={n}: n_kv_tiles should be {n_tiles}, tmem_offsets: s0={kernel.tmem_s0_offset} ' - f'p0={kernel.tmem_p0_offset} o0={kernel.tmem_o0_offset} ' - f'alloc={kernel.num_tmem_alloc_cols}', flush=True) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - - out = c[:, :, 0].float() - cos = torch.nn.functional.cosine_similarity( - out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0) - ).item() - max_abs = (out - ref).abs().max().item() - n_tiles = n // 128 - print(f'FMHA Stage-C Multi-Tile n={n} ({n_tiles} kv tiles): ' - f'cosine {cos:.6f} max_abs {max_abs:.4f} ' - f'{"PASS" if cos >= 0.99 else "FAIL"}') - if cos < 0.99: - print(f' out[0,:4]={out[0,:4].tolist()} ref[0,:4]={ref[0,:4].tolist()}') - - -if __name__ == '__main__': - test() \ No newline at end of file diff --git a/tests/archive/fmha_v3_stage_c_example2.py b/tests/archive/fmha_v3_stage_c_example2.py deleted file mode 100644 index 33a5edaa..00000000 --- a/tests/archive/fmha_v3_stage_c_example2.py +++ /dev/null @@ -1,451 +0,0 @@ -""" -FMHA v3 Stage-C Multi-Tile (combined K+V barrier). - -Replaces the interleaved K-then-V acquires with a single acquire per kt that -loads K and V onto the SAME barrier slot. tx_count is sized for K+V together. -With one acquire per tile, the pipeline `count` returned by acquire_and_advance -goes 0, 1, 2, ... and matches the KV tile index directly — no interleaving -problem, and no need for Python ints or integer-division gymnastics in the -TMA coordinate. kvh.count stays a first-class pipeline state value, which is -the form cute.copy accepts. - -Changes vs the single-tile file: - -1. s_k MUST equal actual n. v_fmha layout uses s_k as the V sequence dim. - -2. kv pipeline carries combined K+V per stage: - - tx_count = K_bytes + V_bytes - - producer: one acquire per kt, K and V copies share kvh.barrier - - consumer: one wait per kt, kvh.index used for both sK and sV reads - - release happens after PV (no separate K-early-release path) - Bonus: this also fixes the unused-SMEM-slot quirk where the original kernel - only ever used sK[0] and sV[1] because of the interleaved count. - -3. O rescale between KV tiles re-enabled (gated on kt > 0). Lives in softmax - body BEFORE softmax_done_bar.arrive(), so MMA's PV[kt] reads a rescaled O. - -4. Explicit MMA→softmax sync before the final normalize. - final_o_bar is a NamedBarrier with 32 MMA + 128 softmax threads. MMA - .arrive() after acc_pipe.producer_tail; softmax .arrive_and_wait() before - reading O. Without this, softmax can race MMA's PV[N-1] and divide a - partially-accumulated O by row_sum. The single-tile test masked the race - because the timing happened to work. - -Alternative if combined-barrier ever bites: keep the interleaved pipeline and -index GMEM by `kh.count // 2` / `vh.count // 2`. Requires CuTeDSL to support -Int32 floor-division in a TMA coordinate. Not used here. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct -import math - -HEAD_DIM = 64 - - -class FmhaV3StageCMulti: - def __init__(self, s_k=128, scale_softmax=None): - # s_k MUST equal actual sequence length n. - self.s_k = s_k - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0,1,2,3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.num_c_stage = 2 - self.kv_stage = 2; self.q_stage = 1; self.num_c_stage = 2 - self.scale_softmax = scale_softmax if scale_softmax is not None else 1.0 / math.sqrt(HEAD_DIM) - self.scale_softmax_log2 = self.scale_softmax * math.log2(math.e) - - def _setup(self, qk_mma, pv_mma): - qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (128, 128, qk_ik * 4) - pv_ik = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (128, HEAD_DIM, pv_ik * (128 // pv_ik)) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0]//cute.size(qk_mma.thr_id.shape), HEAD_DIM, self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - self.q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, self.q_dtype, self.q_stage) - self.k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, self.q_dtype, self.kv_stage) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, self.kv_stage) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - qk_thr = qk_mma.get_slice(0); qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - pv_thr = pv_mma.get_slice(0); pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - self.tmem_s0_offset = 0; self.tmem_p0_offset = 32 - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - p_end = self.tmem_p0_offset + p_cols_fp32 - s_cols = self.qk_mma_tiler[1] - o_after = max(s_cols, p_end) - self.tmem_o0_offset = ((o_after + 31) // 32) * 32 - o_cols = find_tmem_tensor_col_offset(tOtO) - total = self.tmem_o0_offset + o_cols - self.num_tmem_alloc_cols = 1 - while self.num_tmem_alloc_cols < total: - self.num_tmem_alloc_cols *= 2 - cta = cute.size(qk_mma.thr_id.shape) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)) - k_s = cute.slice_(self.k_smem_s,(None,None,None,0)) - v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - self.q_tx_bytes = cute.size_in_bytes(self.q_dtype, q_s) * cta - # Combined barrier: tx_count covers BOTH K and V transfers per acquire. - self.kv_tx_bytes = (cute.size_in_bytes(self.q_dtype, k_s) + - cute.size_in_bytes(self.q_dtype, v_s)) * cta - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - v_fmha = cute.make_tensor( - v.iterator, - cute.make_layout( - (HEAD_DIM, self.s_k, 1), - stride=(1, HEAD_DIM, HEAD_DIM * self.s_k), - ), - ) - self.v_major = LayoutEnum.from_tensor(v_fmha).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, self.a_major, self.b_major, self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, self.qk_acc_dtype, self.cta_group, (128,HEAD_DIM), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)); v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id),q,q_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id),k,k_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v_fmha,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape) - epi_s = cute.select(self.c_smem_s,mode=[0,1]) - tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile) - self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.v_smem_s,self.p_tmem_s,self.c_smem_s,self.epi_tile).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, q_smem_s, k_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx,_,_ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k); cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - q_bar: cute.struct.MemRange[cutlass.Int64, self.q_stage*2] - kv_bar: cute.struct.MemRange[cutlass.Int64, self.kv_stage*2] - s_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage*2] - tmem_dealloc: cutlass.Int64; holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - qp,qc = pipeline.PipelineTmaUmma.create(barrier_storage=st.q_bar.data_ptr(),num_stages=self.q_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.q_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - # Combined K+V pipeline: each stage carries BOTH K and V loaded together. - kvp,kvc = pipeline.PipelineTmaUmma.create(barrier_storage=st.kv_bar.data_ptr(),num_stages=self.kv_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.kv_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - s_prod,s_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.s_bar.data_ptr(),num_stages=1,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.epilogue_warp_id))).make_participants() - softmax_done_bar = pipeline.NamedBarrier(barrier_id=3, num_threads=32 + 32*len(self.epilogue_warp_id)) - # Final-O sync: MMA arrives once after acc_pipe.producer_tail; softmax - # arrives_and_waits before reading O for the final normalize. - final_o_bar = pipeline.NamedBarrier(barrier_id=4, num_threads=32 + 32*len(self.epilogue_warp_id)) - acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(),num_stages=self.num_acc_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,len(self.epilogue_warp_id)),cta_layout_vmnk=cl_vmnk,defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2,num_threads=32*len((self.mma_warp_id,*self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr,barrier_for_retrieve=tmem_bar,allocator_warp_id=self.epilogue_warp_id[0],is_two_cta=cute.size(qk_mma.thr_id.shape)==2,two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk,is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype,layout=q_smem_s.outer,byte_alignment=128,swizzle=q_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype,layout=k_smem_s.outer,byte_alignment=128,swizzle=k_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype,layout=v_smem_s.outer,byte_alignment=128,swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype,layout=c_smem_s.outer,byte_alignment=128,swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None)) - gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None)) - gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None)) - gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None)) - n_kv_tiles = cute.size(gK, mode=[3]) - - qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - tCgV = pv_thr.partition_B(gV); tCgC = pv_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,0,None,0)).shape) - tAsQ,tAgQ = cpasync.tma_partition(tma_q,0,a_lay,cute.group_modes(sQ,0,3),cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape) - tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3)) - tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)]; tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None,None,None,0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ===== TMA LOAD warp ===== - # One acquire per kt; K and V both target kvh.barrier. kvh.count == kt. - if warp_idx == self.tma_warp_id: - qp.reset(); qh = qp.acquire_and_advance() - cute.copy(tma_q, tAgQ[(None, qh.count)], tAsQ[(None, qh.index)], tma_bar_ptr=qh.barrier) - qp.tail() - kvp.reset(); pk = kvp.try_acquire() - for kt in cutlass.range(n_kv_tiles, unroll=1): - kvh = kvp.acquire_and_advance(pk) - # Both transfers decrement the same barrier's tx_count. - # kvh.count is a pipeline-state Int32 (the form cute.copy accepts). - cute.copy(tma_k, tBgK[(None, kvh.count)], tBsK[(None, kvh.index)], tma_bar_ptr=kvh.barrier) - cute.copy(tma_v, tVgV[(None, kvh.count)], tVsV[(None, kvh.index)], tma_bar_ptr=kvh.barrier) - pk = cutlass.Boolean(1) - kvp.tail() - - # ===== MMA warp ===== - # One wait per kt; same slot index used for both K (QK) and V (PV). - # Release happens AFTER PV — combined slot stays held across QK+PV. - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - qc.reset(); qh = qc.wait_and_advance(); qh.release() - kvc.reset(); pk = kvc.try_wait() - acc_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_stage) - acc_pipe.producer_acquire(acc_st) - for kt in range(n_kv_tiles): - kvh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - sh = s_prod.acquire_and_advance() - qk_mma.set(tcgen05.Field.ACCUMULATE, False) - for kb in cutlass.range(cute.size(tCrQ, mode=[2]), unroll_full=True): - cute.gemm(qk_mma, tStS0, tCrQ[(None,None,kb,0)], tCrK[(None,None,kb,kvh.index)], tStS0) - qk_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - sh.commit() - softmax_done_bar.arrive_and_wait() - pv_mma.set(tcgen05.Field.ACCUMULATE, kt != 0) - for kb in cutlass.range(cute.size(tOrP0, mode=[2]), unroll_full=True): - cute.gemm(pv_mma, tOtO0, tOrP0[(None,None,kb)], tCrV[(None,None,kb,kvh.index)], tOtO0) - pv_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - kvh.release() - acc_pipe.producer_commit(acc_st); acc_st.advance() - acc_pipe.producer_tail(acc_st) - # Signal softmax that all PVs are committed and O is final in TMEM. - final_o_bar.arrive() - - # ===== SOFTMAX + EPILOGUE warps ===== - if warp_idx < self.mma_warp_id: - tmem.allocate(self.num_tmem_alloc_cols) - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) - sfw_idx = tidx % (32 * len(self.epilogue_warp_id)) - - # S load - tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS0) - thr_load = tiled_tmem_load.get_slice(sfw_idx) - tTMEM_LOADtS = thr_load.partition_S(tStS0) - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tTMEM_LOADcS = thr_load.partition_D(tScS) - - # P store - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - tStP_layout = cute.composition(tStS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tStP0 = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStP_layout) - tmem_store_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStP0) - thr_store = tiled_tmem_store.get_slice(sfw_idx) - tTMEM_STOREtP = thr_store.partition_D(tStP0) - tScP_layout = cute.composition(tScS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tScP = cute.make_tensor(tScS.iterator, tScP_layout) - tTMEM_STOREcP = thr_store.partition_S(tScP) - - # O rescale / normalize path - cO = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1])) - tOcO = pv_thr.partition_C(cO) - corr_tile_size = 16 - tOtO_i_layout = cute.composition(tOtO.layout, cute.make_layout((128, corr_tile_size))) - tOcO_i_layout = cute.composition(tOcO.layout, cute.make_layout((128, corr_tile_size))) - tOtO_i = cute.make_tensor(tOtO.iterator, tOtO_i_layout) - tOcO_i = cute.make_tensor(tOcO.iterator, tOcO_i_layout) - tmem_load_o_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), self.acc_dtype) - tmem_store_o_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), self.acc_dtype) - tiled_tmem_load_o = tcgen05.make_tmem_copy(tmem_load_o_atom, tOtO_i) - tiled_tmem_store_o = tcgen05.make_tmem_copy(tmem_store_o_atom, tOtO_i) - thr_load_o = tiled_tmem_load_o.get_slice(sfw_idx) - thr_store_o = tiled_tmem_store_o.get_slice(sfw_idx) - tTMEM_LOAD_OtO = thr_load_o.partition_S(tOtO_i) - tTMEM_LOAD_OcO = thr_load_o.partition_D(tOcO_i) - tTMEM_STORE_OtO = thr_store_o.partition_D(tOtO_i) - - o_col_tiles = self.pv_mma_tiler[1] // corr_tile_size - - row_max = -Float32.inf - row_sum = Float32(0.0) - scale_log2 = Float32(self.scale_softmax_log2) - - for kt in range(n_kv_tiles): - si_handle = s_cons.wait_and_advance() - - # Load S[kt] - tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS) - cute.arch.fence_view_async_tmem_load() - - # Pass 1: update row_max - old_row_max = row_max - frg_cnt = 4 - frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt - tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile)) - for j in range(frg_cnt): - for k in range(cute.size(tTMEM_LOADrS_frg, mode=[0])): - row_max = cute.arch.fmax(row_max, tTMEM_LOADrS_frg[k, j] * scale_log2) - - row_max_safe = row_max - if row_max == -cutlass.Float32.inf: - row_max_safe = Float32(0.0) - - # acc_scale used for both row_sum rescale and O rescale. - acc_scale_ = scale_log2 * (old_row_max - row_max_safe) - acc_scale = cute.math.exp2(acc_scale_, fastmath=True) - if old_row_max == -cutlass.Float32.inf: - acc_scale = Float32(0.0) - row_sum *= acc_scale - - # Pass 2: P = exp2((S - new_max) * log2), accumulate row_sum, - # store BF16 P through the FP32-backed register bridge. - rP_words = cute.make_rmem_tensor(tTMEM_STOREcP.shape, self.qk_acc_dtype) - rP_bf16 = cute.make_tensor(cute.recast_ptr(rP_words.iterator, dtype=self.q_dtype), tTMEM_LOADrS.layout) - minus_row_max_scale = (Float32(0.0) - row_max_safe) * scale_log2 - - rP_bf16_frg = cute.logical_divide(rP_bf16, cute.make_layout(frg_tile)) - for j in range(frg_cnt): - for k in range(cute.size(tTMEM_LOADrS_frg, mode=[0])): - tTMEM_LOADrS_frg[k, j] = tTMEM_LOADrS_frg[k, j] * scale_log2 + minus_row_max_scale - tTMEM_LOADrS_frg[k, j] = cute.math.exp2(tTMEM_LOADrS_frg[k, j], fastmath=True) - row_sum = row_sum + tTMEM_LOADrS_frg[k, j] - s_vec = tTMEM_LOADrS_frg[None, j].load() - rP_bf16_frg[None, j].store(s_vec.to(self.q_dtype)) - - cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP) - cute.arch.fence_view_async_tmem_store() - - # O rescale for kt > 0. Reads O written by MMA's PV[kt-1]; - # visibility is provided by s_cons.wait_and_advance above - # (acquires on MMA's S[kt] commit, which orders PV[kt-1] before). - if kt > 0: - for i in range(o_col_tiles): - tTMEM_LOAD_O_i = cute.make_tensor( - tTMEM_LOAD_OtO.iterator + i * corr_tile_size, - tTMEM_LOAD_OtO.layout, - ) - tTMEM_STORE_O_i = cute.make_tensor( - tTMEM_STORE_OtO.iterator + i * corr_tile_size, - tTMEM_STORE_OtO.layout, - ) - tTMrO = cute.make_rmem_tensor(tTMEM_LOAD_OcO.shape, self.acc_dtype) - cute.copy(tiled_tmem_load_o, tTMEM_LOAD_O_i, tTMrO) - cute.arch.fence_view_async_tmem_load() - for k in cutlass.range(cute.size(tTMrO), vectorize=True): - tTMrO[k] = tTMrO[k] * acc_scale - cute.copy(tiled_tmem_store_o, tTMrO, tTMEM_STORE_O_i) - cute.arch.fence_view_async_tmem_store() - - si_handle.release() - softmax_done_bar.arrive() - - # Wait for MMA's last PV to commit before reading O for normalize. - # Without this barrier softmax can race MMA's PV[N-1]. - final_o_bar.arrive_and_wait() - - # Final O = O / row_sum - inv_row_sum = Float32(1.0) / row_sum - for i in range(o_col_tiles): - tTMEM_LOAD_O_i = cute.make_tensor( - tTMEM_LOAD_OtO.iterator + i * corr_tile_size, - tTMEM_LOAD_OtO.layout, - ) - tTMEM_STORE_O_i = cute.make_tensor( - tTMEM_STORE_OtO.iterator + i * corr_tile_size, - tTMEM_STORE_OtO.layout, - ) - tTMrO = cute.make_rmem_tensor(tTMEM_LOAD_OcO.shape, self.acc_dtype) - cute.copy(tiled_tmem_load_o, tTMEM_LOAD_O_i, tTMrO) - cute.arch.fence_view_async_tmem_load() - for k in cutlass.range(cute.size(tTMrO), vectorize=True): - tTMrO[k] = tTMrO[k] * inv_row_sum - cute.copy(tiled_tmem_store_o, tTMrO, tTMEM_STORE_O_i) - cute.arch.fence_view_async_tmem_store() - - # Epilogue: TMEM -> SMEM -> GMEM via TMA store - tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout) - acc_cons_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage) - c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) - c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp) - acc_cons_st = utils.gemm.sm100.epilogue_tma_store(self, tidx, warp_idx, tma_c, tCtO_base, sC, tCgC, epi_tile, 0, const_expr(lambda x: x), (0,0,0), acc_cons_st, acc_pipe, c_pipe) - c_pipe.producer_tail() - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def test(): - torch.manual_seed(42) - for n in [128, 256, 512, 1024]: - torch.manual_seed(42) - m, hd = 128, HEAD_DIM - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(n, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(n, hd, dtype=torch.bfloat16, device='cuda') - v_kernel = v.unsqueeze(-1) - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - - qf = q[:, :, 0].float() - kf = k[:, :, 0].float() - scale = 1.0 / math.sqrt(hd) - attn = qf @ kf.T * scale - attn = torch.softmax(attn, dim=-1) - ref = attn @ v.float() - - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - # Each n requires its own compiled kernel (s_k is compile-time). - kernel = FmhaV3StageCMulti(s_k=n) - print(f'n={n}: Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - print(f'n={n}: tmem s0={kernel.tmem_s0_offset} p0={kernel.tmem_p0_offset} ' - f'o0={kernel.tmem_o0_offset} alloc={kernel.num_tmem_alloc_cols} ' - f'kv_tx_bytes={kernel.kv_tx_bytes}', flush=True) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - - out = c[:, :, 0].float() - cos = torch.nn.functional.cosine_similarity( - out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0) - ).item() - max_abs = (out - ref).abs().max().item() - n_tiles = n // 128 - print(f'FMHA Stage-C Multi n={n} ({n_tiles} kv tiles): ' - f'cos {cos:.6f} max_abs {max_abs:.4f} ' - f'{"PASS" if cos >= 0.99 else "FAIL"}') - if cos < 0.99: - print(f' out[0,:4]={out[0,:4].tolist()}') - print(f' ref[0,:4]={ref[0,:4].tolist()}') - - -if __name__ == '__main__': - test() \ No newline at end of file diff --git a/tests/archive/fmha_v3_stage_c_example3.py b/tests/archive/fmha_v3_stage_c_example3.py deleted file mode 100644 index 881761e3..00000000 --- a/tests/archive/fmha_v3_stage_c_example3.py +++ /dev/null @@ -1,456 +0,0 @@ -""" -FMHA v3 Stage-C Multi-Tile (combined K+V barrier). - -Replaces the interleaved K-then-V acquires with a single acquire per kt that -loads K and V onto the SAME barrier slot. tx_count is sized for K+V together. -With one acquire per tile, the pipeline `count` returned by acquire_and_advance -goes 0, 1, 2, ... and matches the KV tile index directly — no interleaving -problem, and no need for Python ints or integer-division gymnastics in the -TMA coordinate. kvh.count stays a first-class pipeline state value, which is -the form cute.copy accepts. - -Changes vs the single-tile file: - -1. s_k MUST equal actual n. v_fmha layout uses s_k as the V sequence dim. - -2. kv pipeline carries combined K+V per stage: - - tx_count = K_bytes + V_bytes - - producer: one acquire per kt, K and V copies share kvh.barrier - - consumer: one wait per kt, kvh.index used for both sK and sV reads - - release happens after PV (no separate K-early-release path) - Bonus: this also fixes the unused-SMEM-slot quirk where the original kernel - only ever used sK[0] and sV[1] because of the interleaved count. - -3. O rescale between KV tiles re-enabled (gated on kt > 0). Lives in softmax - body BEFORE softmax_done_bar.arrive(), so MMA's PV[kt] reads a rescaled O. - -4. Explicit MMA→softmax sync before the final normalize. - final_o_bar is a NamedBarrier with 32 MMA + 128 softmax threads. MMA - .arrive() between acc_pipe.producer_commit and producer_tail; softmax - .arrive_and_wait() before reading O. Without this, softmax can race MMA's PV[N-1] and divide a - partially-accumulated O by row_sum. The single-tile test masked the race - because the timing happened to work. - -Alternative if combined-barrier ever bites: keep the interleaved pipeline and -index GMEM by `kh.count // 2` / `vh.count // 2`. Requires CuTeDSL to support -Int32 floor-division in a TMA coordinate. Not used here. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct -import math - -HEAD_DIM = 64 - - -class FmhaV3StageCMulti: - def __init__(self, s_k=128, scale_softmax=None): - # s_k MUST equal actual sequence length n. - self.s_k = s_k - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0,1,2,3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.num_c_stage = 2 - self.kv_stage = 2; self.q_stage = 1; self.num_c_stage = 2 - self.scale_softmax = scale_softmax if scale_softmax is not None else 1.0 / math.sqrt(HEAD_DIM) - self.scale_softmax_log2 = self.scale_softmax * math.log2(math.e) - - def _setup(self, qk_mma, pv_mma): - qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (128, 128, qk_ik * 4) - pv_ik = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (128, HEAD_DIM, pv_ik * (128 // pv_ik)) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0]//cute.size(qk_mma.thr_id.shape), HEAD_DIM, self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - self.q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, self.q_dtype, self.q_stage) - self.k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, self.q_dtype, self.kv_stage) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, self.kv_stage) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - qk_thr = qk_mma.get_slice(0); qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - pv_thr = pv_mma.get_slice(0); pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - self.tmem_s0_offset = 0; self.tmem_p0_offset = 32 - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - p_end = self.tmem_p0_offset + p_cols_fp32 - s_cols = self.qk_mma_tiler[1] - o_after = max(s_cols, p_end) - self.tmem_o0_offset = ((o_after + 31) // 32) * 32 - o_cols = find_tmem_tensor_col_offset(tOtO) - total = self.tmem_o0_offset + o_cols - self.num_tmem_alloc_cols = 1 - while self.num_tmem_alloc_cols < total: - self.num_tmem_alloc_cols *= 2 - cta = cute.size(qk_mma.thr_id.shape) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)) - k_s = cute.slice_(self.k_smem_s,(None,None,None,0)) - v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - self.q_tx_bytes = cute.size_in_bytes(self.q_dtype, q_s) * cta - # Combined barrier: tx_count covers BOTH K and V transfers per acquire. - self.kv_tx_bytes = (cute.size_in_bytes(self.q_dtype, k_s) + - cute.size_in_bytes(self.q_dtype, v_s)) * cta - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - v_fmha = cute.make_tensor( - v.iterator, - cute.make_layout( - (HEAD_DIM, self.s_k, 1), - stride=(1, HEAD_DIM, HEAD_DIM * self.s_k), - ), - ) - self.v_major = LayoutEnum.from_tensor(v_fmha).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, self.a_major, self.b_major, self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, self.qk_acc_dtype, self.cta_group, (128,HEAD_DIM), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)); v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id),q,q_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id),k,k_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v_fmha,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape) - epi_s = cute.select(self.c_smem_s,mode=[0,1]) - tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile) - self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.v_smem_s,self.p_tmem_s,self.c_smem_s,self.epi_tile).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, q_smem_s, k_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx,_,_ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k); cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - q_bar: cute.struct.MemRange[cutlass.Int64, self.q_stage*2] - kv_bar: cute.struct.MemRange[cutlass.Int64, self.kv_stage*2] - s_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage*2] - tmem_dealloc: cutlass.Int64; holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - qp,qc = pipeline.PipelineTmaUmma.create(barrier_storage=st.q_bar.data_ptr(),num_stages=self.q_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.q_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - # Combined K+V pipeline: each stage carries BOTH K and V loaded together. - kvp,kvc = pipeline.PipelineTmaUmma.create(barrier_storage=st.kv_bar.data_ptr(),num_stages=self.kv_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.kv_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - s_prod,s_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.s_bar.data_ptr(),num_stages=1,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.epilogue_warp_id))).make_participants() - softmax_done_bar = pipeline.NamedBarrier(barrier_id=3, num_threads=32 + 32*len(self.epilogue_warp_id)) - # Final-O sync: MMA arrives between producer_commit and producer_tail; - # softmax arrives_and_waits before reading O for the final normalize. - final_o_bar = pipeline.NamedBarrier(barrier_id=4, num_threads=32 + 32*len(self.epilogue_warp_id)) - acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(),num_stages=self.num_acc_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,len(self.epilogue_warp_id)),cta_layout_vmnk=cl_vmnk,defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2,num_threads=32*len((self.mma_warp_id,*self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr,barrier_for_retrieve=tmem_bar,allocator_warp_id=self.epilogue_warp_id[0],is_two_cta=cute.size(qk_mma.thr_id.shape)==2,two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk,is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype,layout=q_smem_s.outer,byte_alignment=128,swizzle=q_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype,layout=k_smem_s.outer,byte_alignment=128,swizzle=k_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype,layout=v_smem_s.outer,byte_alignment=128,swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype,layout=c_smem_s.outer,byte_alignment=128,swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None)) - gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None)) - gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None)) - gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None)) - n_kv_tiles = cute.size(gK, mode=[3]) - - qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - tCgV = pv_thr.partition_B(gV); tCgC = pv_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,0,None,0)).shape) - tAsQ,tAgQ = cpasync.tma_partition(tma_q,0,a_lay,cute.group_modes(sQ,0,3),cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape) - tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3)) - tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)]; tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None,None,None,0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ===== TMA LOAD warp ===== - # One acquire per kt; K and V both target kvh.barrier. kvh.count == kt. - if warp_idx == self.tma_warp_id: - qp.reset(); qh = qp.acquire_and_advance() - cute.copy(tma_q, tAgQ[(None, qh.count)], tAsQ[(None, qh.index)], tma_bar_ptr=qh.barrier) - qp.tail() - kvp.reset(); pk = kvp.try_acquire() - for kt in cutlass.range(n_kv_tiles, unroll=1): - kvh = kvp.acquire_and_advance(pk) - # Both transfers decrement the same barrier's tx_count. - # kvh.count is a pipeline-state Int32 (the form cute.copy accepts). - cute.copy(tma_k, tBgK[(None, kvh.count)], tBsK[(None, kvh.index)], tma_bar_ptr=kvh.barrier) - cute.copy(tma_v, tVgV[(None, kvh.count)], tVsV[(None, kvh.index)], tma_bar_ptr=kvh.barrier) - pk = cutlass.Boolean(1) - kvp.tail() - - # ===== MMA warp ===== - # One wait per kt; same slot index used for both K (QK) and V (PV). - # Release happens AFTER PV — combined slot stays held across QK+PV. - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - qc.reset(); qh = qc.wait_and_advance(); qh.release() - kvc.reset(); pk = kvc.try_wait() - acc_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_stage) - acc_pipe.producer_acquire(acc_st) - for kt in range(n_kv_tiles): - kvh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - sh = s_prod.acquire_and_advance() - qk_mma.set(tcgen05.Field.ACCUMULATE, False) - for kb in cutlass.range(cute.size(tCrQ, mode=[2]), unroll_full=True): - cute.gemm(qk_mma, tStS0, tCrQ[(None,None,kb,0)], tCrK[(None,None,kb,kvh.index)], tStS0) - qk_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - sh.commit() - softmax_done_bar.arrive_and_wait() - pv_mma.set(tcgen05.Field.ACCUMULATE, kt != 0) - for kb in cutlass.range(cute.size(tOrP0, mode=[2]), unroll_full=True): - cute.gemm(pv_mma, tOtO0, tOrP0[(None,None,kb)], tCrV[(None,None,kb,kvh.index)], tOtO0) - pv_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - kvh.release() - acc_pipe.producer_commit(acc_st); acc_st.advance() - # Signal softmax FIRST so it can run normalize + epilogue. Then - # wait for the epilogue's consumer-release in producer_tail. - # Reverse order deadlocks: producer_tail blocks waiting for - # consumer release; softmax blocks at final_o_bar waiting for - # MMA arrive; the epilogue (which does the release) is gated - # behind softmax's final_o_bar wait. Cycle. - final_o_bar.arrive() - acc_pipe.producer_tail(acc_st) - - # ===== SOFTMAX + EPILOGUE warps ===== - if warp_idx < self.mma_warp_id: - tmem.allocate(self.num_tmem_alloc_cols) - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) - sfw_idx = tidx % (32 * len(self.epilogue_warp_id)) - - # S load - tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS0) - thr_load = tiled_tmem_load.get_slice(sfw_idx) - tTMEM_LOADtS = thr_load.partition_S(tStS0) - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tTMEM_LOADcS = thr_load.partition_D(tScS) - - # P store - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - tStP_layout = cute.composition(tStS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tStP0 = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStP_layout) - tmem_store_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStP0) - thr_store = tiled_tmem_store.get_slice(sfw_idx) - tTMEM_STOREtP = thr_store.partition_D(tStP0) - tScP_layout = cute.composition(tScS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tScP = cute.make_tensor(tScS.iterator, tScP_layout) - tTMEM_STOREcP = thr_store.partition_S(tScP) - - # O rescale / normalize path - cO = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1])) - tOcO = pv_thr.partition_C(cO) - corr_tile_size = 16 - tOtO_i_layout = cute.composition(tOtO.layout, cute.make_layout((128, corr_tile_size))) - tOcO_i_layout = cute.composition(tOcO.layout, cute.make_layout((128, corr_tile_size))) - tOtO_i = cute.make_tensor(tOtO.iterator, tOtO_i_layout) - tOcO_i = cute.make_tensor(tOcO.iterator, tOcO_i_layout) - tmem_load_o_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), self.acc_dtype) - tmem_store_o_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), self.acc_dtype) - tiled_tmem_load_o = tcgen05.make_tmem_copy(tmem_load_o_atom, tOtO_i) - tiled_tmem_store_o = tcgen05.make_tmem_copy(tmem_store_o_atom, tOtO_i) - thr_load_o = tiled_tmem_load_o.get_slice(sfw_idx) - thr_store_o = tiled_tmem_store_o.get_slice(sfw_idx) - tTMEM_LOAD_OtO = thr_load_o.partition_S(tOtO_i) - tTMEM_LOAD_OcO = thr_load_o.partition_D(tOcO_i) - tTMEM_STORE_OtO = thr_store_o.partition_D(tOtO_i) - - o_col_tiles = self.pv_mma_tiler[1] // corr_tile_size - - row_max = -Float32.inf - row_sum = Float32(0.0) - scale_log2 = Float32(self.scale_softmax_log2) - - for kt in range(n_kv_tiles): - si_handle = s_cons.wait_and_advance() - - # Load S[kt] - tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS) - cute.arch.fence_view_async_tmem_load() - - # Pass 1: update row_max - old_row_max = row_max - frg_cnt = 4 - frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt - tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile)) - for j in range(frg_cnt): - for k in range(cute.size(tTMEM_LOADrS_frg, mode=[0])): - row_max = cute.arch.fmax(row_max, tTMEM_LOADrS_frg[k, j] * scale_log2) - - row_max_safe = row_max - if row_max == -cutlass.Float32.inf: - row_max_safe = Float32(0.0) - - # acc_scale used for both row_sum rescale and O rescale. - acc_scale_ = scale_log2 * (old_row_max - row_max_safe) - acc_scale = cute.math.exp2(acc_scale_, fastmath=True) - if old_row_max == -cutlass.Float32.inf: - acc_scale = Float32(0.0) - row_sum *= acc_scale - - # Pass 2: P = exp2((S - new_max) * log2), accumulate row_sum, - # store BF16 P through the FP32-backed register bridge. - rP_words = cute.make_rmem_tensor(tTMEM_STOREcP.shape, self.qk_acc_dtype) - rP_bf16 = cute.make_tensor(cute.recast_ptr(rP_words.iterator, dtype=self.q_dtype), tTMEM_LOADrS.layout) - minus_row_max_scale = (Float32(0.0) - row_max_safe) * scale_log2 - - rP_bf16_frg = cute.logical_divide(rP_bf16, cute.make_layout(frg_tile)) - for j in range(frg_cnt): - for k in range(cute.size(tTMEM_LOADrS_frg, mode=[0])): - tTMEM_LOADrS_frg[k, j] = tTMEM_LOADrS_frg[k, j] * scale_log2 + minus_row_max_scale - tTMEM_LOADrS_frg[k, j] = cute.math.exp2(tTMEM_LOADrS_frg[k, j], fastmath=True) - row_sum = row_sum + tTMEM_LOADrS_frg[k, j] - s_vec = tTMEM_LOADrS_frg[None, j].load() - rP_bf16_frg[None, j].store(s_vec.to(self.q_dtype)) - - cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP) - cute.arch.fence_view_async_tmem_store() - - # O rescale for kt > 0. Reads O written by MMA's PV[kt-1]; - # visibility is provided by s_cons.wait_and_advance above - # (acquires on MMA's S[kt] commit, which orders PV[kt-1] before). - if kt > 0: - for i in range(o_col_tiles): - tTMEM_LOAD_O_i = cute.make_tensor( - tTMEM_LOAD_OtO.iterator + i * corr_tile_size, - tTMEM_LOAD_OtO.layout, - ) - tTMEM_STORE_O_i = cute.make_tensor( - tTMEM_STORE_OtO.iterator + i * corr_tile_size, - tTMEM_STORE_OtO.layout, - ) - tTMrO = cute.make_rmem_tensor(tTMEM_LOAD_OcO.shape, self.acc_dtype) - cute.copy(tiled_tmem_load_o, tTMEM_LOAD_O_i, tTMrO) - cute.arch.fence_view_async_tmem_load() - for k in cutlass.range(cute.size(tTMrO), vectorize=True): - tTMrO[k] = tTMrO[k] * acc_scale - cute.copy(tiled_tmem_store_o, tTMrO, tTMEM_STORE_O_i) - cute.arch.fence_view_async_tmem_store() - - si_handle.release() - softmax_done_bar.arrive() - - # Wait for MMA's last PV to commit before reading O for normalize. - # Without this barrier softmax can race MMA's PV[N-1]. - final_o_bar.arrive_and_wait() - - # Final O = O / row_sum - inv_row_sum = Float32(1.0) / row_sum - for i in range(o_col_tiles): - tTMEM_LOAD_O_i = cute.make_tensor( - tTMEM_LOAD_OtO.iterator + i * corr_tile_size, - tTMEM_LOAD_OtO.layout, - ) - tTMEM_STORE_O_i = cute.make_tensor( - tTMEM_STORE_OtO.iterator + i * corr_tile_size, - tTMEM_STORE_OtO.layout, - ) - tTMrO = cute.make_rmem_tensor(tTMEM_LOAD_OcO.shape, self.acc_dtype) - cute.copy(tiled_tmem_load_o, tTMEM_LOAD_O_i, tTMrO) - cute.arch.fence_view_async_tmem_load() - for k in cutlass.range(cute.size(tTMrO), vectorize=True): - tTMrO[k] = tTMrO[k] * inv_row_sum - cute.copy(tiled_tmem_store_o, tTMrO, tTMEM_STORE_O_i) - cute.arch.fence_view_async_tmem_store() - - # Epilogue: TMEM -> SMEM -> GMEM via TMA store - tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout) - acc_cons_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage) - c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) - c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp) - acc_cons_st = utils.gemm.sm100.epilogue_tma_store(self, tidx, warp_idx, tma_c, tCtO_base, sC, tCgC, epi_tile, 0, const_expr(lambda x: x), (0,0,0), acc_cons_st, acc_pipe, c_pipe) - c_pipe.producer_tail() - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def test(): - torch.manual_seed(42) - for n in [128, 256, 512, 1024]: - torch.manual_seed(42) - m, hd = 128, HEAD_DIM - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(n, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(n, hd, dtype=torch.bfloat16, device='cuda') - v_kernel = v.unsqueeze(-1) - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - - qf = q[:, :, 0].float() - kf = k[:, :, 0].float() - scale = 1.0 / math.sqrt(hd) - attn = qf @ kf.T * scale - attn = torch.softmax(attn, dim=-1) - ref = attn @ v.float() - - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - # Each n requires its own compiled kernel (s_k is compile-time). - kernel = FmhaV3StageCMulti(s_k=n) - print(f'n={n}: Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - print(f'n={n}: tmem s0={kernel.tmem_s0_offset} p0={kernel.tmem_p0_offset} ' - f'o0={kernel.tmem_o0_offset} alloc={kernel.num_tmem_alloc_cols} ' - f'kv_tx_bytes={kernel.kv_tx_bytes}', flush=True) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - - out = c[:, :, 0].float() - cos = torch.nn.functional.cosine_similarity( - out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0) - ).item() - max_abs = (out - ref).abs().max().item() - n_tiles = n // 128 - print(f'FMHA Stage-C Multi n={n} ({n_tiles} kv tiles): ' - f'cos {cos:.6f} max_abs {max_abs:.4f} ' - f'{"PASS" if cos >= 0.99 else "FAIL"}') - if cos < 0.99: - print(f' out[0,:4]={out[0,:4].tolist()}') - print(f' ref[0,:4]={ref[0,:4].tolist()}') - - -if __name__ == '__main__': - test() \ No newline at end of file diff --git a/tests/archive/fmha_v3_stage_c_example4.py b/tests/archive/fmha_v3_stage_c_example4.py deleted file mode 100644 index e902ed94..00000000 --- a/tests/archive/fmha_v3_stage_c_example4.py +++ /dev/null @@ -1,469 +0,0 @@ -""" -FMHA v3 Stage-C Multi-Tile (combined K+V barrier, manual GMEM counter). - -The combined K+V barrier (one acquire per kt, K and V on the same barrier slot, -tx_count = K_bytes + V_bytes) eliminates the K/V interleaving in the pipeline -state. That part is fine. - -What does NOT work — and isn't really documented anywhere — is using the -pipeline handle's `.count` field as a GMEM tile coordinate. Whatever .count -actually is at runtime (phase, wrapped slot index, internal state machine), it -is not a global tile counter and TMA copies do not consume it as one. The -CUTLASS Blackwell FMHA reference confirms this by example: every TMA copy in -the load warp uses a manually-tracked Int32 (q0_coord, kv_coord) for GMEM and -handle.index for SMEM. .count is never used as a coordinate, anywhere. - -Rule of thumb for any CuTeDSL pipeline-driven TMA: handle.index is for SMEM -ring buffer slots; GMEM tile coordinates must be tracked separately, e.g. with -an Int32 register variable that you increment yourself. - -Changes vs the single-tile file: - -1. s_k MUST equal actual n. v_fmha layout uses s_k as the V sequence dim. - -2. kv pipeline carries combined K+V per stage: - - tx_count = K_bytes + V_bytes - - producer: one acquire per kt, K and V copies share kvh.barrier - - producer: GMEM coord is a manual Int32 counter (kv_coord), incremented - at the end of each iteration. SMEM slot is kvh.index. - - consumer: one wait per kt, kvh.index used for both sK and sV reads - - release happens after PV - -3. O rescale between KV tiles re-enabled (gated on kt > 0). Lives in softmax - body BEFORE softmax_done_bar.arrive(), so MMA's PV[kt] reads a rescaled O. - -4. Explicit MMA→softmax sync before the final normalize. - final_o_bar is a NamedBarrier with 32 MMA + 128 softmax threads. MMA - .arrive() between acc_pipe.producer_commit and producer_tail; softmax - .arrive_and_wait() before reading O. Without this, softmax can race - MMA's PV[N-1] and divide a partially-accumulated O by row_sum. - -5. final_o_bar.arrive() must come BEFORE acc_pipe.producer_tail in the MMA - warp. producer_tail blocks waiting for the consumer to release the stage; - the consumer release happens inside epilogue_tma_store, which softmax can - only reach after passing final_o_bar. Reverse the order and the kernel - deadlocks. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct -import math - -HEAD_DIM = 64 - - -class FmhaV3StageCMulti: - def __init__(self, s_k=128, scale_softmax=None): - # s_k MUST equal actual sequence length n. - self.s_k = s_k - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0,1,2,3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.num_c_stage = 2 - self.kv_stage = 2; self.q_stage = 1; self.num_c_stage = 2 - self.scale_softmax = scale_softmax if scale_softmax is not None else 1.0 / math.sqrt(HEAD_DIM) - self.scale_softmax_log2 = self.scale_softmax * math.log2(math.e) - - def _setup(self, qk_mma, pv_mma): - qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (128, 128, qk_ik * 4) - pv_ik = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (128, HEAD_DIM, pv_ik * (128 // pv_ik)) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0]//cute.size(qk_mma.thr_id.shape), HEAD_DIM, self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - self.q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, self.q_dtype, self.q_stage) - self.k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, self.q_dtype, self.kv_stage) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, self.kv_stage) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - qk_thr = qk_mma.get_slice(0); qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - pv_thr = pv_mma.get_slice(0); pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - self.tmem_s0_offset = 0; self.tmem_p0_offset = 32 - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - p_end = self.tmem_p0_offset + p_cols_fp32 - s_cols = self.qk_mma_tiler[1] - o_after = max(s_cols, p_end) - self.tmem_o0_offset = ((o_after + 31) // 32) * 32 - o_cols = find_tmem_tensor_col_offset(tOtO) - total = self.tmem_o0_offset + o_cols - self.num_tmem_alloc_cols = 1 - while self.num_tmem_alloc_cols < total: - self.num_tmem_alloc_cols *= 2 - cta = cute.size(qk_mma.thr_id.shape) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)) - k_s = cute.slice_(self.k_smem_s,(None,None,None,0)) - v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - self.q_tx_bytes = cute.size_in_bytes(self.q_dtype, q_s) * cta - # Combined barrier: tx_count covers BOTH K and V transfers per acquire. - self.kv_tx_bytes = (cute.size_in_bytes(self.q_dtype, k_s) + - cute.size_in_bytes(self.q_dtype, v_s)) * cta - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - v_fmha = cute.make_tensor( - v.iterator, - cute.make_layout( - (HEAD_DIM, self.s_k, 1), - stride=(1, HEAD_DIM, HEAD_DIM * self.s_k), - ), - ) - self.v_major = LayoutEnum.from_tensor(v_fmha).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, self.a_major, self.b_major, self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, self.qk_acc_dtype, self.cta_group, (128,HEAD_DIM), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)); v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id),q,q_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id),k,k_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v_fmha,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape) - epi_s = cute.select(self.c_smem_s,mode=[0,1]) - tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile) - self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.v_smem_s,self.p_tmem_s,self.c_smem_s,self.epi_tile).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, q_smem_s, k_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx,_,_ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k); cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - q_bar: cute.struct.MemRange[cutlass.Int64, self.q_stage*2] - kv_bar: cute.struct.MemRange[cutlass.Int64, self.kv_stage*2] - s_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage*2] - tmem_dealloc: cutlass.Int64; holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - qp,qc = pipeline.PipelineTmaUmma.create(barrier_storage=st.q_bar.data_ptr(),num_stages=self.q_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.q_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - # Combined K+V pipeline: each stage carries BOTH K and V loaded together. - kvp,kvc = pipeline.PipelineTmaUmma.create(barrier_storage=st.kv_bar.data_ptr(),num_stages=self.kv_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.kv_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - s_prod,s_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.s_bar.data_ptr(),num_stages=1,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.epilogue_warp_id))).make_participants() - softmax_done_bar = pipeline.NamedBarrier(barrier_id=3, num_threads=32 + 32*len(self.epilogue_warp_id)) - # Final-O sync: MMA arrives between producer_commit and producer_tail; - # softmax arrives_and_waits before reading O for the final normalize. - final_o_bar = pipeline.NamedBarrier(barrier_id=4, num_threads=32 + 32*len(self.epilogue_warp_id)) - acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(),num_stages=self.num_acc_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,len(self.epilogue_warp_id)),cta_layout_vmnk=cl_vmnk,defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2,num_threads=32*len((self.mma_warp_id,*self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr,barrier_for_retrieve=tmem_bar,allocator_warp_id=self.epilogue_warp_id[0],is_two_cta=cute.size(qk_mma.thr_id.shape)==2,two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk,is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype,layout=q_smem_s.outer,byte_alignment=128,swizzle=q_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype,layout=k_smem_s.outer,byte_alignment=128,swizzle=k_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype,layout=v_smem_s.outer,byte_alignment=128,swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype,layout=c_smem_s.outer,byte_alignment=128,swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None)) - gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None)) - gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None)) - gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None)) - n_kv_tiles = cute.size(gK, mode=[3]) - - qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - tCgV = pv_thr.partition_B(gV); tCgC = pv_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,0,None,0)).shape) - tAsQ,tAgQ = cpasync.tma_partition(tma_q,0,a_lay,cute.group_modes(sQ,0,3),cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape) - tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3)) - tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)]; tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None,None,None,0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ===== TMA LOAD warp ===== - # GMEM tile coordinate: manual Int32 counter (kv_coord). SMEM slot: - # kvh.index. The pipeline state .count field is NOT a usable GMEM - # coordinate — CUTLASS reference FMHA tracks its own kv_coord variable - # for the exact same reason; whatever .count actually means at runtime - # (phase, wrapped slot, internal state) it's not what tBgK expects. - if warp_idx == self.tma_warp_id: - qp.reset(); qh = qp.acquire_and_advance() - cute.copy(tma_q, tAgQ[(None, Int32(0))], tAsQ[(None, qh.index)], tma_bar_ptr=qh.barrier) - qp.tail() - kvp.reset(); pk = kvp.try_acquire() - kv_coord = Int32(0) - for kt in cutlass.range(n_kv_tiles, unroll=1): - kvh = kvp.acquire_and_advance(pk) - cute.copy(tma_k, tBgK[(None, kv_coord)], tBsK[(None, kvh.index)], tma_bar_ptr=kvh.barrier) - cute.copy(tma_v, tVgV[(None, kv_coord)], tVsV[(None, kvh.index)], tma_bar_ptr=kvh.barrier) - kv_coord += 1 - pk = cutlass.Boolean(1) - kvp.tail() - - # ===== MMA warp ===== - # One wait per kt; same slot index used for both K (QK) and V (PV). - # Release happens AFTER PV — combined slot stays held across QK+PV. - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - qc.reset(); qh = qc.wait_and_advance(); qh.release() - kvc.reset(); pk = kvc.try_wait() - acc_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_stage) - acc_pipe.producer_acquire(acc_st) - for kt in range(n_kv_tiles): - kvh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - sh = s_prod.acquire_and_advance() - qk_mma.set(tcgen05.Field.ACCUMULATE, False) - for kb in cutlass.range(cute.size(tCrQ, mode=[2]), unroll_full=True): - cute.gemm(qk_mma, tStS0, tCrQ[(None,None,kb,0)], tCrK[(None,None,kb,kvh.index)], tStS0) - qk_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - sh.commit() - softmax_done_bar.arrive_and_wait() - pv_mma.set(tcgen05.Field.ACCUMULATE, kt != 0) - for kb in cutlass.range(cute.size(tOrP0, mode=[2]), unroll_full=True): - cute.gemm(pv_mma, tOtO0, tOrP0[(None,None,kb)], tCrV[(None,None,kb,kvh.index)], tOtO0) - pv_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - kvh.release() - acc_pipe.producer_commit(acc_st); acc_st.advance() - # Signal softmax FIRST so it can run normalize + epilogue. Then - # wait for the epilogue's consumer-release in producer_tail. - # Reverse order deadlocks: producer_tail blocks waiting for - # consumer release; softmax blocks at final_o_bar waiting for - # MMA arrive; the epilogue (which does the release) is gated - # behind softmax's final_o_bar wait. Cycle. - final_o_bar.arrive() - acc_pipe.producer_tail(acc_st) - - # ===== SOFTMAX + EPILOGUE warps ===== - if warp_idx < self.mma_warp_id: - tmem.allocate(self.num_tmem_alloc_cols) - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) - sfw_idx = tidx % (32 * len(self.epilogue_warp_id)) - - # S load - tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS0) - thr_load = tiled_tmem_load.get_slice(sfw_idx) - tTMEM_LOADtS = thr_load.partition_S(tStS0) - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tTMEM_LOADcS = thr_load.partition_D(tScS) - - # P store - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - tStP_layout = cute.composition(tStS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tStP0 = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStP_layout) - tmem_store_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStP0) - thr_store = tiled_tmem_store.get_slice(sfw_idx) - tTMEM_STOREtP = thr_store.partition_D(tStP0) - tScP_layout = cute.composition(tScS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tScP = cute.make_tensor(tScS.iterator, tScP_layout) - tTMEM_STOREcP = thr_store.partition_S(tScP) - - # O rescale / normalize path - cO = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1])) - tOcO = pv_thr.partition_C(cO) - corr_tile_size = 16 - tOtO_i_layout = cute.composition(tOtO.layout, cute.make_layout((128, corr_tile_size))) - tOcO_i_layout = cute.composition(tOcO.layout, cute.make_layout((128, corr_tile_size))) - tOtO_i = cute.make_tensor(tOtO.iterator, tOtO_i_layout) - tOcO_i = cute.make_tensor(tOcO.iterator, tOcO_i_layout) - tmem_load_o_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), self.acc_dtype) - tmem_store_o_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), self.acc_dtype) - tiled_tmem_load_o = tcgen05.make_tmem_copy(tmem_load_o_atom, tOtO_i) - tiled_tmem_store_o = tcgen05.make_tmem_copy(tmem_store_o_atom, tOtO_i) - thr_load_o = tiled_tmem_load_o.get_slice(sfw_idx) - thr_store_o = tiled_tmem_store_o.get_slice(sfw_idx) - tTMEM_LOAD_OtO = thr_load_o.partition_S(tOtO_i) - tTMEM_LOAD_OcO = thr_load_o.partition_D(tOcO_i) - tTMEM_STORE_OtO = thr_store_o.partition_D(tOtO_i) - - o_col_tiles = self.pv_mma_tiler[1] // corr_tile_size - - row_max = -Float32.inf - row_sum = Float32(0.0) - scale_log2 = Float32(self.scale_softmax_log2) - - for kt in range(n_kv_tiles): - si_handle = s_cons.wait_and_advance() - - # Load S[kt] - tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS) - cute.arch.fence_view_async_tmem_load() - - # Pass 1: update row_max - old_row_max = row_max - frg_cnt = 4 - frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt - tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile)) - for j in range(frg_cnt): - for k in range(cute.size(tTMEM_LOADrS_frg, mode=[0])): - row_max = cute.arch.fmax(row_max, tTMEM_LOADrS_frg[k, j] * scale_log2) - - row_max_safe = row_max - if row_max == -cutlass.Float32.inf: - row_max_safe = Float32(0.0) - - # acc_scale used for both row_sum rescale and O rescale. - acc_scale_ = scale_log2 * (old_row_max - row_max_safe) - acc_scale = cute.math.exp2(acc_scale_, fastmath=True) - if old_row_max == -cutlass.Float32.inf: - acc_scale = Float32(0.0) - row_sum *= acc_scale - - # Pass 2: P = exp2((S - new_max) * log2), accumulate row_sum, - # store BF16 P through the FP32-backed register bridge. - rP_words = cute.make_rmem_tensor(tTMEM_STOREcP.shape, self.qk_acc_dtype) - rP_bf16 = cute.make_tensor(cute.recast_ptr(rP_words.iterator, dtype=self.q_dtype), tTMEM_LOADrS.layout) - minus_row_max_scale = (Float32(0.0) - row_max_safe) * scale_log2 - - rP_bf16_frg = cute.logical_divide(rP_bf16, cute.make_layout(frg_tile)) - for j in range(frg_cnt): - for k in range(cute.size(tTMEM_LOADrS_frg, mode=[0])): - tTMEM_LOADrS_frg[k, j] = tTMEM_LOADrS_frg[k, j] * scale_log2 + minus_row_max_scale - tTMEM_LOADrS_frg[k, j] = cute.math.exp2(tTMEM_LOADrS_frg[k, j], fastmath=True) - row_sum = row_sum + tTMEM_LOADrS_frg[k, j] - s_vec = tTMEM_LOADrS_frg[None, j].load() - rP_bf16_frg[None, j].store(s_vec.to(self.q_dtype)) - - cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP) - cute.arch.fence_view_async_tmem_store() - - # O rescale for kt > 0. Reads O written by MMA's PV[kt-1]; - # visibility is provided by s_cons.wait_and_advance above - # (acquires on MMA's S[kt] commit, which orders PV[kt-1] before). - if kt > 0: - for i in range(o_col_tiles): - tTMEM_LOAD_O_i = cute.make_tensor( - tTMEM_LOAD_OtO.iterator + i * corr_tile_size, - tTMEM_LOAD_OtO.layout, - ) - tTMEM_STORE_O_i = cute.make_tensor( - tTMEM_STORE_OtO.iterator + i * corr_tile_size, - tTMEM_STORE_OtO.layout, - ) - tTMrO = cute.make_rmem_tensor(tTMEM_LOAD_OcO.shape, self.acc_dtype) - cute.copy(tiled_tmem_load_o, tTMEM_LOAD_O_i, tTMrO) - cute.arch.fence_view_async_tmem_load() - for k in cutlass.range(cute.size(tTMrO), vectorize=True): - tTMrO[k] = tTMrO[k] * acc_scale - cute.copy(tiled_tmem_store_o, tTMrO, tTMEM_STORE_O_i) - cute.arch.fence_view_async_tmem_store() - - si_handle.release() - softmax_done_bar.arrive() - - # Wait for MMA's last PV to commit before reading O for normalize. - # Without this barrier softmax can race MMA's PV[N-1]. - final_o_bar.arrive_and_wait() - - # Final O = O / row_sum - inv_row_sum = Float32(1.0) / row_sum - for i in range(o_col_tiles): - tTMEM_LOAD_O_i = cute.make_tensor( - tTMEM_LOAD_OtO.iterator + i * corr_tile_size, - tTMEM_LOAD_OtO.layout, - ) - tTMEM_STORE_O_i = cute.make_tensor( - tTMEM_STORE_OtO.iterator + i * corr_tile_size, - tTMEM_STORE_OtO.layout, - ) - tTMrO = cute.make_rmem_tensor(tTMEM_LOAD_OcO.shape, self.acc_dtype) - cute.copy(tiled_tmem_load_o, tTMEM_LOAD_O_i, tTMrO) - cute.arch.fence_view_async_tmem_load() - for k in cutlass.range(cute.size(tTMrO), vectorize=True): - tTMrO[k] = tTMrO[k] * inv_row_sum - cute.copy(tiled_tmem_store_o, tTMrO, tTMEM_STORE_O_i) - cute.arch.fence_view_async_tmem_store() - - # Epilogue: TMEM -> SMEM -> GMEM via TMA store - tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout) - acc_cons_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage) - c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) - c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp) - acc_cons_st = utils.gemm.sm100.epilogue_tma_store(self, tidx, warp_idx, tma_c, tCtO_base, sC, tCgC, epi_tile, 0, const_expr(lambda x: x), (0,0,0), acc_cons_st, acc_pipe, c_pipe) - c_pipe.producer_tail() - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def test(): - torch.manual_seed(42) - for n in [128, 256, 512, 1024]: - torch.manual_seed(42) - m, hd = 128, HEAD_DIM - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(n, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(n, hd, dtype=torch.bfloat16, device='cuda') - v_kernel = v.unsqueeze(-1) - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - - qf = q[:, :, 0].float() - kf = k[:, :, 0].float() - scale = 1.0 / math.sqrt(hd) - attn = qf @ kf.T * scale - attn = torch.softmax(attn, dim=-1) - ref = attn @ v.float() - - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - # Each n requires its own compiled kernel (s_k is compile-time). - kernel = FmhaV3StageCMulti(s_k=n) - print(f'n={n}: Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - print(f'n={n}: tmem s0={kernel.tmem_s0_offset} p0={kernel.tmem_p0_offset} ' - f'o0={kernel.tmem_o0_offset} alloc={kernel.num_tmem_alloc_cols} ' - f'kv_tx_bytes={kernel.kv_tx_bytes}', flush=True) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - - out = c[:, :, 0].float() - cos = torch.nn.functional.cosine_similarity( - out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0) - ).item() - max_abs = (out - ref).abs().max().item() - n_tiles = n // 128 - print(f'FMHA Stage-C Multi n={n} ({n_tiles} kv tiles): ' - f'cos {cos:.6f} max_abs {max_abs:.4f} ' - f'{"PASS" if cos >= 0.99 else "FAIL"}') - if cos < 0.99: - print(f' out[0,:4]={out[0,:4].tolist()}') - print(f' ref[0,:4]={ref[0,:4].tolist()}') - - -if __name__ == '__main__': - test() \ No newline at end of file diff --git a/tests/archive/fmha_v3_stage_c_example5.py b/tests/archive/fmha_v3_stage_c_example5.py deleted file mode 100644 index a13938db..00000000 --- a/tests/archive/fmha_v3_stage_c_example5.py +++ /dev/null @@ -1,469 +0,0 @@ -""" -FMHA v3 Stage-C Multi-Tile (combined K+V barrier, manual GMEM counter). - -The combined K+V barrier (one acquire per kt, K and V on the same barrier slot, -tx_count = K_bytes + V_bytes) eliminates the K/V interleaving in the pipeline -state. That part is fine. - -What does NOT work — and isn't really documented anywhere — is using the -pipeline handle's `.count` field as a GMEM tile coordinate. Whatever .count -actually is at runtime (phase, wrapped slot index, internal state machine), it -is not a global tile counter and TMA copies do not consume it as one. The -CUTLASS Blackwell FMHA reference confirms this by example: every TMA copy in -the load warp uses a manually-tracked Int32 (q0_coord, kv_coord) for GMEM and -handle.index for SMEM. .count is never used as a coordinate, anywhere. - -Rule of thumb for any CuTeDSL pipeline-driven TMA: handle.index is for SMEM -ring buffer slots; GMEM tile coordinates must be tracked separately, e.g. with -an Int32 register variable that you increment yourself. - -Changes vs the single-tile file: - -1. s_k MUST equal actual n. v_fmha layout uses s_k as the V sequence dim. - -2. kv pipeline carries combined K+V per stage: - - tx_count = K_bytes + V_bytes - - producer: one acquire per kt, K and V copies share kvh.barrier - - producer: GMEM coord is a manual Int32 counter (kv_coord), incremented - at the end of each iteration. SMEM slot is kvh.index. - - consumer: one wait per kt, kvh.index used for both sK and sV reads - - release happens after PV - -3. O rescale between KV tiles re-enabled (gated on kt > 0). Lives in softmax - body BEFORE softmax_done_bar.arrive(), so MMA's PV[kt] reads a rescaled O. - -4. Explicit MMA→softmax sync before the final normalize. - final_o_bar is a NamedBarrier with 32 MMA + 128 softmax threads. MMA - .arrive() between acc_pipe.producer_commit and producer_tail; softmax - .arrive_and_wait() before reading O. Without this, softmax can race - MMA's PV[N-1] and divide a partially-accumulated O by row_sum. - -5. final_o_bar.arrive() must come BEFORE acc_pipe.producer_tail in the MMA - warp. producer_tail blocks waiting for the consumer to release the stage; - the consumer release happens inside epilogue_tma_store, which softmax can - only reach after passing final_o_bar. Reverse the order and the kernel - deadlocks. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct -import math - -HEAD_DIM = 64 - - -class FmhaV3StageCMulti: - def __init__(self, s_k=128, scale_softmax=None): - # s_k MUST equal actual sequence length n. - self.s_k = s_k - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0,1,2,3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.num_c_stage = 2 - self.kv_stage = 2; self.q_stage = 1; self.num_c_stage = 2 - self.scale_softmax = scale_softmax if scale_softmax is not None else 1.0 / math.sqrt(HEAD_DIM) - self.scale_softmax_log2 = self.scale_softmax * math.log2(math.e) - - def _setup(self, qk_mma, pv_mma): - qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (128, 128, qk_ik * 4) - pv_ik = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (128, HEAD_DIM, pv_ik * (128 // pv_ik)) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0]//cute.size(qk_mma.thr_id.shape), HEAD_DIM, self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - self.q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, self.q_dtype, self.q_stage) - self.k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, self.q_dtype, self.kv_stage) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, self.kv_stage) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - qk_thr = qk_mma.get_slice(0); qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - pv_thr = pv_mma.get_slice(0); pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - self.tmem_s0_offset = 0; self.tmem_p0_offset = 32 - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - p_end = self.tmem_p0_offset + p_cols_fp32 - s_cols = self.qk_mma_tiler[1] - o_after = max(s_cols, p_end) - self.tmem_o0_offset = ((o_after + 31) // 32) * 32 - o_cols = find_tmem_tensor_col_offset(tOtO) - total = self.tmem_o0_offset + o_cols - self.num_tmem_alloc_cols = 1 - while self.num_tmem_alloc_cols < total: - self.num_tmem_alloc_cols *= 2 - cta = cute.size(qk_mma.thr_id.shape) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)) - k_s = cute.slice_(self.k_smem_s,(None,None,None,0)) - v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - self.q_tx_bytes = cute.size_in_bytes(self.q_dtype, q_s) * cta - # Combined barrier: tx_count covers BOTH K and V transfers per acquire. - self.kv_tx_bytes = (cute.size_in_bytes(self.q_dtype, k_s) + - cute.size_in_bytes(self.q_dtype, v_s)) * cta - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - v_fmha = cute.make_tensor( - v.iterator, - cute.make_layout( - (HEAD_DIM, self.s_k, 1), - stride=(1, HEAD_DIM, HEAD_DIM * self.s_k), - ), - ) - self.v_major = LayoutEnum.from_tensor(v_fmha).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, self.a_major, self.b_major, self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, self.qk_acc_dtype, self.cta_group, (128,HEAD_DIM), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)); v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id),q,q_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id),k,k_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v_fmha,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape) - epi_s = cute.select(self.c_smem_s,mode=[0,1]) - tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile) - self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.v_smem_s,self.p_tmem_s,self.c_smem_s,self.epi_tile).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, q_smem_s, k_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx,_,_ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k); cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - q_bar: cute.struct.MemRange[cutlass.Int64, self.q_stage*2] - kv_bar: cute.struct.MemRange[cutlass.Int64, self.kv_stage*2] - s_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage*2] - tmem_dealloc: cutlass.Int64; holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - qp,qc = pipeline.PipelineTmaUmma.create(barrier_storage=st.q_bar.data_ptr(),num_stages=self.q_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.q_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - # Combined K+V pipeline: each stage carries BOTH K and V loaded together. - kvp,kvc = pipeline.PipelineTmaUmma.create(barrier_storage=st.kv_bar.data_ptr(),num_stages=self.kv_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.kv_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - s_prod,s_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.s_bar.data_ptr(),num_stages=1,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.epilogue_warp_id))).make_participants() - softmax_done_bar = pipeline.NamedBarrier(barrier_id=3, num_threads=32 + 32*len(self.epilogue_warp_id)) - # Final-O sync: MMA arrives between producer_commit and producer_tail; - # softmax arrives_and_waits before reading O for the final normalize. - final_o_bar = pipeline.NamedBarrier(barrier_id=4, num_threads=32 + 32*len(self.epilogue_warp_id)) - acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(),num_stages=self.num_acc_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,len(self.epilogue_warp_id)),cta_layout_vmnk=cl_vmnk,defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2,num_threads=32*len((self.mma_warp_id,*self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr,barrier_for_retrieve=tmem_bar,allocator_warp_id=self.epilogue_warp_id[0],is_two_cta=cute.size(qk_mma.thr_id.shape)==2,two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk,is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype,layout=q_smem_s.outer,byte_alignment=128,swizzle=q_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype,layout=k_smem_s.outer,byte_alignment=128,swizzle=k_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype,layout=v_smem_s.outer,byte_alignment=128,swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype,layout=c_smem_s.outer,byte_alignment=128,swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None)) - gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None)) - gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None)) - gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None)) - n_kv_tiles = cute.size(gK, mode=[3]) - - qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - tCgV = pv_thr.partition_B(gV); tCgC = pv_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,0,None,0)).shape) - tAsQ,tAgQ = cpasync.tma_partition(tma_q,0,a_lay,cute.group_modes(sQ,0,3),cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape) - tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3)) - tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)]; tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None,None,None,0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ===== TMA LOAD warp ===== - # GMEM tile coordinate: use the cutlass.range induction variable kt - # directly. CuTeDSL's `cutlass.range` doesn't auto-detect a Python `+=` - # rebinding as a loop-carried iter_args update — the JIT traces the - # body once and captures whatever value `kv_coord` had at trace time, - # so an outer `kv_coord = Int32(0)` plus a `kv_coord += 1` inside the - # loop bakes 0 into every iteration's TMA descriptor at runtime. - # The induction variable IS the loop-carried state, properly tracked. - if warp_idx == self.tma_warp_id: - qp.reset(); qh = qp.acquire_and_advance() - cute.copy(tma_q, tAgQ[(None, Int32(0))], tAsQ[(None, qh.index)], tma_bar_ptr=qh.barrier) - qp.tail() - kvp.reset(); pk = kvp.try_acquire() - for kt in cutlass.range(0, n_kv_tiles, 1, unroll=1): - kvh = kvp.acquire_and_advance(pk) - cute.copy(tma_k, tBgK[(None, kt)], tBsK[(None, kvh.index)], tma_bar_ptr=kvh.barrier) - cute.copy(tma_v, tVgV[(None, kt)], tVsV[(None, kvh.index)], tma_bar_ptr=kvh.barrier) - pk = cutlass.Boolean(1) - kvp.tail() - - # ===== MMA warp ===== - # One wait per kt; same slot index used for both K (QK) and V (PV). - # Release happens AFTER PV — combined slot stays held across QK+PV. - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - qc.reset(); qh = qc.wait_and_advance(); qh.release() - kvc.reset(); pk = kvc.try_wait() - acc_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_stage) - acc_pipe.producer_acquire(acc_st) - for kt in range(n_kv_tiles): - kvh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - sh = s_prod.acquire_and_advance() - qk_mma.set(tcgen05.Field.ACCUMULATE, False) - for kb in cutlass.range(cute.size(tCrQ, mode=[2]), unroll_full=True): - cute.gemm(qk_mma, tStS0, tCrQ[(None,None,kb,0)], tCrK[(None,None,kb,kvh.index)], tStS0) - qk_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - sh.commit() - softmax_done_bar.arrive_and_wait() - pv_mma.set(tcgen05.Field.ACCUMULATE, kt != 0) - for kb in cutlass.range(cute.size(tOrP0, mode=[2]), unroll_full=True): - cute.gemm(pv_mma, tOtO0, tOrP0[(None,None,kb)], tCrV[(None,None,kb,kvh.index)], tOtO0) - pv_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - kvh.release() - acc_pipe.producer_commit(acc_st); acc_st.advance() - # Signal softmax FIRST so it can run normalize + epilogue. Then - # wait for the epilogue's consumer-release in producer_tail. - # Reverse order deadlocks: producer_tail blocks waiting for - # consumer release; softmax blocks at final_o_bar waiting for - # MMA arrive; the epilogue (which does the release) is gated - # behind softmax's final_o_bar wait. Cycle. - final_o_bar.arrive() - acc_pipe.producer_tail(acc_st) - - # ===== SOFTMAX + EPILOGUE warps ===== - if warp_idx < self.mma_warp_id: - tmem.allocate(self.num_tmem_alloc_cols) - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) - sfw_idx = tidx % (32 * len(self.epilogue_warp_id)) - - # S load - tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS0) - thr_load = tiled_tmem_load.get_slice(sfw_idx) - tTMEM_LOADtS = thr_load.partition_S(tStS0) - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tTMEM_LOADcS = thr_load.partition_D(tScS) - - # P store - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - tStP_layout = cute.composition(tStS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tStP0 = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStP_layout) - tmem_store_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStP0) - thr_store = tiled_tmem_store.get_slice(sfw_idx) - tTMEM_STOREtP = thr_store.partition_D(tStP0) - tScP_layout = cute.composition(tScS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tScP = cute.make_tensor(tScS.iterator, tScP_layout) - tTMEM_STOREcP = thr_store.partition_S(tScP) - - # O rescale / normalize path - cO = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1])) - tOcO = pv_thr.partition_C(cO) - corr_tile_size = 16 - tOtO_i_layout = cute.composition(tOtO.layout, cute.make_layout((128, corr_tile_size))) - tOcO_i_layout = cute.composition(tOcO.layout, cute.make_layout((128, corr_tile_size))) - tOtO_i = cute.make_tensor(tOtO.iterator, tOtO_i_layout) - tOcO_i = cute.make_tensor(tOcO.iterator, tOcO_i_layout) - tmem_load_o_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), self.acc_dtype) - tmem_store_o_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), self.acc_dtype) - tiled_tmem_load_o = tcgen05.make_tmem_copy(tmem_load_o_atom, tOtO_i) - tiled_tmem_store_o = tcgen05.make_tmem_copy(tmem_store_o_atom, tOtO_i) - thr_load_o = tiled_tmem_load_o.get_slice(sfw_idx) - thr_store_o = tiled_tmem_store_o.get_slice(sfw_idx) - tTMEM_LOAD_OtO = thr_load_o.partition_S(tOtO_i) - tTMEM_LOAD_OcO = thr_load_o.partition_D(tOcO_i) - tTMEM_STORE_OtO = thr_store_o.partition_D(tOtO_i) - - o_col_tiles = self.pv_mma_tiler[1] // corr_tile_size - - row_max = -Float32.inf - row_sum = Float32(0.0) - scale_log2 = Float32(self.scale_softmax_log2) - - for kt in range(n_kv_tiles): - si_handle = s_cons.wait_and_advance() - - # Load S[kt] - tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS) - cute.arch.fence_view_async_tmem_load() - - # Pass 1: update row_max - old_row_max = row_max - frg_cnt = 4 - frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt - tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile)) - for j in range(frg_cnt): - for k in range(cute.size(tTMEM_LOADrS_frg, mode=[0])): - row_max = cute.arch.fmax(row_max, tTMEM_LOADrS_frg[k, j] * scale_log2) - - row_max_safe = row_max - if row_max == -cutlass.Float32.inf: - row_max_safe = Float32(0.0) - - # acc_scale used for both row_sum rescale and O rescale. - acc_scale_ = scale_log2 * (old_row_max - row_max_safe) - acc_scale = cute.math.exp2(acc_scale_, fastmath=True) - if old_row_max == -cutlass.Float32.inf: - acc_scale = Float32(0.0) - row_sum *= acc_scale - - # Pass 2: P = exp2((S - new_max) * log2), accumulate row_sum, - # store BF16 P through the FP32-backed register bridge. - rP_words = cute.make_rmem_tensor(tTMEM_STOREcP.shape, self.qk_acc_dtype) - rP_bf16 = cute.make_tensor(cute.recast_ptr(rP_words.iterator, dtype=self.q_dtype), tTMEM_LOADrS.layout) - minus_row_max_scale = (Float32(0.0) - row_max_safe) * scale_log2 - - rP_bf16_frg = cute.logical_divide(rP_bf16, cute.make_layout(frg_tile)) - for j in range(frg_cnt): - for k in range(cute.size(tTMEM_LOADrS_frg, mode=[0])): - tTMEM_LOADrS_frg[k, j] = tTMEM_LOADrS_frg[k, j] * scale_log2 + minus_row_max_scale - tTMEM_LOADrS_frg[k, j] = cute.math.exp2(tTMEM_LOADrS_frg[k, j], fastmath=True) - row_sum = row_sum + tTMEM_LOADrS_frg[k, j] - s_vec = tTMEM_LOADrS_frg[None, j].load() - rP_bf16_frg[None, j].store(s_vec.to(self.q_dtype)) - - cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP) - cute.arch.fence_view_async_tmem_store() - - # O rescale for kt > 0. Reads O written by MMA's PV[kt-1]; - # visibility is provided by s_cons.wait_and_advance above - # (acquires on MMA's S[kt] commit, which orders PV[kt-1] before). - if kt > 0: - for i in range(o_col_tiles): - tTMEM_LOAD_O_i = cute.make_tensor( - tTMEM_LOAD_OtO.iterator + i * corr_tile_size, - tTMEM_LOAD_OtO.layout, - ) - tTMEM_STORE_O_i = cute.make_tensor( - tTMEM_STORE_OtO.iterator + i * corr_tile_size, - tTMEM_STORE_OtO.layout, - ) - tTMrO = cute.make_rmem_tensor(tTMEM_LOAD_OcO.shape, self.acc_dtype) - cute.copy(tiled_tmem_load_o, tTMEM_LOAD_O_i, tTMrO) - cute.arch.fence_view_async_tmem_load() - for k in cutlass.range(cute.size(tTMrO), vectorize=True): - tTMrO[k] = tTMrO[k] * acc_scale - cute.copy(tiled_tmem_store_o, tTMrO, tTMEM_STORE_O_i) - cute.arch.fence_view_async_tmem_store() - - si_handle.release() - softmax_done_bar.arrive() - - # Wait for MMA's last PV to commit before reading O for normalize. - # Without this barrier softmax can race MMA's PV[N-1]. - final_o_bar.arrive_and_wait() - - # Final O = O / row_sum - inv_row_sum = Float32(1.0) / row_sum - for i in range(o_col_tiles): - tTMEM_LOAD_O_i = cute.make_tensor( - tTMEM_LOAD_OtO.iterator + i * corr_tile_size, - tTMEM_LOAD_OtO.layout, - ) - tTMEM_STORE_O_i = cute.make_tensor( - tTMEM_STORE_OtO.iterator + i * corr_tile_size, - tTMEM_STORE_OtO.layout, - ) - tTMrO = cute.make_rmem_tensor(tTMEM_LOAD_OcO.shape, self.acc_dtype) - cute.copy(tiled_tmem_load_o, tTMEM_LOAD_O_i, tTMrO) - cute.arch.fence_view_async_tmem_load() - for k in cutlass.range(cute.size(tTMrO), vectorize=True): - tTMrO[k] = tTMrO[k] * inv_row_sum - cute.copy(tiled_tmem_store_o, tTMrO, tTMEM_STORE_O_i) - cute.arch.fence_view_async_tmem_store() - - # Epilogue: TMEM -> SMEM -> GMEM via TMA store - tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout) - acc_cons_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage) - c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) - c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp) - acc_cons_st = utils.gemm.sm100.epilogue_tma_store(self, tidx, warp_idx, tma_c, tCtO_base, sC, tCgC, epi_tile, 0, const_expr(lambda x: x), (0,0,0), acc_cons_st, acc_pipe, c_pipe) - c_pipe.producer_tail() - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def test(): - torch.manual_seed(42) - for n in [128, 256, 512, 1024]: - torch.manual_seed(42) - m, hd = 128, HEAD_DIM - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(n, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(n, hd, dtype=torch.bfloat16, device='cuda') - v_kernel = v.unsqueeze(-1) - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - - qf = q[:, :, 0].float() - kf = k[:, :, 0].float() - scale = 1.0 / math.sqrt(hd) - attn = qf @ kf.T * scale - attn = torch.softmax(attn, dim=-1) - ref = attn @ v.float() - - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - # Each n requires its own compiled kernel (s_k is compile-time). - kernel = FmhaV3StageCMulti(s_k=n) - print(f'n={n}: Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - print(f'n={n}: tmem s0={kernel.tmem_s0_offset} p0={kernel.tmem_p0_offset} ' - f'o0={kernel.tmem_o0_offset} alloc={kernel.num_tmem_alloc_cols} ' - f'kv_tx_bytes={kernel.kv_tx_bytes}', flush=True) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - - out = c[:, :, 0].float() - cos = torch.nn.functional.cosine_similarity( - out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0) - ).item() - max_abs = (out - ref).abs().max().item() - n_tiles = n // 128 - print(f'FMHA Stage-C Multi n={n} ({n_tiles} kv tiles): ' - f'cos {cos:.6f} max_abs {max_abs:.4f} ' - f'{"PASS" if cos >= 0.99 else "FAIL"}') - if cos < 0.99: - print(f' out[0,:4]={out[0,:4].tolist()}') - print(f' ref[0,:4]={ref[0,:4].tolist()}') - - -if __name__ == '__main__': - test() \ No newline at end of file diff --git a/tests/archive/fmha_v3_stage_c_example6.py b/tests/archive/fmha_v3_stage_c_example6.py deleted file mode 100644 index b56738af..00000000 --- a/tests/archive/fmha_v3_stage_c_example6.py +++ /dev/null @@ -1,472 +0,0 @@ -""" -FMHA v3 Stage-C Multi-Tile (combined K+V barrier, manual GMEM counter). - -The combined K+V barrier (one acquire per kt, K and V on the same barrier slot, -tx_count = K_bytes + V_bytes) eliminates the K/V interleaving in the pipeline -state. That part is fine. - -What does NOT work — and isn't really documented anywhere — is using the -pipeline handle's `.count` field as a GMEM tile coordinate. Whatever .count -actually is at runtime (phase, wrapped slot index, internal state machine), it -is not a global tile counter and TMA copies do not consume it as one. The -CUTLASS Blackwell FMHA reference confirms this by example: every TMA copy in -the load warp uses a manually-tracked Int32 (q0_coord, kv_coord) for GMEM and -handle.index for SMEM. .count is never used as a coordinate, anywhere. - -Rule of thumb for any CuTeDSL pipeline-driven TMA: handle.index is for SMEM -ring buffer slots; GMEM tile coordinates must be tracked separately, e.g. with -an Int32 register variable that you increment yourself. - -Changes vs the single-tile file: - -1. s_k MUST equal actual n. v_fmha layout uses s_k as the V sequence dim. - -2. kv pipeline carries combined K+V per stage: - - tx_count = K_bytes + V_bytes - - producer: one acquire per kt, K and V copies share kvh.barrier - - producer: GMEM coord is a manual Int32 counter (kv_coord), incremented - at the end of each iteration. SMEM slot is kvh.index. - - consumer: one wait per kt, kvh.index used for both sK and sV reads - - release happens after PV - -3. O rescale between KV tiles re-enabled (gated on kt > 0). Lives in softmax - body BEFORE softmax_done_bar.arrive(), so MMA's PV[kt] reads a rescaled O. - -4. Explicit MMA→softmax sync before the final normalize. - final_o_bar is a NamedBarrier with 32 MMA + 128 softmax threads. MMA - .arrive() between acc_pipe.producer_commit and producer_tail; softmax - .arrive_and_wait() before reading O. Without this, softmax can race - MMA's PV[N-1] and divide a partially-accumulated O by row_sum. - -5. final_o_bar.arrive() must come BEFORE acc_pipe.producer_tail in the MMA - warp. producer_tail blocks waiting for the consumer to release the stage; - the consumer release happens inside epilogue_tma_store, which softmax can - only reach after passing final_o_bar. Reverse the order and the kernel - deadlocks. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct -import math - -HEAD_DIM = 64 - - -class FmhaV3StageCMulti: - def __init__(self, s_k=128, scale_softmax=None): - # s_k MUST equal actual sequence length n. - self.s_k = s_k - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0,1,2,3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.num_c_stage = 2 - self.kv_stage = 2; self.q_stage = 1; self.num_c_stage = 2 - self.scale_softmax = scale_softmax if scale_softmax is not None else 1.0 / math.sqrt(HEAD_DIM) - self.scale_softmax_log2 = self.scale_softmax * math.log2(math.e) - - def _setup(self, qk_mma, pv_mma): - qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (128, 128, qk_ik * 4) - pv_ik = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (128, HEAD_DIM, pv_ik * (128 // pv_ik)) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0]//cute.size(qk_mma.thr_id.shape), HEAD_DIM, self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - self.q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, self.q_dtype, self.q_stage) - self.k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, self.q_dtype, self.kv_stage) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, self.kv_stage) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - qk_thr = qk_mma.get_slice(0); qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - pv_thr = pv_mma.get_slice(0); pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - self.tmem_s0_offset = 0; self.tmem_p0_offset = 32 - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - p_end = self.tmem_p0_offset + p_cols_fp32 - s_cols = self.qk_mma_tiler[1] - o_after = max(s_cols, p_end) - self.tmem_o0_offset = ((o_after + 31) // 32) * 32 - o_cols = find_tmem_tensor_col_offset(tOtO) - total = self.tmem_o0_offset + o_cols - self.num_tmem_alloc_cols = 1 - while self.num_tmem_alloc_cols < total: - self.num_tmem_alloc_cols *= 2 - cta = cute.size(qk_mma.thr_id.shape) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)) - k_s = cute.slice_(self.k_smem_s,(None,None,None,0)) - v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - self.q_tx_bytes = cute.size_in_bytes(self.q_dtype, q_s) * cta - # Combined barrier: tx_count covers BOTH K and V transfers per acquire. - self.kv_tx_bytes = (cute.size_in_bytes(self.q_dtype, k_s) + - cute.size_in_bytes(self.q_dtype, v_s)) * cta - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - v_fmha = cute.make_tensor( - v.iterator, - cute.make_layout( - (HEAD_DIM, self.s_k, 1), - stride=(1, HEAD_DIM, HEAD_DIM * self.s_k), - ), - ) - self.v_major = LayoutEnum.from_tensor(v_fmha).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, self.a_major, self.b_major, self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, self.qk_acc_dtype, self.cta_group, (128,HEAD_DIM), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)); v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id),q,q_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id),k,k_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v_fmha,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape) - epi_s = cute.select(self.c_smem_s,mode=[0,1]) - tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile) - self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.v_smem_s,self.p_tmem_s,self.c_smem_s,self.epi_tile).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, q_smem_s, k_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx,_,_ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k); cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - q_bar: cute.struct.MemRange[cutlass.Int64, self.q_stage*2] - kv_bar: cute.struct.MemRange[cutlass.Int64, self.kv_stage*2] - s_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage*2] - tmem_dealloc: cutlass.Int64; holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - qp,qc = pipeline.PipelineTmaUmma.create(barrier_storage=st.q_bar.data_ptr(),num_stages=self.q_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.q_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - # Combined K+V pipeline: each stage carries BOTH K and V loaded together. - kvp,kvc = pipeline.PipelineTmaUmma.create(barrier_storage=st.kv_bar.data_ptr(),num_stages=self.kv_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.kv_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - s_prod,s_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.s_bar.data_ptr(),num_stages=1,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.epilogue_warp_id))).make_participants() - softmax_done_bar = pipeline.NamedBarrier(barrier_id=3, num_threads=32 + 32*len(self.epilogue_warp_id)) - # Final-O sync: MMA arrives between producer_commit and producer_tail; - # softmax arrives_and_waits before reading O for the final normalize. - final_o_bar = pipeline.NamedBarrier(barrier_id=4, num_threads=32 + 32*len(self.epilogue_warp_id)) - acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(),num_stages=self.num_acc_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,len(self.epilogue_warp_id)),cta_layout_vmnk=cl_vmnk,defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2,num_threads=32*len((self.mma_warp_id,*self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr,barrier_for_retrieve=tmem_bar,allocator_warp_id=self.epilogue_warp_id[0],is_two_cta=cute.size(qk_mma.thr_id.shape)==2,two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk,is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype,layout=q_smem_s.outer,byte_alignment=128,swizzle=q_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype,layout=k_smem_s.outer,byte_alignment=128,swizzle=k_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype,layout=v_smem_s.outer,byte_alignment=128,swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype,layout=c_smem_s.outer,byte_alignment=128,swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None)) - gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None)) - gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None)) - gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None)) - n_kv_tiles = cute.size(gK, mode=[3]) - - qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - tCgV = pv_thr.partition_B(gV); tCgC = pv_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,0,None,0)).shape) - tAsQ,tAgQ = cpasync.tma_partition(tma_q,0,a_lay,cute.group_modes(sQ,0,3),cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape) - tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3)) - tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,None,0,0)]; tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None,None,None,0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ===== TMA LOAD warp ===== - # Combined K+V barrier pattern matching working test_fmha_v3_diag.py. - # K uses (None,None,0,0) pre-slice to keep GMEM tile dim free. - # V uses (None,0,None,0) — GMEM tile dim accessible via kv_coord. - if warp_idx == self.tma_warp_id: - qp.reset(); qh = qp.acquire_and_advance() - cute.copy(tma_q, tAgQ[(None, Int32(0))], tAsQ[(None, qh.index)], tma_bar_ptr=qh.barrier) - qp.tail() - kvp.reset(); pk = kvp.try_acquire() - kv_coord = Int32(0 + 0) - for kt in cutlass.range(n_kv_tiles, unroll=1): - kvh = kvp.acquire_and_advance(pk) - cute.copy(tma_k, tBgK[(None, kv_coord)], tBsK[(None, kvh.index)], tma_bar_ptr=kvh.barrier) - cute.copy(tma_v, tVgV[(None, kv_coord)], tVsV[(None, kvh.index)], tma_bar_ptr=kvh.barrier) - kv_coord += 1 - pk = cutlass.Boolean(1) - kvp.tail() - - # ===== MMA warp ===== - # One wait per kt; same slot index used for both K (QK) and V (PV). - # Release happens AFTER PV — combined slot stays held across QK+PV. - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - qc.reset(); qh = qc.wait_and_advance(); qh.release() - kvc.reset(); pk = kvc.try_wait() - acc_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_stage) - acc_pipe.producer_acquire(acc_st) - for kt in range(n_kv_tiles): - kvh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - sh = s_prod.acquire_and_advance() - qk_mma.set(tcgen05.Field.ACCUMULATE, False) - for kb in cutlass.range(cute.size(tCrQ, mode=[2]), unroll_full=True): - cute.gemm(qk_mma, tStS0, tCrQ[(None,None,kb,0)], tCrK[(None,None,kb,kvh.index)], tStS0) - qk_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - sh.commit() - softmax_done_bar.arrive_and_wait() - pv_mma.set(tcgen05.Field.ACCUMULATE, kt != 0) - for kb in cutlass.range(cute.size(tOrP0, mode=[2]), unroll_full=True): - cute.gemm(pv_mma, tOtO0, tOrP0[(None,None,kb)], tCrV[(None,None,kb,kvh.index)], tOtO0) - pv_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - kvh.release() - acc_pipe.producer_commit(acc_st); acc_st.advance() - # Signal softmax FIRST so it can run normalize + epilogue. Then - # wait for the epilogue's consumer-release in producer_tail. - # Reverse order deadlocks: producer_tail blocks waiting for - # consumer release; softmax blocks at final_o_bar waiting for - # MMA arrive; the epilogue (which does the release) is gated - # behind softmax's final_o_bar wait. Cycle. - final_o_bar.arrive() - acc_pipe.producer_tail(acc_st) - - # ===== SOFTMAX + EPILOGUE warps ===== - if warp_idx < self.mma_warp_id: - tmem.allocate(self.num_tmem_alloc_cols) - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) - sfw_idx = tidx % (32 * len(self.epilogue_warp_id)) - - # S load - tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS0) - thr_load = tiled_tmem_load.get_slice(sfw_idx) - tTMEM_LOADtS = thr_load.partition_S(tStS0) - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tTMEM_LOADcS = thr_load.partition_D(tScS) - - # P store - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - tStP_layout = cute.composition(tStS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tStP0 = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStP_layout) - tmem_store_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStP0) - thr_store = tiled_tmem_store.get_slice(sfw_idx) - tTMEM_STOREtP = thr_store.partition_D(tStP0) - tScP_layout = cute.composition(tScS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tScP = cute.make_tensor(tScS.iterator, tScP_layout) - tTMEM_STOREcP = thr_store.partition_S(tScP) - - # O rescale / normalize path. - # CRITICAL: use tOtO0.iterator (offsetted to actual O location) - # NOT tOtO.iterator (which is base TMEM 0 — where S/P live). - # PV writes O via tOtO0; we must read/write the same region. - cO = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1])) - tOcO = pv_thr.partition_C(cO) - corr_tile_size = 16 - tOtO_i_layout = cute.composition(tOtO.layout, cute.make_layout((128, corr_tile_size))) - tOcO_i_layout = cute.composition(tOcO.layout, cute.make_layout((128, corr_tile_size))) - tOtO_i = cute.make_tensor(tOtO0.iterator, tOtO_i_layout) - tOcO_i = cute.make_tensor(tOcO.iterator, tOcO_i_layout) - tmem_load_o_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), self.acc_dtype) - tmem_store_o_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), self.acc_dtype) - tiled_tmem_load_o = tcgen05.make_tmem_copy(tmem_load_o_atom, tOtO_i) - tiled_tmem_store_o = tcgen05.make_tmem_copy(tmem_store_o_atom, tOtO_i) - thr_load_o = tiled_tmem_load_o.get_slice(sfw_idx) - thr_store_o = tiled_tmem_store_o.get_slice(sfw_idx) - tTMEM_LOAD_OtO = thr_load_o.partition_S(tOtO_i) - tTMEM_LOAD_OcO = thr_load_o.partition_D(tOcO_i) - tTMEM_STORE_OtO = thr_store_o.partition_D(tOtO_i) - - o_col_tiles = self.pv_mma_tiler[1] // corr_tile_size - - row_max = -Float32.inf - row_sum = Float32(0.0) - scale_log2 = Float32(self.scale_softmax_log2) - - for kt in range(n_kv_tiles): - si_handle = s_cons.wait_and_advance() - - # Load S[kt] - tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS) - cute.arch.fence_view_async_tmem_load() - - # Pass 1: update row_max - old_row_max = row_max - frg_cnt = 4 - frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt - tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile)) - for j in range(frg_cnt): - for k in range(cute.size(tTMEM_LOADrS_frg, mode=[0])): - row_max = cute.arch.fmax(row_max, tTMEM_LOADrS_frg[k, j] * scale_log2) - - row_max_safe = row_max - if row_max == -cutlass.Float32.inf: - row_max_safe = Float32(0.0) - - # acc_scale used for both row_sum rescale and O rescale. - # row_max is already in scaled domain (S * scale_log2), so - # acc_scale = exp2(old_max - new_max) with no extra scale_log2. - acc_scale_ = old_row_max - row_max_safe - acc_scale = cute.math.exp2(acc_scale_, fastmath=True) - if old_row_max == -cutlass.Float32.inf: - acc_scale = Float32(0.0) - row_sum *= acc_scale - - # Pass 2: P = exp2((S - new_max) * log2), accumulate row_sum, - # store BF16 P through the FP32-backed register bridge. - rP_words = cute.make_rmem_tensor(tTMEM_STOREcP.shape, self.qk_acc_dtype) - rP_bf16 = cute.make_tensor(cute.recast_ptr(rP_words.iterator, dtype=self.q_dtype), tTMEM_LOADrS.layout) - minus_row_max = Float32(0.0) - row_max_safe - - rP_bf16_frg = cute.logical_divide(rP_bf16, cute.make_layout(frg_tile)) - for j in range(frg_cnt): - for k in range(cute.size(tTMEM_LOADrS_frg, mode=[0])): - tTMEM_LOADrS_frg[k, j] = tTMEM_LOADrS_frg[k, j] * scale_log2 + minus_row_max - tTMEM_LOADrS_frg[k, j] = cute.math.exp2(tTMEM_LOADrS_frg[k, j], fastmath=True) - row_sum = row_sum + tTMEM_LOADrS_frg[k, j] - s_vec = tTMEM_LOADrS_frg[None, j].load() - rP_bf16_frg[None, j].store(s_vec.to(self.q_dtype)) - - cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP) - cute.arch.fence_view_async_tmem_store() - - # O rescale for kt > 0. Reads O written by MMA's PV[kt-1]; - # visibility is provided by s_cons.wait_and_advance above - # (acquires on MMA's S[kt] commit, which orders PV[kt-1] before). - if kt > 0: - for i in range(o_col_tiles): - tTMEM_LOAD_O_i = cute.make_tensor( - tTMEM_LOAD_OtO.iterator + i * corr_tile_size, - tTMEM_LOAD_OtO.layout, - ) - tTMEM_STORE_O_i = cute.make_tensor( - tTMEM_STORE_OtO.iterator + i * corr_tile_size, - tTMEM_STORE_OtO.layout, - ) - tTMrO = cute.make_rmem_tensor(tTMEM_LOAD_OcO.shape, self.acc_dtype) - cute.copy(tiled_tmem_load_o, tTMEM_LOAD_O_i, tTMrO) - cute.arch.fence_view_async_tmem_load() - for k in cutlass.range(cute.size(tTMrO), vectorize=True): - tTMrO[k] = tTMrO[k] * acc_scale - cute.copy(tiled_tmem_store_o, tTMrO, tTMEM_STORE_O_i) - cute.arch.fence_view_async_tmem_store() - - si_handle.release() - softmax_done_bar.arrive() - - # Wait for MMA's last PV to commit before reading O for normalize. - # Without this barrier softmax can race MMA's PV[N-1]. - final_o_bar.arrive_and_wait() - - # Final O = O / row_sum - inv_row_sum = Float32(1.0) / row_sum - for i in range(o_col_tiles): - tTMEM_LOAD_O_i = cute.make_tensor( - tTMEM_LOAD_OtO.iterator + i * corr_tile_size, - tTMEM_LOAD_OtO.layout, - ) - tTMEM_STORE_O_i = cute.make_tensor( - tTMEM_STORE_OtO.iterator + i * corr_tile_size, - tTMEM_STORE_OtO.layout, - ) - tTMrO = cute.make_rmem_tensor(tTMEM_LOAD_OcO.shape, self.acc_dtype) - cute.copy(tiled_tmem_load_o, tTMEM_LOAD_O_i, tTMrO) - cute.arch.fence_view_async_tmem_load() - for k in cutlass.range(cute.size(tTMrO), vectorize=True): - tTMrO[k] = tTMrO[k] * inv_row_sum - cute.copy(tiled_tmem_store_o, tTMrO, tTMEM_STORE_O_i) - cute.arch.fence_view_async_tmem_store() - - # Epilogue: TMEM -> SMEM -> GMEM via TMA store - tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout) - acc_cons_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage) - c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) - c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp) - acc_cons_st = utils.gemm.sm100.epilogue_tma_store(self, tidx, warp_idx, tma_c, tCtO_base, sC, tCgC, epi_tile, 0, const_expr(lambda x: x), (0,0,0), acc_cons_st, acc_pipe, c_pipe) - c_pipe.producer_tail() - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def test(): - torch.manual_seed(42) - for n in [128, 256, 512, 1024]: - torch.manual_seed(42) - m, hd = 128, HEAD_DIM - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(n, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(n, hd, dtype=torch.bfloat16, device='cuda') - v_kernel = v.unsqueeze(-1) - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - - qf = q[:, :, 0].float() - kf = k[:, :, 0].float() - scale = 1.0 / math.sqrt(hd) - attn = qf @ kf.T * scale - attn = torch.softmax(attn, dim=-1) - ref = attn @ v.float() - - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - # Each n requires its own compiled kernel (s_k is compile-time). - kernel = FmhaV3StageCMulti(s_k=n) - print(f'n={n}: Compiling... [SLICE_FIX_v5]', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - print(f'n={n}: tmem s0={kernel.tmem_s0_offset} p0={kernel.tmem_p0_offset} ' - f'o0={kernel.tmem_o0_offset} alloc={kernel.num_tmem_alloc_cols} ' - f'kv_tx_bytes={kernel.kv_tx_bytes}', flush=True) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - - out = c[:, :, 0].float() - cos = torch.nn.functional.cosine_similarity( - out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0) - ).item() - max_abs = (out - ref).abs().max().item() - n_tiles = n // 128 - print(f'FMHA Stage-C Multi n={n} ({n_tiles} kv tiles): ' - f'cos {cos:.6f} max_abs {max_abs:.4f} ' - f'{"PASS" if cos >= 0.99 else "FAIL"}') - if cos < 0.99: - print(f' out[0,:4]={out[0,:4].tolist()}') - print(f' ref[0,:4]={ref[0,:4].tolist()}') - - -if __name__ == '__main__': - test() \ No newline at end of file diff --git a/tests/archive/fmha_v3_stage_c_example8.py b/tests/archive/fmha_v3_stage_c_example8.py deleted file mode 100644 index bf7ce364..00000000 --- a/tests/archive/fmha_v3_stage_c_example8.py +++ /dev/null @@ -1,486 +0,0 @@ -""" -FMHA v3 Stage-C Multi-Tile (paired TMEM/SMEM atoms, reference-style epilogue). - -Two structural rules we had to learn the hard way: - -(A) Pipeline handle's `.count` is NOT a GMEM tile coordinate. Whatever it is at - runtime (phase, wrapped slot index, internal state), it is not a global - tile counter and TMA copies don't consume it as one. Use the loop - induction variable for GMEM, handle.index for SMEM. - -(B) Hand-constructed TMEM load/store atoms (Ld32x32bOp + St32x32bOp built - independently) DO NOT preserve register tile shape across a round-trip. - Use paired atoms (or, as we discovered: independently constructed atoms - DO work if they're built from the SAME `Repetition(N)` count — the - Ld32x32bOp(Rep(16)) + St32x32bOp(Rep(16)) pair preserves the register - tile shape exactly because the atom width matches). This is what the - CUTLASS Blackwell FMHA reference does in `correction_rescale`. - -(C) Multi-tile GMEM indexing: `kt` from cutlass.range constant-folds at trace - time, so all TMA loads address tile 0. Workaround: track an Int32 - coordinate manually, BUT seed it from an SSA expression - (`n_kv_tiles - n_kv_tiles`) rather than a literal `Int32(0)`, so the JIT - sees it as a runtime register and propagates the `+= 1` as a tracked - loop-carried iter_args update. - -Kernel structure: - -1. Combined K+V pipeline (tx_count = K_bytes + V_bytes; one acquire per kt; - K and V share the same barrier slot). SMEM slot via kvh.index, GMEM via - manually-tracked kv_coord (SSA-seeded). - -2. Reference-style scaled epilogue: TMEM correction_rescale (O *= 1/row_sum - via paired Ld32x32b + St32x32b atoms), then standard epilogue_tma_store - to send O from TMEM through SMEM to GMEM. - -3. Per-tile O rescale (multiplying existing O by exp2(old_max - new_max) - before PV[kt]) lives in the softmax warp BEFORE softmax_done_bar.arrive(). - Reuses the same paired-atom pattern as the final normalize. - -4. final_o_bar (32 MMA + 128 softmax threads). MMA arrives between - acc_pipe.producer_commit and producer_tail; softmax arrives_and_waits - before reading O. Order: producer_commit → final_o_bar.arrive() → - producer_tail (reverse deadlocks). -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct -import math - -HEAD_DIM = 64 - - -class FmhaV3StageCMulti: - def __init__(self, s_k=128, scale_softmax=None): - # s_k MUST equal actual sequence length n. - self.s_k = s_k - self.n_kv_tiles = s_k // 128 - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0,1,2,3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.num_c_stage = 2 - self.kv_stage = 2; self.q_stage = 1; self.num_c_stage = 2 - self.scale_softmax = scale_softmax if scale_softmax is not None else 1.0 / math.sqrt(HEAD_DIM) - self.scale_softmax_log2 = self.scale_softmax * math.log2(math.e) - - def _setup(self, qk_mma, pv_mma): - qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (128, 128, qk_ik * 4) - pv_ik = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (128, HEAD_DIM, pv_ik * (128 // pv_ik)) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0]//cute.size(qk_mma.thr_id.shape), HEAD_DIM, self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - self.q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, self.q_dtype, self.q_stage) - self.k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, self.q_dtype, self.kv_stage) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, self.kv_stage) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - qk_thr = qk_mma.get_slice(0); qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - pv_thr = pv_mma.get_slice(0); pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - self.tmem_s0_offset = 0; self.tmem_p0_offset = 32 - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - p_end = self.tmem_p0_offset + p_cols_fp32 - s_cols = self.qk_mma_tiler[1] - o_after = max(s_cols, p_end) - self.tmem_o0_offset = ((o_after + 31) // 32) * 32 - o_cols = find_tmem_tensor_col_offset(tOtO) - total = self.tmem_o0_offset + o_cols - self.num_tmem_alloc_cols = 1 - while self.num_tmem_alloc_cols < total: - self.num_tmem_alloc_cols *= 2 - cta = cute.size(qk_mma.thr_id.shape) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)) - k_s = cute.slice_(self.k_smem_s,(None,None,None,0)) - v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - self.q_tx_bytes = cute.size_in_bytes(self.q_dtype, q_s) * cta - # Combined barrier: tx_count covers BOTH K and V transfers per acquire. - self.kv_tx_bytes = (cute.size_in_bytes(self.q_dtype, k_s) + - cute.size_in_bytes(self.q_dtype, v_s)) * cta - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - v_fmha = cute.make_tensor( - v.iterator, - cute.make_layout( - (HEAD_DIM, self.s_k, 1), - stride=(1, HEAD_DIM, HEAD_DIM * self.s_k), - ), - ) - self.v_major = LayoutEnum.from_tensor(v_fmha).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, self.a_major, self.b_major, self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, self.qk_acc_dtype, self.cta_group, (128,HEAD_DIM), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)); v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id),q,q_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id),k,k_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v_fmha,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape) - epi_s = cute.select(self.c_smem_s,mode=[0,1]) - tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile) - self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.v_smem_s,self.p_tmem_s,self.c_smem_s,self.epi_tile).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, q_smem_s, k_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx,_,_ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k); cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - q_bar: cute.struct.MemRange[cutlass.Int64, self.q_stage*2] - kv_bar: cute.struct.MemRange[cutlass.Int64, self.kv_stage*2] - s_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage*2] - tmem_dealloc: cutlass.Int64; holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - qp,qc = pipeline.PipelineTmaUmma.create(barrier_storage=st.q_bar.data_ptr(),num_stages=self.q_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.q_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - kvp,kvc = pipeline.PipelineTmaUmma.create(barrier_storage=st.kv_bar.data_ptr(),num_stages=self.kv_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.kv_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - s_prod,s_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.s_bar.data_ptr(),num_stages=1,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.epilogue_warp_id))).make_participants() - softmax_done_bar = pipeline.NamedBarrier(barrier_id=3, num_threads=32 + 32*len(self.epilogue_warp_id)) - final_o_bar = pipeline.NamedBarrier(barrier_id=4, num_threads=32 + 32*len(self.epilogue_warp_id)) - acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(),num_stages=self.num_acc_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,len(self.epilogue_warp_id)),cta_layout_vmnk=cl_vmnk,defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2,num_threads=32*len((self.mma_warp_id,*self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr,barrier_for_retrieve=tmem_bar,allocator_warp_id=self.epilogue_warp_id[0],is_two_cta=cute.size(qk_mma.thr_id.shape)==2,two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk,is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype,layout=q_smem_s.outer,byte_alignment=128,swizzle=q_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype,layout=k_smem_s.outer,byte_alignment=128,swizzle=k_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype,layout=v_smem_s.outer,byte_alignment=128,swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype,layout=c_smem_s.outer,byte_alignment=128,swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None)) - gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None)) - gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None)) - gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None)) - n_kv_tiles = cute.size(gK, mode=[3]) - - qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - tCgV = pv_thr.partition_B(gV); tCgC = pv_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,0,None,0)).shape) - tAsQ,tAgQ = cpasync.tma_partition(tma_q,0,a_lay,cute.group_modes(sQ,0,3),cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape) - tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3)) - tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,None,0,0)]; tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None,None,None,0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ===== TMA LOAD warp ===== - # Multi-tile GMEM indexing trick: - # - kt from cutlass.range constant-folds at trace time → all TMA - # loads address tile 0 in compiled code. This is the actual - # observed behavior in CuTeDSL 4.5.1, not a hypothesis. - # - Manual kv_coord works IF its initial value is an SSA Int32 - # (a runtime register) rather than a literal Int32(0). - # `n_kv_tiles - n_kv_tiles` is an opaque SSA zero — n_kv_tiles is - # itself an SSA value from cute.size(gK, mode=[3]). With the seed - # in SSA, the JIT treats kv_coord as a tracked loop-carried iter - # variable and propagates `kv_coord = kv_coord + 1` properly. - # - Read kv_coord BEFORE the increment; assignment via `=` (not - # augmented `+=`) avoids any in-place mutation ambiguity. - if warp_idx == self.tma_warp_id: - qp.reset(); qh = qp.acquire_and_advance() - cute.copy(tma_q, tAgQ[(None, Int32(0))], tAsQ[(None, qh.index)], tma_bar_ptr=qh.barrier) - qp.tail() - kvp.reset(); pk = kvp.try_acquire() - kv_coord = n_kv_tiles - n_kv_tiles # SSA runtime zero - for kt in cutlass.range(0, n_kv_tiles, 1, unroll=1): - kvh = kvp.acquire_and_advance(pk) - cute.copy(tma_k, tBgK[(None, kv_coord)], tBsK[(None, kvh.index)], tma_bar_ptr=kvh.barrier) - cute.copy(tma_v, tVgV[(None, kv_coord)], tVsV[(None, kvh.index)], tma_bar_ptr=kvh.barrier) - kv_coord = kv_coord + 1 - pk = cutlass.Boolean(1) - kvp.tail() - - # ===== MMA warp ===== - # One wait per kt; same slot index used for both K (QK) and V (PV). - # Release happens AFTER PV — combined slot stays held across QK+PV. - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - qc.reset(); qh = qc.wait_and_advance(); qh.release() - kvc.reset(); pk = kvc.try_wait() - acc_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_stage) - acc_pipe.producer_acquire(acc_st) - for kt in range(n_kv_tiles): - kvh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - sh = s_prod.acquire_and_advance() - qk_mma.set(tcgen05.Field.ACCUMULATE, False) - for kb in cutlass.range(cute.size(tCrQ, mode=[2]), unroll_full=True): - cute.gemm(qk_mma, tStS0, tCrQ[(None,None,kb,0)], tCrK[(None,None,kb,kvh.index)], tStS0) - qk_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - sh.commit() - softmax_done_bar.arrive_and_wait() - pv_mma.set(tcgen05.Field.ACCUMULATE, kt != 0) - for kb in cutlass.range(cute.size(tOrP0, mode=[2]), unroll_full=True): - cute.gemm(pv_mma, tOtO0, tOrP0[(None,None,kb)], tCrV[(None,None,kb,kvh.index)], tOtO0) - pv_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - kvh.release() - acc_pipe.producer_commit(acc_st); acc_st.advance() - final_o_bar.arrive() - acc_pipe.producer_tail(acc_st) - - # ===== SOFTMAX + EPILOGUE warps ===== - if warp_idx < self.mma_warp_id: - tmem.allocate(self.num_tmem_alloc_cols) - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) - sfw_idx = tidx % (32 * len(self.epilogue_warp_id)) - - # S load - tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS0) - thr_load = tiled_tmem_load.get_slice(sfw_idx) - tTMEM_LOADtS = thr_load.partition_S(tStS0) - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tTMEM_LOADcS = thr_load.partition_D(tScS) - - # P store - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - tStP_layout = cute.composition(tStS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tStP0 = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStP_layout) - tmem_store_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStP0) - thr_store = tiled_tmem_store.get_slice(sfw_idx) - tTMEM_STOREtP = thr_store.partition_D(tStP0) - tScP_layout = cute.composition(tScS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tScP = cute.make_tensor(tScS.iterator, tScP_layout) - tTMEM_STOREcP = thr_store.partition_S(tScP) - - # === O rescale path setup (used per-tile AND for final normalize) === - corr_tile_size = 16 - cO = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1])) - tOcO = pv_thr.partition_C(cO) - tOtO_i_layout = cute.composition(tOtO0.layout, cute.make_layout((128, corr_tile_size))) - tOcO_i_layout = cute.composition(tOcO.layout, cute.make_layout((128, corr_tile_size))) - tOtO_i = cute.make_tensor(tOtO0.iterator, tOtO_i_layout) - tOcO_i = cute.make_tensor(tOcO.iterator, tOcO_i_layout) - tmem_load_o_atom = cute.make_copy_atom( - tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), - self.acc_dtype, - ) - tmem_store_o_atom = cute.make_copy_atom( - tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), - self.acc_dtype, - ) - tiled_tmem_load_o = tcgen05.make_tmem_copy(tmem_load_o_atom, tOtO_i) - tiled_tmem_store_o = tcgen05.make_tmem_copy(tmem_store_o_atom, tOtO_i) - thr_tmem_load_o = tiled_tmem_load_o.get_slice(sfw_idx) - thr_tmem_store_o = tiled_tmem_store_o.get_slice(sfw_idx) - tTMEM_LOADtO = thr_tmem_load_o.partition_S(tOtO_i) - tTMEM_LOADcO = thr_tmem_load_o.partition_D(tOcO_i) - tTMEM_STOREtO = thr_tmem_store_o.partition_D(tOtO_i) - n_corr_tiles = HEAD_DIM // corr_tile_size - - row_max = -Float32.inf - row_sum = Float32(0.0) - scale_log2 = Float32(self.scale_softmax_log2) - - # Per-tile softmax loop with online rescale. - for kt in range(n_kv_tiles): - si_handle = s_cons.wait_and_advance() - - # Load S[kt] - tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS) - cute.arch.fence_view_async_tmem_load() - - # Pass 1: update row_max in log2-domain. - old_row_max = row_max - frg_cnt = 4 - frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt - tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile)) - for j in range(frg_cnt): - for k in range(cute.size(tTMEM_LOADrS_frg, mode=[0])): - row_max = cute.arch.fmax(row_max, tTMEM_LOADrS_frg[k, j] * scale_log2) - - row_max_safe = row_max - if row_max == -cutlass.Float32.inf: - row_max_safe = Float32(0.0) - - # acc_scale = exp2(old_max - new_max). On first tile this is 0 - # (old_max = -inf), so row_sum stays 0 and rescale is skipped. - # row_max is already in scaled domain, so no extra scale_log2. - acc_scale_ = old_row_max - row_max_safe - acc_scale = cute.math.exp2(acc_scale_, fastmath=True) - if old_row_max == -cutlass.Float32.inf: - acc_scale = Float32(0.0) - row_sum *= acc_scale - - # Pass 2: P = exp2((S - new_max) * log2), accumulate row_sum, - # cast to BF16 via FP32-backed register bridge. - rP_words = cute.make_rmem_tensor(tTMEM_STOREcP.shape, self.qk_acc_dtype) - rP_bf16 = cute.make_tensor(cute.recast_ptr(rP_words.iterator, dtype=self.q_dtype), tTMEM_LOADrS.layout) - minus_row_max = Float32(0.0) - row_max_safe - - rP_bf16_frg = cute.logical_divide(rP_bf16, cute.make_layout(frg_tile)) - for j in range(frg_cnt): - for k in range(cute.size(tTMEM_LOADrS_frg, mode=[0])): - tTMEM_LOADrS_frg[k, j] = tTMEM_LOADrS_frg[k, j] * scale_log2 + minus_row_max - tTMEM_LOADrS_frg[k, j] = cute.math.exp2(tTMEM_LOADrS_frg[k, j], fastmath=True) - row_sum = row_sum + tTMEM_LOADrS_frg[k, j] - s_vec = tTMEM_LOADrS_frg[None, j].load() - rP_bf16_frg[None, j].store(s_vec.to(self.q_dtype)) - - cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP) - cute.arch.fence_view_async_tmem_store() - - # === Per-tile O rescale: O *= acc_scale for kt > 0 === - # Uses the SAME paired-atom pattern as the final normalize. - # Must run BEFORE softmax_done_bar.arrive() so MMA's PV[kt] - # reads the rescaled O. - # Visibility of MMA's PV[kt-1] writes: provided by - # s_cons.wait_and_advance at the top of this iteration, which - # acquires on MMA's S[kt] commit. S[kt] is sequenced after - # PV[kt-1] in MMA's iteration, so PV[kt-1]'s tmem_store_fence - # has been observed by the time we read O here. - if kt > 0: - for i in range(n_corr_tiles): - tTMEM_LOADtO_i = cute.make_tensor( - tTMEM_LOADtO.iterator + i * corr_tile_size, - tTMEM_LOADtO.layout, - ) - tTMEM_STOREtO_i = cute.make_tensor( - tTMEM_STOREtO.iterator + i * corr_tile_size, - tTMEM_STOREtO.layout, - ) - tTMrO = cute.make_rmem_tensor(tTMEM_LOADcO.shape, self.acc_dtype) - cute.copy(tiled_tmem_load_o, tTMEM_LOADtO_i, tTMrO) - cute.arch.fence_view_async_tmem_load() - for k in cutlass.range(cute.size(tTMrO), vectorize=True): - tTMrO[k] = tTMrO[k] * acc_scale - cute.copy(tiled_tmem_store_o, tTMrO, tTMEM_STOREtO_i) - cute.arch.fence_view_async_tmem_store() - - si_handle.release() - softmax_done_bar.arrive() - - # Wait for MMA's PV[N-1] to commit before reading O for normalize. - final_o_bar.arrive_and_wait() - - # === Final O normalization: O *= 1/row_sum === - inv_row_sum = Float32(1.0) / row_sum - for i in range(n_corr_tiles): - tTMEM_LOADtO_i = cute.make_tensor( - tTMEM_LOADtO.iterator + i * corr_tile_size, - tTMEM_LOADtO.layout, - ) - tTMEM_STOREtO_i = cute.make_tensor( - tTMEM_STOREtO.iterator + i * corr_tile_size, - tTMEM_STOREtO.layout, - ) - tTMrO = cute.make_rmem_tensor(tTMEM_LOADcO.shape, self.acc_dtype) - cute.copy(tiled_tmem_load_o, tTMEM_LOADtO_i, tTMrO) - cute.arch.fence_view_async_tmem_load() - for k in cutlass.range(cute.size(tTMrO), vectorize=True): - tTMrO[k] = tTMrO[k] * inv_row_sum - cute.copy(tiled_tmem_store_o, tTMrO, tTMEM_STOREtO_i) - cute.arch.fence_view_async_tmem_store() - - # Standard epilogue: TMEM → SMEM → GMEM via TMA store. - # O in TMEM is now scaled by 1/row_sum. - tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout) - acc_cons_st = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Consumer, self.num_acc_stage - ) - c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) - c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp) - acc_cons_st = utils.gemm.sm100.epilogue_tma_store( - self, tidx, warp_idx, tma_c, tCtO_base, sC, tCgC, epi_tile, - 0, const_expr(lambda x: x), (0, 0, 0), - acc_cons_st, acc_pipe, c_pipe, - ) - c_pipe.producer_tail() - - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def test(): - torch.manual_seed(42) - for n in [128, 256, 512, 1024]: - torch.manual_seed(42) - m, hd = 128, HEAD_DIM - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(n, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(n, hd, dtype=torch.bfloat16, device='cuda') - v_kernel = v.unsqueeze(-1) - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - - qf = q[:, :, 0].float() - kf = k[:, :, 0].float() - scale = 1.0 / math.sqrt(hd) - attn = qf @ kf.T * scale - attn = torch.softmax(attn, dim=-1) - ref = attn @ v.float() - - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - kernel = FmhaV3StageCMulti(s_k=n) - print(f'n={n}: Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - print(f'n={n}: tmem s0={kernel.tmem_s0_offset} p0={kernel.tmem_p0_offset} ' - f'o0={kernel.tmem_o0_offset} alloc={kernel.num_tmem_alloc_cols} ' - f'kv_tx_bytes={kernel.kv_tx_bytes}', flush=True) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - - out = c[:, :, 0].float() - cos = torch.nn.functional.cosine_similarity( - out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0) - ).item() - max_abs = (out - ref).abs().max().item() - n_tiles = n // 128 - print(f'FMHA Stage-C Multi n={n} ({n_tiles} kv tiles): ' - f'cos {cos:.6f} max_abs {max_abs:.4f} ' - f'{"PASS" if cos >= 0.99 else "FAIL"}') - if cos < 0.99: - print(f' out[0,:4]={out[0,:4].tolist()}') - print(f' ref[0,:4]={ref[0,:4].tolist()}') - - -if __name__ == '__main__': - test() \ No newline at end of file diff --git a/tests/archive/native_stage_c_patch.py b/tests/archive/native_stage_c_patch.py deleted file mode 100644 index 734c0408..00000000 --- a/tests/archive/native_stage_c_patch.py +++ /dev/null @@ -1,351 +0,0 @@ -""" -Native Blackwell Stage-C drop-in implementation for FmhaV3. - -This file is intentionally a patch module rather than a second full copy of the -whole test. It contains the exact role layout, pipeline definitions, and the -native correction helpers to splice into your current test_fmha_v3.py / Stage-B -kernel. - -Architecture: - softmax warps 0-3 : S(TMEM) -> P(TMEM), vec(TMEM) - correction warps 4-7 : vec(TMEM) + O(TMEM) -> corrected O(SMEM) - MMA warp 8 : QK and PV - TMA/load warp 9 : Q/K/V load - epilogue warp 10 : corrected O SMEM -> GMEM - empty warp 11 : tmem dealloc mbar init/helper - -The key design decision: softmax NEVER rescales O. Correction is the only code -path that touches O for online rescale/final normalization. Epilogue never -reads TMEM. -""" - -# ------------------------- -# 1) role layout -# ------------------------- -# Put this in __init__ -SOFTMAX_WARP_IDS = (0, 1, 2, 3) -CORRECTION_WARP_IDS = (4, 5, 6, 7) -MMA_WARP_ID = 8 -TMA_WARP_ID = 9 -EPILOGUE_WARP_ID = 10 -EMPTY_WARP_ID = 11 -THREADS_PER_CTA = 32 * 12 - -# Pipeline stages to add to your class -MMA_SOFTMAX_STAGE = 1 -SOFTMAX_CORR_STAGE = 1 -MMA_CORR_STAGE = 2 -EPI_STAGE = 2 - -# TMEM offsets for a single 128x128 S/P tile and 128x64 O tile. -# vec reuses S region only after the S tile has been loaded by softmax. -TMEM_S0_OFFSET = 0 -TMEM_VEC0_OFFSET = 0 -TMEM_P0_OFFSET = 32 -TMEM_O0_OFFSET = 128 - - -# ------------------------- -# 2) shared storage fields -# ------------------------- -SHARED_STORAGE_FIELDS = r''' -@cute.struct -class SS: - q_bar: cute.struct.MemRange[cutlass.Int64, self.q_stage * 2] - kv_bar: cute.struct.MemRange[cutlass.Int64, self.kv_stage * 2] - mma_s_bar: cute.struct.MemRange[cutlass.Int64, self.mma_softmax_stage * 2] - s_corr_bar: cute.struct.MemRange[cutlass.Int64, self.softmax_corr_stage * 2] - mma_corr_bar: cute.struct.MemRange[cutlass.Int64, self.mma_corr_stage * 2] - corr_epi_bar: cute.struct.MemRange[cutlass.Int64, self.epi_stage * 2] - tmem_dealloc: cute.struct.MemRange[cutlass.Int64, 1] - holding: cutlass.Int32 -''' - - -# ------------------------- -# 3) pipeline creation -# ------------------------- -PIPELINE_CREATION = r''' -def cg_threads(n: int): - return pipeline.CooperativeGroup(pipeline.Agent.Thread, n) - -q_prod, q_cons = pipeline.PipelineTmaUmma.create( - barrier_storage=st.q_bar.data_ptr(), - num_stages=self.q_stage, - producer_group=cg_threads(1), - consumer_group=cg_threads(1), - tx_count=self.q_tx_bytes, - cta_layout_vmnk=cl_vmnk, - defer_sync=True, -).make_participants() - -kv_prod, kv_cons = pipeline.PipelineTmaUmma.create( - barrier_storage=st.kv_bar.data_ptr(), - num_stages=self.kv_stage, - producer_group=cg_threads(1), - consumer_group=cg_threads(1), - tx_count=self.kv_tx_bytes, - cta_layout_vmnk=cl_vmnk, - defer_sync=True, -).make_participants() - -# MMA publishes S; softmax consumes S and releases the handle only after P is ready. -mma_s_prod, mma_s_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_s_bar.data_ptr(), - num_stages=self.mma_softmax_stage, - producer_group=cg_threads(1), - consumer_group=cg_threads(32 * len(self.softmax_warp_ids)), - cta_layout_vmnk=cl_vmnk, - defer_sync=True, -).make_participants() - -# Softmax publishes vec[row] = [old_max,new_max] each tile and [row_sum,row_max] at final. -s_corr_prod, s_corr_cons = pipeline.PipelineAsync.create( - barrier_storage=st.s_corr_bar.data_ptr(), - num_stages=self.softmax_corr_stage, - producer_group=cg_threads(32 * len(self.softmax_warp_ids)), - consumer_group=cg_threads(32 * len(self.correction_warp_ids)), -).make_participants() - -# MMA publishes O after each PV. Correction consumes O for online rescale / final epilog. -mma_corr_prod, mma_corr_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_corr_bar.data_ptr(), - num_stages=self.mma_corr_stage, - producer_group=cg_threads(1), - consumer_group=cg_threads(32 * len(self.correction_warp_ids)), - cta_layout_vmnk=cl_vmnk, - defer_sync=True, -).make_participants() - -# Correction publishes final converted O in SMEM; epilogue warp TMA-stores it. -corr_epi_prod, corr_epi_cons = pipeline.PipelineAsync.create( - barrier_storage=st.corr_epi_bar.data_ptr(), - num_stages=self.epi_stage, - producer_group=cg_threads(32 * len(self.correction_warp_ids)), - consumer_group=cg_threads(32), -).make_participants() - -# TMEM allocation/retrieve participants are softmax + correction + MMA only. -tmem_bar = pipeline.NamedBarrier( - barrier_id=2, - num_threads=32 * len((*self.softmax_warp_ids, *self.correction_warp_ids, self.mma_warp_id)), -) -tmem = utils.TmemAllocator( - st.holding.ptr, - barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.softmax_warp_ids[0], - is_two_cta=cute.size(qk_mma.thr_id.shape) == 2, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr, -) - -# Deallocation mbarrier is NOT a named barrier. Only softmax+correction arrive; -# MMA waits and deallocates after relinquish_tmem_alloc_permit(). -if warp_idx == self.empty_warp_id: - cute.arch.mbarrier_init( - st.tmem_dealloc.data_ptr(), - 32 * len((*self.softmax_warp_ids, *self.correction_warp_ids)), - ) - cute.arch.mbarrier_init_fence() -''' - - -# ------------------------- -# 4) correction helpers -# ------------------------- -CORRECTION_HELPERS = r''' -@cute.jit -def correction_rescale(self, pv_thr, tOtO, scale: Float32): - """Correction warpgroup: O[row,:] *= scale[row]. O stays in TMEM.""" - cO = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1])) - tOcO = pv_thr.partition_C(cO) - - corr_tile_size = 16 - tmem_load_atom = cute.make_copy_atom( - tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), - self.pv_acc_dtype, - ) - tmem_store_atom = cute.make_copy_atom( - tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), - self.pv_acc_dtype, - ) - - tOtO_i_layout = cute.composition(tOtO.layout, cute.make_layout((128, corr_tile_size))) - tOcO_i_layout = cute.composition(tOcO.layout, cute.make_layout((128, corr_tile_size))) - tOtO_i = cute.make_tensor(tOtO.iterator, tOtO_i_layout) - tOcO_i = cute.make_tensor(tOcO.iterator, tOcO_i_layout) - - tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tOtO_i) - tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tOtO_i) - - tidx, _, _ = cute.arch.thread_idx() - thread_idx = tidx % (32 * len(self.correction_warp_ids)) - thr_load = tiled_tmem_load.get_slice(thread_idx) - thr_store = tiled_tmem_store.get_slice(thread_idx) - - tTMEM_LOADtO = thr_load.partition_S(tOtO_i) - tTMEM_LOADcO = thr_load.partition_D(tOcO_i) - tTMEM_STOREtO = thr_store.partition_D(tOtO_i) - - o_col_tiles = self.pv_mma_tiler[1] // corr_tile_size - tTMrO = cute.make_rmem_tensor((tTMEM_LOADcO.shape, o_col_tiles), self.pv_acc_dtype) - for i in range(o_col_tiles): - tTMrO_i_ = tTMrO[None, i] - tTMrO_i_layout = cute.composition(tTMrO_i_.layout, cute.make_layout(tTMrO.shape[0])) - tTMrO_i = cute.make_tensor(tTMrO_i_.iterator, tTMrO_i_layout) - tTMEM_LOADtO_i = cute.make_tensor(tTMEM_LOADtO.iterator + i * corr_tile_size, tTMEM_LOADtO.layout) - tTMEM_STOREtO_i = cute.make_tensor(tTMEM_STOREtO.iterator + i * corr_tile_size, tTMEM_STOREtO.layout) - cute.copy(tiled_tmem_load, tTMEM_LOADtO_i, tTMrO_i) - for j in cutlass.range(cute.size(tTMrO_i), vectorize=True): - tTMrO_i[j] = tTMrO_i[j] * scale - cute.copy(tiled_tmem_store, tTMrO_i, tTMEM_STOREtO_i) - - -@cute.jit -def correction_epilog(self, pv_thr, tOtO, scale: Float32, sO): - """Correction final: load O from TMEM, normalize/convert, write SMEM.""" - cO = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1])) - corr_tile_size = 32 * 8 // self.o_dtype.width - - tOsO = pv_thr.partition_C(sO) - tOcO = pv_thr.partition_C(cO) - tOtO_i = cute.logical_divide(tOtO, cute.make_layout((128, corr_tile_size))) - tOcO_i = cute.logical_divide(tOcO, cute.make_layout((128, corr_tile_size))) - tOsO_i = cute.logical_divide(tOsO, cute.make_layout((128, corr_tile_size))) - - tidx, _, _ = cute.arch.thread_idx() - thread_idx = tidx % (32 * len(self.correction_warp_ids)) - epi_subtile = (self.epi_tile[0], corr_tile_size) - - tmem_copy_atom = utils.sm100.get_tmem_load_op( - self.pv_mma_tiler, - self.c_layout, - self.o_dtype, - self.pv_acc_dtype, - epi_subtile, - use_2cta_instrs=False, - ) - tiled_tmem_load = tcgen05.make_tmem_copy(tmem_copy_atom, tOtO_i[(None, None), 0]) - thr_tmem_load = tiled_tmem_load.get_slice(thread_idx) - - smem_copy_atom = utils.sm100.get_smem_store_op( - self.c_layout, - self.o_dtype, - self.pv_acc_dtype, - tiled_tmem_load, - ) - tiled_smem_store = cute.make_tiled_copy_D(smem_copy_atom, tiled_tmem_load) - - tTMEM_LOADtO = thr_tmem_load.partition_S(tOtO_i[(None, None), None]) - tTMEM_LOADsO = thr_tmem_load.partition_D(tOsO_i[(None, None), None]) - tTMEM_LOADoO = thr_tmem_load.partition_D(tOcO_i[(None, None), None]) - - for i in range(self.pv_mma_tiler[1] // corr_tile_size): - tTMrO = cute.make_rmem_tensor(tTMEM_LOADoO[None, 0, 0, i].shape, self.pv_acc_dtype) - cute.copy(tiled_tmem_load, tTMEM_LOADtO[None, 0, 0, i], tTMrO) - for j in cutlass.range(cute.size(tTMrO), vectorize=True): - tTMrO[j] = tTMrO[j] * scale - tSMrO = cute.make_rmem_tensor(tTMrO.shape, self.o_dtype) - tSMrO.store(tTMrO.load().to(self.o_dtype)) - cute.copy(tiled_smem_store, tSMrO, tTMEM_LOADsO[None, 0, 0, i]) - - cute.arch.fence_proxy("async.shared", space="cta") -''' - - -# ------------------------- -# 5) native Stage-C flow -# ------------------------- -NATIVE_STAGE_C_FLOW = r''' -# MMA loop pseudocode, replacing softmax_done_bar/pv_done_bar: -if is_mma: - tmem.wait_for_alloc() - ... load Q/K/V pipeline setup ... - s_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.mma_softmax_stage) - o_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.mma_corr_stage) - mma_s_prod.producer_acquire(s_state) - mma_corr_prod.producer_acquire(o_state) - for kt in range(n_kv_tiles): - # QK -> S - ... cute.gemm(qk_mma, ..., tStS0) ... - cute.arch.fence_view_async_tmem_store() - mma_s_prod.producer_commit(s_state) - s_state.advance() - - # IMPORTANT: there is no named barrier here. The softmax consumer releases - # the same S handle only after P has been stored, so this is P-ready. - if kt + 1 < n_kv_tiles: - mma_s_prod.producer_acquire(s_state) - - # PV -> O - ... cute.gemm(pv_mma, tOtO0, tOrP0, V, tOtO0) ... - cute.arch.fence_view_async_tmem_store() - mma_corr_prod.producer_commit(o_state) - o_state.advance() - if kt + 1 < n_kv_tiles: - mma_corr_prod.producer_acquire(o_state) - mma_s_prod.producer_tail(s_state) - mma_corr_prod.producer_tail(o_state) - - cute.arch.relinquish_tmem_alloc_permit() - cute.arch.mbarrier_wait(st.tmem_dealloc.data_ptr(), 0) - tmem_ptr = cute.arch.retrieve_tmem_ptr(self.qk_acc_dtype, alignment=16, ptr_to_buffer_holding_addr=st.holding.ptr) - cute.arch.dealloc_tmem(tmem_ptr, Int32(self.num_tmem_alloc_cols)) - - -# Softmax loop: -if is_softmax: - tmem.allocate(self.num_tmem_alloc_cols) - tmem.wait_for_alloc() - vec_handle = s_corr_prod.acquire_and_advance() - row_max = -Float32.inf - row_sum = Float32(0.0) - for kt in range(n_kv_tiles): - si = mma_s_cons.wait_and_advance() - # load S, compute old/new row max, store vec=[old_max,new_max] - # compute P=exp2((S-new_max)*scale), store P to TMEM - # update row_sum - vec_handle.commit() - si.release() # P-ready signal to MMA - vec_handle = s_corr_prod.acquire_and_advance() - # final vec=[row_sum,row_max] - final_si = mma_s_cons.wait_and_advance() - ... store final vec ... - vec_handle.commit() - s_corr_prod.acquire() # balance final pipe step - final_si.release() - cute.arch.mbarrier_arrive(st.tmem_dealloc.data_ptr()) - - -# Correction loop: -if is_correction: - tmem.wait_for_alloc() - first_vec = s_corr_cons.wait_and_advance() - first_vec.release() # first tile has no previous O to rescale - for kt in range(n_kv_tiles - 1): - vec = s_corr_cons.wait_and_advance() - ... load [old_max,new_max] ... - scale = exp2(scale_log2 * (old_max - new_max)) - o = mma_corr_cons.wait_and_advance() - self.correction_rescale(pv_thr, tOtO0, scale) - cute.arch.fence_view_async_tmem_store() - o.release() - vec.release() - final_vec = s_corr_cons.wait_and_advance() - ... load [row_sum,row_max] ... - final_vec.release() - final_o = mma_corr_cons.wait_and_advance() - epi = corr_epi_prod.acquire_and_advance() - self.correction_epilog(pv_thr, tOtO0, self.scale_output / row_sum, sC[(None, None, epi.index)]) - final_o.release() - epi.commit() - cute.arch.mbarrier_arrive(st.tmem_dealloc.data_ptr()) - - -# Epilogue warp: -if is_epi: - h = corr_epi_cons.wait_and_advance() - cute.copy(tma_c, tCsC[(None, h.index)], tCgC_epi[(None, 0)]) - cute.arch.cp_async_bulk_commit_group() - cute.arch.cp_async_bulk_wait_group(0, read=True) - h.release() -''' diff --git a/tests/archive/quick_v3_multitile.py b/tests/archive/quick_v3_multitile.py deleted file mode 100644 index 8687cdcb..00000000 --- a/tests/archive/quick_v3_multitile.py +++ /dev/null @@ -1,53 +0,0 @@ -"""Quick test: run the working test_fmha_v3.py with n=256 to check multi-tile.""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct -import math - -HEAD_DIM = 64 - -# Import the working kernel class -import sys -sys.path.insert(0, '/root/dsv4-nvfp4-workspace/kernel') -from tests.unit.test_fmha_v3 import FmhaV3 - -for n in [128, 256]: - torch.manual_seed(42) - m, hd = 128, HEAD_DIM - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(n, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.ones(n, hd, dtype=torch.bfloat16, device='cuda') # V=ones like the working test - v_kernel = v.unsqueeze(-1) - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - - qf = q[:, :, 0].float() - kf = k[:, :, 0].float() - scale = 1.0 / math.sqrt(hd) - attn = qf @ kf.T * scale - attn = torch.softmax(attn, dim=-1) - ref = attn @ v.float() - - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - kernel = FmhaV3(s_k=n) - print(f'n={n}: Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - - out = c[:, :, 0].float() - cos = torch.nn.functional.cosine_similarity( - out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0) - ).item() - print(f'FMHA v3 n={n}: cos {cos:.6f} {"PASS" if cos >= 0.99 else "FAIL"}') - if cos < 0.99: - print(f' out[0,:4]={out[0,:4].tolist()}') - print(f' ref[0,:4]={ref[0,:4].tolist()}') diff --git a/tests/archive/stage_b_debug5.py b/tests/archive/stage_b_debug5.py deleted file mode 100644 index e0357e8e..00000000 --- a/tests/archive/stage_b_debug5.py +++ /dev/null @@ -1,187 +0,0 @@ -"""Stage B debug v5: Minimal two-MMA with debug printf to find deadlock location.""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -import cuda.bindings.driver as cuda - -class StageBDebug5: - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE; self.use_2cta_instrs = False - self.epilogue_warp_id = (0, 1, 2, 3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.epilog_sync_bar_id = 1; self.num_c_stage = 2 - - def _setup(self, qk_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - self.mma_tiler = self.qk_mma_tiler - self.cta_tile_shape_mnk = tuple(self.qk_mma_tiler) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, BFloat16) - self.num_ab_stage = 1; self.num_acc_stage = 1 - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, BFloat16, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, BFloat16, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(BFloat16, self.c_layout, self.epi_tile, 2) - # Use QK fragment for tmem allocation - acc_shape = qk_mma.partition_shape_C(self.mma_tiler_mn) - tCtAcc_fake = qk_mma.make_fragment_C(cute.append(acc_shape, 1)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(tCtAcc_fake, arch="sm_100") - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = (cute.size_in_bytes(BFloat16, a_smem) + cute.size_in_bytes(BFloat16, b_smem)) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, a, b, c, stream): - self.a_dtype = a.element_type; self.b_dtype = b.element_type; self.c_dtype = c.element_type - self.a_major = LayoutEnum.from_tensor(a).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(b).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.a_major, self.b_major, self.acc_dtype, self.cta_group, self.mma_tiler_mn) - self._setup(qk_mma) - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - tma_a, tma_ta = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - a, a_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_b, tma_tb = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - b, b_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - self._kernel(qk_mma, tma_a, tma_ta, tma_b, tma_tb, tma_c, tma_tc, - self.cluster_layout_vmnk, self.a_smem_s, self.b_smem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[192,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, tma_a, mA, tma_b, mB, tma_c, mC, cl_vmnk, - a_smem_s, b_smem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_a); cpasync.prefetch_descriptor(tma_b); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, len(self.epilogue_warp_id)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2, num_threads=160) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=0, is_two_cta=False, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sA = smem.allocate_tensor(element_type=BFloat16, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sB = smem.allocate_tensor(element_type=BFloat16, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sC = smem.allocate_tensor(element_type=BFloat16, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - gA = cute.local_tile(mA, cute.slice_(self.mma_tiler, (None,0,None)), (None,None,None)) - gB = cute.local_tile(mB, cute.slice_(self.mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gA, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - tCgA = qk_thr.partition_A(gA); tCgB = qk_thr.partition_B(gB); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsA, tAgA = cpasync.tma_partition(tma_a, 0, a_lay, cute.group_modes(sA,0,3), cute.group_modes(tCgA,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsB, tBgB = cpasync.tma_partition(tma_b, 0, b_lay, cute.group_modes(sB,0,3), cute.group_modes(tCgB,0,3)) - tAgA = tAgA[(None,0,None,0)]; tBgB = tBgB[(None,0,None,0)] - - tCrA = qk_mma.make_fragment_A(sA); tCrB = qk_mma.make_fragment_B(sB) - # NO pv_mma.make_fragment_B - - acc_shape = qk_mma.partition_shape_C(self.mma_tiler_mn) - tCtAcc_fake = qk_mma.make_fragment_C(cute.append(acc_shape, 1)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # TMA - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_a, tAgA[(None,h.count)], tAsA[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_b, tBgB[(None,h.count)], tBsB[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else "FAIL"}') - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_128_16_bigP.py b/tests/archive/test_128_16_bigP.py deleted file mode 100644 index fc47cf7a..00000000 --- a/tests/archive/test_128_16_bigP.py +++ /dev/null @@ -1,383 +0,0 @@ -""" -Minimal PV-only test: Load P from GMEM to TMEM via QK-style MMA, then PV from TMEM. -Step 1: QK MMA writes FP32 S to TMEM (we know this works) -Step 2: Softmax packing writes BF16 P to TMEM (test this) -Step 3: PV MMA reads BF16 P from TMEM and V from SMEM, produces O - -But to isolate the bug, let me test just the PV MMA in isolation. -I'll write known BF16 values to TMEM using the softmax packing path, -then immediately read them back using the PV A-fragment path, -and compare. - -Actually, the simplest isolation test: -1. Do QK MMA to get S in TMEM (cosine 0.999999 verified) -2. Do softmax packing: S → P in TMEM (at offset 32) -3. Skip PV entirely — read P from TMEM using the C-fragment composition LOAD path -4. Output P to GMEM and compare against S.to(BF16) - -This tests whether the softmax packing writes P correctly to the same TMEM -that the PV would read from. - -But we can't easily read P from TMEM using the standard epilogue path -because the epilogue expects FP32 accumulator data. - -Alternative: Use the PV MMA with V=I (identity). If P is correct, -then P @ I = P. But V needs to be MN-major and (128, 128), not (128, 64). -The output would be (128, 128) which doesn't match our (128, 64) c tensor. - -Let me use V that selects the first 64 columns: V[k, n] = delta(k, n) for k in [0,63]. -This gives P @ V = P[:, :64], and the output is (128, 64). -But V is (128, 128) in the MMA K,N dims. V[k, n] for k in [0,127], n in [0,63]. -Hmm, this is getting complicated. Let me just do the identity approach with a (128, 128) output. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class Test128x16Tiler: - """QK + softmax packing + PV with V=I to isolate PV MMA correctness. - Output should be P = S.to(BF16), i.e. (Q@K^T).bfloat16() - With V=I, O = P @ I = P. - But V is (K=128, N=128) in the MMA. We need a 128x128 identity in MN-major. - Output tensor is (128, 128). - """ - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False # needed by epilogue_tma_store - self.epilog_sync_bar_id = 1 # needed by epilogue_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = int(cute.size(qk_mma.shape_mnk, mode=[2])) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - # PV with V=I: output is (128, 128), same as QK - self.pv_mma_tiler = (self.qk_mma_tiler[0], qk_inst_k, self.qk_mma_tiler[1]) - # pv_mma_tiler = (128, 128, 128) since V is 128x128 - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - (self.pv_mma_tiler[0], self.pv_mma_tiler[1], self.pv_mma_tiler[2]), False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 128 - self.tmem_o0_offset = 256 # After S (128 cols) and P (128 cols) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = 512 # Ensure enough TMEM for S+P+O - - # ⛔⛔⛔ CRITICAL: num_tma_load_bytes MUST include ALL TMA-loaded tensors (Q + K + V). Missing V → DEADLOCK. See FOOTGUN #0 in README. - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - # PV with 128x128 output (V=I) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, (128, 16), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_128_16_debug.py b/tests/archive/test_128_16_debug.py deleted file mode 100644 index d1738c1e..00000000 --- a/tests/archive/test_128_16_debug.py +++ /dev/null @@ -1,384 +0,0 @@ -""" -Minimal PV-only test: Load P from GMEM to TMEM via QK-style MMA, then PV from TMEM. -Step 1: QK MMA writes FP32 S to TMEM (we know this works) -Step 2: Softmax packing writes BF16 P to TMEM (test this) -Step 3: PV MMA reads BF16 P from TMEM and V from SMEM, produces O - -But to isolate the bug, let me test just the PV MMA in isolation. -I'll write known BF16 values to TMEM using the softmax packing path, -then immediately read them back using the PV A-fragment path, -and compare. - -Actually, the simplest isolation test: -1. Do QK MMA to get S in TMEM (cosine 0.999999 verified) -2. Do softmax packing: S → P in TMEM (at offset 32) -3. Skip PV entirely — read P from TMEM using the C-fragment composition LOAD path -4. Output P to GMEM and compare against S.to(BF16) - -This tests whether the softmax packing writes P correctly to the same TMEM -that the PV would read from. - -But we can't easily read P from TMEM using the standard epilogue path -because the epilogue expects FP32 accumulator data. - -Alternative: Use the PV MMA with V=I (identity). If P is correct, -then P @ I = P. But V needs to be MN-major and (128, 128), not (128, 64). -The output would be (128, 128) which doesn't match our (128, 64) c tensor. - -Let me use V that selects the first 64 columns: V[k, n] = delta(k, n) for k in [0,63]. -This gives P @ V = P[:, :64], and the output is (128, 64). -But V is (128, 128) in the MMA K,N dims. V[k, n] for k in [0,127], n in [0,63]. -Hmm, this is getting complicated. Let me just do the identity approach with a (128, 128) output. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class Test128x16Tiler: - """QK + softmax packing + PV with V=I to isolate PV MMA correctness. - Output should be P = S.to(BF16), i.e. (Q@K^T).bfloat16() - With V=I, O = P @ I = P. - But V is (K=128, N=128) in the MMA. We need a 128x128 identity in MN-major. - Output tensor is (128, 128). - """ - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False # needed by epilogue_tma_store - self.epilog_sync_bar_id = 1 # needed by epilogue_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = int(cute.size(qk_mma.shape_mnk, mode=[2])) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - # PV with V=I: output is (128, 128), same as QK - self.pv_mma_tiler = (self.qk_mma_tiler[0], qk_inst_k, self.qk_mma_tiler[1]) - # pv_mma_tiler = (128, 128, 128) since V is 128x128 - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - (self.pv_mma_tiler[0], self.pv_mma_tiler[1], self.pv_mma_tiler[2]), False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - print(f"tilePlikeFP32={self.tilePlikeFP32}, pv_mma_tiler={self.pv_mma_tiler}, qk_mma_tiler={self.qk_mma_tiler}") - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - # ⛔⛔⛔ CRITICAL: num_tma_load_bytes MUST include ALL TMA-loaded tensors (Q + K + V). Missing V → DEADLOCK. See FOOTGUN #0 in README. - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - # PV with 128x128 output (V=I) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, (128, 16), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_128_16_debug2.py b/tests/archive/test_128_16_debug2.py deleted file mode 100644 index 21d7c40a..00000000 --- a/tests/archive/test_128_16_debug2.py +++ /dev/null @@ -1,389 +0,0 @@ -""" -Minimal PV-only test: Load P from GMEM to TMEM via QK-style MMA, then PV from TMEM. -Step 1: QK MMA writes FP32 S to TMEM (we know this works) -Step 2: Softmax packing writes BF16 P to TMEM (test this) -Step 3: PV MMA reads BF16 P from TMEM and V from SMEM, produces O - -But to isolate the bug, let me test just the PV MMA in isolation. -I'll write known BF16 values to TMEM using the softmax packing path, -then immediately read them back using the PV A-fragment path, -and compare. - -Actually, the simplest isolation test: -1. Do QK MMA to get S in TMEM (cosine 0.999999 verified) -2. Do softmax packing: S → P in TMEM (at offset 32) -3. Skip PV entirely — read P from TMEM using the C-fragment composition LOAD path -4. Output P to GMEM and compare against S.to(BF16) - -This tests whether the softmax packing writes P correctly to the same TMEM -that the PV would read from. - -But we can't easily read P from TMEM using the standard epilogue path -because the epilogue expects FP32 accumulator data. - -Alternative: Use the PV MMA with V=I (identity). If P is correct, -then P @ I = P. But V needs to be MN-major and (128, 128), not (128, 64). -The output would be (128, 128) which doesn't match our (128, 64) c tensor. - -Let me use V that selects the first 64 columns: V[k, n] = delta(k, n) for k in [0,63]. -This gives P @ V = P[:, :64], and the output is (128, 64). -But V is (128, 128) in the MMA K,N dims. V[k, n] for k in [0,127], n in [0,63]. -Hmm, this is getting complicated. Let me just do the identity approach with a (128, 128) output. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class Test128x16Tiler: - """QK + softmax packing + PV with V=I to isolate PV MMA correctness. - Output should be P = S.to(BF16), i.e. (Q@K^T).bfloat16() - With V=I, O = P @ I = P. - But V is (K=128, N=128) in the MMA. We need a 128x128 identity in MN-major. - Output tensor is (128, 128). - """ - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False # needed by epilogue_tma_store - self.epilog_sync_bar_id = 1 # needed by epilogue_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = int(cute.size(qk_mma.shape_mnk, mode=[2])) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - # PV with V=I: output is (128, 128), same as QK - self.pv_mma_tiler = (self.qk_mma_tiler[0], qk_inst_k, self.qk_mma_tiler[1]) - # pv_mma_tiler = (128, 128, 128) since V is 128x128 - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - (self.pv_mma_tiler[0], self.pv_mma_tiler[1], self.pv_mma_tiler[2]), False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - print(f"p_tmem_s.outer shape={cute.size(p_tmem_s.outer)}") - pv_thr2 = pv_mma.get_slice(0) - pv_acc2 = pv_thr2.partition_shape_C(self.pv_mma_tiler[:2]) - tP2 = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP2 = pv_thr2.make_fragment_A(tP2)[None, None, None, 0] - print(f"tOrP layout shape={cute.size(tOrP2.layout)}, tStS shape={cute.size(tStS.layout)}") - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - # ⛔⛔⛔ CRITICAL: num_tma_load_bytes MUST include ALL TMA-loaded tensors (Q + K + V). Missing V → DEADLOCK. See FOOTGUN #0 in README. - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - # PV with 128x128 output (V=I) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, (128, 16), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_128_16_debug3.py b/tests/archive/test_128_16_debug3.py deleted file mode 100644 index 68ee9f5d..00000000 --- a/tests/archive/test_128_16_debug3.py +++ /dev/null @@ -1,384 +0,0 @@ -""" -Minimal PV-only test: Load P from GMEM to TMEM via QK-style MMA, then PV from TMEM. -Step 1: QK MMA writes FP32 S to TMEM (we know this works) -Step 2: Softmax packing writes BF16 P to TMEM (test this) -Step 3: PV MMA reads BF16 P from TMEM and V from SMEM, produces O - -But to isolate the bug, let me test just the PV MMA in isolation. -I'll write known BF16 values to TMEM using the softmax packing path, -then immediately read them back using the PV A-fragment path, -and compare. - -Actually, the simplest isolation test: -1. Do QK MMA to get S in TMEM (cosine 0.999999 verified) -2. Do softmax packing: S → P in TMEM (at offset 32) -3. Skip PV entirely — read P from TMEM using the C-fragment composition LOAD path -4. Output P to GMEM and compare against S.to(BF16) - -This tests whether the softmax packing writes P correctly to the same TMEM -that the PV would read from. - -But we can't easily read P from TMEM using the standard epilogue path -because the epilogue expects FP32 accumulator data. - -Alternative: Use the PV MMA with V=I (identity). If P is correct, -then P @ I = P. But V needs to be MN-major and (128, 128), not (128, 64). -The output would be (128, 128) which doesn't match our (128, 64) c tensor. - -Let me use V that selects the first 64 columns: V[k, n] = delta(k, n) for k in [0,63]. -This gives P @ V = P[:, :64], and the output is (128, 64). -But V is (128, 128) in the MMA K,N dims. V[k, n] for k in [0,127], n in [0,63]. -Hmm, this is getting complicated. Let me just do the identity approach with a (128, 128) output. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class Test128x16Tiler: - """QK + softmax packing + PV with V=I to isolate PV MMA correctness. - Output should be P = S.to(BF16), i.e. (Q@K^T).bfloat16() - With V=I, O = P @ I = P. - But V is (K=128, N=128) in the MMA. We need a 128x128 identity in MN-major. - Output tensor is (128, 128). - """ - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False # needed by epilogue_tma_store - self.epilog_sync_bar_id = 1 # needed by epilogue_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = int(cute.size(qk_mma.shape_mnk, mode=[2])) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - # PV with V=I: output is (128, 128), same as QK - self.pv_mma_tiler = (self.qk_mma_tiler[0], qk_inst_k, self.qk_mma_tiler[1]) - # pv_mma_tiler = (128, 128, 128) since V is 128x128 - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - (self.pv_mma_tiler[0], self.pv_mma_tiler[1], self.pv_mma_tiler[2]), False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - print(f"tmem offsets: s={self.tmem_s0_offset}, p={self.tmem_p0_offset}, o={self.tmem_o0_offset}, s_cols={s_cols}, o_cols={o_cols}") - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - # ⛔⛔⛔ CRITICAL: num_tma_load_bytes MUST include ALL TMA-loaded tensors (Q + K + V). Missing V → DEADLOCK. See FOOTGUN #0 in README. - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - # PV with 128x128 output (V=I) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, (128, 16), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_128_16_fp16.py b/tests/archive/test_128_16_fp16.py deleted file mode 100644 index 65eb4c95..00000000 --- a/tests/archive/test_128_16_fp16.py +++ /dev/null @@ -1,383 +0,0 @@ -""" -Minimal PV-only test: Load P from GMEM to TMEM via QK-style MMA, then PV from TMEM. -Step 1: QK MMA writes FP32 S to TMEM (we know this works) -Step 2: Softmax packing writes BF16 P to TMEM (test this) -Step 3: PV MMA reads BF16 P from TMEM and V from SMEM, produces O - -But to isolate the bug, let me test just the PV MMA in isolation. -I'll write known BF16 values to TMEM using the softmax packing path, -then immediately read them back using the PV A-fragment path, -and compare. - -Actually, the simplest isolation test: -1. Do QK MMA to get S in TMEM (cosine 0.999999 verified) -2. Do softmax packing: S → P in TMEM (at offset 32) -3. Skip PV entirely — read P from TMEM using the C-fragment composition LOAD path -4. Output P to GMEM and compare against S.to(BF16) - -This tests whether the softmax packing writes P correctly to the same TMEM -that the PV would read from. - -But we can't easily read P from TMEM using the standard epilogue path -because the epilogue expects FP32 accumulator data. - -Alternative: Use the PV MMA with V=I (identity). If P is correct, -then P @ I = P. But V needs to be MN-major and (128, 128), not (128, 64). -The output would be (128, 128) which doesn't match our (128, 64) c tensor. - -Let me use V that selects the first 64 columns: V[k, n] = delta(k, n) for k in [0,63]. -This gives P @ V = P[:, :64], and the output is (128, 64). -But V is (128, 128) in the MMA K,N dims. V[k, n] for k in [0,127], n in [0,63]. -Hmm, this is getting complicated. Let me just do the identity approach with a (128, 128) output. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, Float16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class FP16PVKernel: - """QK + softmax packing + PV with V=I to isolate PV MMA correctness. - Output should be P = S.to(BF16), i.e. (Q@K^T).half() - With V=I, O = P @ I = P. - But V is (K=128, N=128) in the MMA. We need a 128x128 identity in MN-major. - Output tensor is (128, 128). - """ - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = Float16; self.o_dtype = Float16; self.c_dtype = Float16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False # needed by epilogue_tma_store - self.epilog_sync_bar_id = 1 # needed by epilogue_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = int(cute.size(qk_mma.shape_mnk, mode=[2])) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - # PV with V=I: output is (128, 128), same as QK - self.pv_mma_tiler = (self.qk_mma_tiler[0], qk_inst_k, self.qk_mma_tiler[1]) - # pv_mma_tiler = (128, 128, 128) since V is 128x128 - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - (self.pv_mma_tiler[0], self.pv_mma_tiler[1], self.pv_mma_tiler[2]), False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - # ⛔⛔⛔ CRITICAL: num_tma_load_bytes MUST include ALL TMA-loaded tensors (Q + K + V). Missing V → DEADLOCK. See FOOTGUN #0 in README. - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - # PV with 128x128 output (V=I) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, (128, 16), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_128_16_full.py b/tests/archive/test_128_16_full.py deleted file mode 100644 index e343bbae..00000000 --- a/tests/archive/test_128_16_full.py +++ /dev/null @@ -1,385 +0,0 @@ -""" -Minimal PV-only test: Load P from GMEM to TMEM via QK-style MMA, then PV from TMEM. -Step 1: QK MMA writes FP32 S to TMEM (we know this works) -Step 2: Softmax packing writes BF16 P to TMEM (test this) -Step 3: PV MMA reads BF16 P from TMEM and V from SMEM, produces O - -But to isolate the bug, let me test just the PV MMA in isolation. -I'll write known BF16 values to TMEM using the softmax packing path, -then immediately read them back using the PV A-fragment path, -and compare. - -Actually, the simplest isolation test: -1. Do QK MMA to get S in TMEM (cosine 0.999999 verified) -2. Do softmax packing: S → P in TMEM (at offset 32) -3. Skip PV entirely — read P from TMEM using the C-fragment composition LOAD path -4. Output P to GMEM and compare against S.to(BF16) - -This tests whether the softmax packing writes P correctly to the same TMEM -that the PV would read from. - -But we can't easily read P from TMEM using the standard epilogue path -because the epilogue expects FP32 accumulator data. - -Alternative: Use the PV MMA with V=I (identity). If P is correct, -then P @ I = P. But V needs to be MN-major and (128, 128), not (128, 64). -The output would be (128, 128) which doesn't match our (128, 64) c tensor. - -Let me use V that selects the first 64 columns: V[k, n] = delta(k, n) for k in [0,63]. -This gives P @ V = P[:, :64], and the output is (128, 64). -But V is (128, 128) in the MMA K,N dims. V[k, n] for k in [0,127], n in [0,63]. -Hmm, this is getting complicated. Let me just do the identity approach with a (128, 128) output. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class Test128x16Tiler: - """QK + softmax packing + PV with V=I to isolate PV MMA correctness. - Output should be P = S.to(BF16), i.e. (Q@K^T).bfloat16() - With V=I, O = P @ I = P. - But V is (K=128, N=128) in the MMA. We need a 128x128 identity in MN-major. - Output tensor is (128, 128). - """ - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False # needed by epilogue_tma_store - self.epilog_sync_bar_id = 1 # needed by epilogue_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = int(cute.size(qk_mma.shape_mnk, mode=[2])) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - # PV with V=I: output is (128, 128), same as QK - self.pv_mma_tiler = (self.qk_mma_tiler[0], qk_inst_k, self.qk_mma_tiler[1]) - # pv_mma_tiler = (128, 128, 128) since V is 128x128 - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - (self.pv_mma_tiler[0], self.pv_mma_tiler[1], self.pv_mma_tiler[2]), False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - # ⛔⛔⛔ CRITICAL: num_tma_load_bytes MUST include ALL TMA-loaded tensors (Q + K + V). Missing V → DEADLOCK. See FOOTGUN #0 in README. - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - # PV with 128x128 output (V=I) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, (128, 16), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_((self.pv_mma_tiler[0], self.pv_mma_tiler[1], self.pv_mma_tiler[2]), (None,0,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_128_16_minimal.py b/tests/archive/test_128_16_minimal.py deleted file mode 100644 index 14bf24f2..00000000 --- a/tests/archive/test_128_16_minimal.py +++ /dev/null @@ -1,385 +0,0 @@ -""" -Minimal PV-only test: Load P from GMEM to TMEM via QK-style MMA, then PV from TMEM. -Step 1: QK MMA writes FP32 S to TMEM (we know this works) -Step 2: Softmax packing writes BF16 P to TMEM (test this) -Step 3: PV MMA reads BF16 P from TMEM and V from SMEM, produces O - -But to isolate the bug, let me test just the PV MMA in isolation. -I'll write known BF16 values to TMEM using the softmax packing path, -then immediately read them back using the PV A-fragment path, -and compare. - -Actually, the simplest isolation test: -1. Do QK MMA to get S in TMEM (cosine 0.999999 verified) -2. Do softmax packing: S → P in TMEM (at offset 32) -3. Skip PV entirely — read P from TMEM using the C-fragment composition LOAD path -4. Output P to GMEM and compare against S.to(BF16) - -This tests whether the softmax packing writes P correctly to the same TMEM -that the PV would read from. - -But we can't easily read P from TMEM using the standard epilogue path -because the epilogue expects FP32 accumulator data. - -Alternative: Use the PV MMA with V=I (identity). If P is correct, -then P @ I = P. But V needs to be MN-major and (128, 128), not (128, 64). -The output would be (128, 128) which doesn't match our (128, 64) c tensor. - -Let me use V that selects the first 64 columns: V[k, n] = delta(k, n) for k in [0,63]. -This gives P @ V = P[:, :64], and the output is (128, 64). -But V is (128, 128) in the MMA K,N dims. V[k, n] for k in [0,127], n in [0,63]. -Hmm, this is getting complicated. Let me just do the identity approach with a (128, 128) output. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class MinimalV8Kernel: - """QK + softmax packing + PV with V=I to isolate PV MMA correctness. - Output should be P = S.to(BF16), i.e. (Q@K^T).bfloat16() - With V=I, O = P @ I = P. - But V is (K=128, N=128) in the MMA. We need a 128x128 identity in MN-major. - Output tensor is (128, 128). - """ - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False # needed by epilogue_tma_store - self.epilog_sync_bar_id = 1 # needed by epilogue_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = int(cute.size(qk_mma.shape_mnk, mode=[2])) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - # PV with V=I: output is (128, 128), same as QK - self.pv_mma_tiler = (self.qk_mma_tiler[0], qk_inst_k, self.qk_mma_tiler[1]) - # pv_mma_tiler = (128, 128, 128) since V is 128x128 - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - (self.pv_mma_tiler[0], self.pv_mma_tiler[1], self.pv_mma_tiler[2]), False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - # ⛔⛔⛔ CRITICAL: num_tma_load_bytes MUST include ALL TMA-loaded tensors (Q + K + V). Missing V → DEADLOCK. See FOOTGUN #0 in README. - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - # PV with 128x128 output (V=I) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, (128, 16), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.pv_mma_tiler, (None,0,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1 O = P[:,:16] = (Q@K^T).bf16()[:,:16] - ref = (qf @ kf.T).bfloat16().float()[:, :16] - - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - kernel = MinimalV8Kernel(mma_tiler_mn=(128, 128)) - print('Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - print('Running...', flush=True) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - out = c[:,:,0].float() - cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - print('PV(128,16) minimal: cosine {:.6f} {}'.format(cos, 'PASS' if cos >= 0.99 else 'FAIL')) - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_128_16_nogC.py b/tests/archive/test_128_16_nogC.py deleted file mode 100644 index e4527eb2..00000000 --- a/tests/archive/test_128_16_nogC.py +++ /dev/null @@ -1,383 +0,0 @@ -""" -Minimal PV-only test: Load P from GMEM to TMEM via QK-style MMA, then PV from TMEM. -Step 1: QK MMA writes FP32 S to TMEM (we know this works) -Step 2: Softmax packing writes BF16 P to TMEM (test this) -Step 3: PV MMA reads BF16 P from TMEM and V from SMEM, produces O - -But to isolate the bug, let me test just the PV MMA in isolation. -I'll write known BF16 values to TMEM using the softmax packing path, -then immediately read them back using the PV A-fragment path, -and compare. - -Actually, the simplest isolation test: -1. Do QK MMA to get S in TMEM (cosine 0.999999 verified) -2. Do softmax packing: S → P in TMEM (at offset 32) -3. Skip PV entirely — read P from TMEM using the C-fragment composition LOAD path -4. Output P to GMEM and compare against S.to(BF16) - -This tests whether the softmax packing writes P correctly to the same TMEM -that the PV would read from. - -But we can't easily read P from TMEM using the standard epilogue path -because the epilogue expects FP32 accumulator data. - -Alternative: Use the PV MMA with V=I (identity). If P is correct, -then P @ I = P. But V needs to be MN-major and (128, 128), not (128, 64). -The output would be (128, 128) which doesn't match our (128, 64) c tensor. - -Let me use V that selects the first 64 columns: V[k, n] = delta(k, n) for k in [0,63]. -This gives P @ V = P[:, :64], and the output is (128, 64). -But V is (128, 128) in the MMA K,N dims. V[k, n] for k in [0,127], n in [0,63]. -Hmm, this is getting complicated. Let me just do the identity approach with a (128, 128) output. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class Test128x16Tiler: - """QK + softmax packing + PV with V=I to isolate PV MMA correctness. - Output should be P = S.to(BF16), i.e. (Q@K^T).bfloat16() - With V=I, O = P @ I = P. - But V is (K=128, N=128) in the MMA. We need a 128x128 identity in MN-major. - Output tensor is (128, 128). - """ - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False # needed by epilogue_tma_store - self.epilog_sync_bar_id = 1 # needed by epilogue_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = int(cute.size(qk_mma.shape_mnk, mode=[2])) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - # PV with V=I: output is (128, 128), same as QK - self.pv_mma_tiler = (self.qk_mma_tiler[0], qk_inst_k, self.qk_mma_tiler[1]) - # pv_mma_tiler = (128, 128, 128) since V is 128x128 - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - (self.pv_mma_tiler[0], self.pv_mma_tiler[1], self.pv_mma_tiler[2]), False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - # ⛔⛔⛔ CRITICAL: num_tma_load_bytes MUST include ALL TMA-loaded tensors (Q + K + V). Missing V → DEADLOCK. See FOOTGUN #0 in README. - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - # PV with 128x128 output (V=I) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, (128, 16), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_128_16_nopack.py b/tests/archive/test_128_16_nopack.py deleted file mode 100644 index c3ca85ae..00000000 --- a/tests/archive/test_128_16_nopack.py +++ /dev/null @@ -1,383 +0,0 @@ -""" -Minimal PV-only test: Load P from GMEM to TMEM via QK-style MMA, then PV from TMEM. -Step 1: QK MMA writes FP32 S to TMEM (we know this works) -Step 2: Softmax packing writes BF16 P to TMEM (test this) -Step 3: PV MMA reads BF16 P from TMEM and V from SMEM, produces O - -But to isolate the bug, let me test just the PV MMA in isolation. -I'll write known BF16 values to TMEM using the softmax packing path, -then immediately read them back using the PV A-fragment path, -and compare. - -Actually, the simplest isolation test: -1. Do QK MMA to get S in TMEM (cosine 0.999999 verified) -2. Do softmax packing: S → P in TMEM (at offset 32) -3. Skip PV entirely — read P from TMEM using the C-fragment composition LOAD path -4. Output P to GMEM and compare against S.to(BF16) - -This tests whether the softmax packing writes P correctly to the same TMEM -that the PV would read from. - -But we can't easily read P from TMEM using the standard epilogue path -because the epilogue expects FP32 accumulator data. - -Alternative: Use the PV MMA with V=I (identity). If P is correct, -then P @ I = P. But V needs to be MN-major and (128, 128), not (128, 64). -The output would be (128, 128) which doesn't match our (128, 64) c tensor. - -Let me use V that selects the first 64 columns: V[k, n] = delta(k, n) for k in [0,63]. -This gives P @ V = P[:, :64], and the output is (128, 64). -But V is (128, 128) in the MMA K,N dims. V[k, n] for k in [0,127], n in [0,63]. -Hmm, this is getting complicated. Let me just do the identity approach with a (128, 128) output. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class NoPackV8Kernel: - """QK + softmax packing + PV with V=I to isolate PV MMA correctness. - Output should be P = S.to(BF16), i.e. (Q@K^T).bfloat16() - With V=I, O = P @ I = P. - But V is (K=128, N=128) in the MMA. We need a 128x128 identity in MN-major. - Output tensor is (128, 128). - """ - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False # needed by epilogue_tma_store - self.epilog_sync_bar_id = 1 # needed by epilogue_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - # PV with V=I: output is (128, 128), same as QK - self.pv_mma_tiler = (self.qk_mma_tiler[0], qk_inst_k, self.qk_mma_tiler[1]) - # pv_mma_tiler = (128, 128, 128) since V is 128x128 - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - (self.pv_mma_tiler[0], self.pv_mma_tiler[1], self.pv_mma_tiler[2]), False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - # ⛔⛔⛔ CRITICAL: num_tma_load_bytes MUST include ALL TMA-loaded tensors (Q + K + V). Missing V → DEADLOCK. See FOOTGUN #0 in README. - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - # PV with 128x128 output (V=I) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, (128, 16), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_128_16_nosoftmax.py b/tests/archive/test_128_16_nosoftmax.py deleted file mode 100644 index 3395c962..00000000 --- a/tests/archive/test_128_16_nosoftmax.py +++ /dev/null @@ -1,383 +0,0 @@ -""" -Minimal PV-only test: Load P from GMEM to TMEM via QK-style MMA, then PV from TMEM. -Step 1: QK MMA writes FP32 S to TMEM (we know this works) -Step 2: Softmax packing writes BF16 P to TMEM (test this) -Step 3: PV MMA reads BF16 P from TMEM and V from SMEM, produces O - -But to isolate the bug, let me test just the PV MMA in isolation. -I'll write known BF16 values to TMEM using the softmax packing path, -then immediately read them back using the PV A-fragment path, -and compare. - -Actually, the simplest isolation test: -1. Do QK MMA to get S in TMEM (cosine 0.999999 verified) -2. Do softmax packing: S → P in TMEM (at offset 32) -3. Skip PV entirely — read P from TMEM using the C-fragment composition LOAD path -4. Output P to GMEM and compare against S.to(BF16) - -This tests whether the softmax packing writes P correctly to the same TMEM -that the PV would read from. - -But we can't easily read P from TMEM using the standard epilogue path -because the epilogue expects FP32 accumulator data. - -Alternative: Use the PV MMA with V=I (identity). If P is correct, -then P @ I = P. But V needs to be MN-major and (128, 128), not (128, 64). -The output would be (128, 128) which doesn't match our (128, 64) c tensor. - -Let me use V that selects the first 64 columns: V[k, n] = delta(k, n) for k in [0,63]. -This gives P @ V = P[:, :64], and the output is (128, 64). -But V is (128, 128) in the MMA K,N dims. V[k, n] for k in [0,127], n in [0,63]. -Hmm, this is getting complicated. Let me just do the identity approach with a (128, 128) output. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class Test128x16NoSoftmax: - """QK + softmax packing + PV with V=I to isolate PV MMA correctness. - Output should be P = S.to(BF16), i.e. (Q@K^T).bfloat16() - With V=I, O = P @ I = P. - But V is (K=128, N=128) in the MMA. We need a 128x128 identity in MN-major. - Output tensor is (128, 128). - """ - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False # needed by epilogue_tma_store - self.epilog_sync_bar_id = 1 # needed by epilogue_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - # PV with V=I: output is (128, 128), same as QK - self.pv_mma_tiler = (self.qk_mma_tiler[0], self.qk_mma_tiler[1], self.qk_mma_tiler[1]) - # pv_mma_tiler = (128, 128, 128) since V is 128x128 - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - # ⛔⛔⛔ CRITICAL: num_tma_load_bytes MUST include ALL TMA-loaded tensors (Q + K + V). Missing V → DEADLOCK. See FOOTGUN #0 in README. - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - # PV with 128x128 output (V=I) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, (128, 16), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_128_16_pAtS.py b/tests/archive/test_128_16_pAtS.py deleted file mode 100644 index 57d7a30b..00000000 --- a/tests/archive/test_128_16_pAtS.py +++ /dev/null @@ -1,367 +0,0 @@ -""" -Minimal PV-only test: Load P from GMEM to TMEM via QK-style MMA, then PV from TMEM. -Step 1: QK MMA writes FP32 S to TMEM (we know this works) -Step 2: Softmax packing writes BF16 P to TMEM (test this) -Step 3: PV MMA reads BF16 P from TMEM and V from SMEM, produces O - -But to isolate the bug, let me test just the PV MMA in isolation. -I'll write known BF16 values to TMEM using the softmax packing path, -then immediately read them back using the PV A-fragment path, -and compare. - -Actually, the simplest isolation test: -1. Do QK MMA to get S in TMEM (cosine 0.999999 verified) -2. Do softmax packing: S → P in TMEM (at offset 32) -3. Skip PV entirely — read P from TMEM using the C-fragment composition LOAD path -4. Output P to GMEM and compare against S.to(BF16) - -This tests whether the softmax packing writes P correctly to the same TMEM -that the PV would read from. - -But we can't easily read P from TMEM using the standard epilogue path -because the epilogue expects FP32 accumulator data. - -Alternative: Use the PV MMA with V=I (identity). If P is correct, -then P @ I = P. But V needs to be MN-major and (128, 128), not (128, 64). -The output would be (128, 128) which doesn't match our (128, 64) c tensor. - -Let me use V that selects the first 64 columns: V[k, n] = delta(k, n) for k in [0,63]. -This gives P @ V = P[:, :64], and the output is (128, 64). -But V is (128, 128) in the MMA K,N dims. V[k, n] for k in [0,127], n in [0,63]. -Hmm, this is getting complicated. Let me just do the identity approach with a (128, 128) output. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class Test128x16Tiler: - """QK + softmax packing + PV with V=I to isolate PV MMA correctness. - Output should be P = S.to(BF16), i.e. (Q@K^T).bfloat16() - With V=I, O = P @ I = P. - But V is (K=128, N=128) in the MMA. We need a 128x128 identity in MN-major. - Output tensor is (128, 128). - """ - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False # needed by epilogue_tma_store - self.epilog_sync_bar_id = 1 # needed by epilogue_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = int(cute.size(qk_mma.shape_mnk, mode=[2])) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - # PV with V=I: output is (128, 128), same as QK - self.pv_mma_tiler = (self.qk_mma_tiler[0], qk_inst_k, self.qk_mma_tiler[1]) - # pv_mma_tiler = (128, 128, 128) since V is 128x128 - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - (self.pv_mma_tiler[0], self.pv_mma_tiler[1], self.pv_mma_tiler[2]), False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 0 - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - # ⛔⛔⛔ CRITICAL: num_tma_load_bytes MUST include ALL TMA-loaded tensors (Q + K + V). Missing V → DEADLOCK. See FOOTGUN #0 in README. - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - # PV with 128x128 output (V=I) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, (128, 16), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_128_16_pvlayout.py b/tests/archive/test_128_16_pvlayout.py deleted file mode 100644 index c11584bf..00000000 --- a/tests/archive/test_128_16_pvlayout.py +++ /dev/null @@ -1,387 +0,0 @@ -""" -Minimal PV-only test: Load P from GMEM to TMEM via QK-style MMA, then PV from TMEM. -Step 1: QK MMA writes FP32 S to TMEM (we know this works) -Step 2: Softmax packing writes BF16 P to TMEM (test this) -Step 3: PV MMA reads BF16 P from TMEM and V from SMEM, produces O - -But to isolate the bug, let me test just the PV MMA in isolation. -I'll write known BF16 values to TMEM using the softmax packing path, -then immediately read them back using the PV A-fragment path, -and compare. - -Actually, the simplest isolation test: -1. Do QK MMA to get S in TMEM (cosine 0.999999 verified) -2. Do softmax packing: S → P in TMEM (at offset 32) -3. Skip PV entirely — read P from TMEM using the C-fragment composition LOAD path -4. Output P to GMEM and compare against S.to(BF16) - -This tests whether the softmax packing writes P correctly to the same TMEM -that the PV would read from. - -But we can't easily read P from TMEM using the standard epilogue path -because the epilogue expects FP32 accumulator data. - -Alternative: Use the PV MMA with V=I (identity). If P is correct, -then P @ I = P. But V needs to be MN-major and (128, 128), not (128, 64). -The output would be (128, 128) which doesn't match our (128, 64) c tensor. - -Let me use V that selects the first 64 columns: V[k, n] = delta(k, n) for k in [0,63]. -This gives P @ V = P[:, :64], and the output is (128, 64). -But V is (128, 128) in the MMA K,N dims. V[k, n] for k in [0,127], n in [0,63]. -Hmm, this is getting complicated. Let me just do the identity approach with a (128, 128) output. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class Test128x16Tiler: - """QK + softmax packing + PV with V=I to isolate PV MMA correctness. - Output should be P = S.to(BF16), i.e. (Q@K^T).bfloat16() - With V=I, O = P @ I = P. - But V is (K=128, N=128) in the MMA. We need a 128x128 identity in MN-major. - Output tensor is (128, 128). - """ - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False # needed by epilogue_tma_store - self.epilog_sync_bar_id = 1 # needed by epilogue_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = int(cute.size(qk_mma.shape_mnk, mode=[2])) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - # PV with V=I: output is (128, 128), same as QK - self.pv_mma_tiler = (self.qk_mma_tiler[0], qk_inst_k, self.qk_mma_tiler[1]) - # pv_mma_tiler = (128, 128, 128) since V is 128x128 - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - (self.pv_mma_tiler[0], self.pv_mma_tiler[1], self.pv_mma_tiler[2]), False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - # ⛔⛔⛔ CRITICAL: num_tma_load_bytes MUST include ALL TMA-loaded tensors (Q + K + V). Missing V → DEADLOCK. See FOOTGUN #0 in README. - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - # PV with 128x128 output (V=I) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, (128, 16), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_128_16_pvpack.py b/tests/archive/test_128_16_pvpack.py deleted file mode 100644 index e4527eb2..00000000 --- a/tests/archive/test_128_16_pvpack.py +++ /dev/null @@ -1,383 +0,0 @@ -""" -Minimal PV-only test: Load P from GMEM to TMEM via QK-style MMA, then PV from TMEM. -Step 1: QK MMA writes FP32 S to TMEM (we know this works) -Step 2: Softmax packing writes BF16 P to TMEM (test this) -Step 3: PV MMA reads BF16 P from TMEM and V from SMEM, produces O - -But to isolate the bug, let me test just the PV MMA in isolation. -I'll write known BF16 values to TMEM using the softmax packing path, -then immediately read them back using the PV A-fragment path, -and compare. - -Actually, the simplest isolation test: -1. Do QK MMA to get S in TMEM (cosine 0.999999 verified) -2. Do softmax packing: S → P in TMEM (at offset 32) -3. Skip PV entirely — read P from TMEM using the C-fragment composition LOAD path -4. Output P to GMEM and compare against S.to(BF16) - -This tests whether the softmax packing writes P correctly to the same TMEM -that the PV would read from. - -But we can't easily read P from TMEM using the standard epilogue path -because the epilogue expects FP32 accumulator data. - -Alternative: Use the PV MMA with V=I (identity). If P is correct, -then P @ I = P. But V needs to be MN-major and (128, 128), not (128, 64). -The output would be (128, 128) which doesn't match our (128, 64) c tensor. - -Let me use V that selects the first 64 columns: V[k, n] = delta(k, n) for k in [0,63]. -This gives P @ V = P[:, :64], and the output is (128, 64). -But V is (128, 128) in the MMA K,N dims. V[k, n] for k in [0,127], n in [0,63]. -Hmm, this is getting complicated. Let me just do the identity approach with a (128, 128) output. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class Test128x16Tiler: - """QK + softmax packing + PV with V=I to isolate PV MMA correctness. - Output should be P = S.to(BF16), i.e. (Q@K^T).bfloat16() - With V=I, O = P @ I = P. - But V is (K=128, N=128) in the MMA. We need a 128x128 identity in MN-major. - Output tensor is (128, 128). - """ - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False # needed by epilogue_tma_store - self.epilog_sync_bar_id = 1 # needed by epilogue_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = int(cute.size(qk_mma.shape_mnk, mode=[2])) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - # PV with V=I: output is (128, 128), same as QK - self.pv_mma_tiler = (self.qk_mma_tiler[0], qk_inst_k, self.qk_mma_tiler[1]) - # pv_mma_tiler = (128, 128, 128) since V is 128x128 - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - (self.pv_mma_tiler[0], self.pv_mma_tiler[1], self.pv_mma_tiler[2]), False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - # ⛔⛔⛔ CRITICAL: num_tma_load_bytes MUST include ALL TMA-loaded tensors (Q + K + V). Missing V → DEADLOCK. See FOOTGUN #0 in README. - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - # PV with 128x128 output (V=I) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, (128, 16), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_128_16_pvwrite.py b/tests/archive/test_128_16_pvwrite.py deleted file mode 100644 index e4527eb2..00000000 --- a/tests/archive/test_128_16_pvwrite.py +++ /dev/null @@ -1,383 +0,0 @@ -""" -Minimal PV-only test: Load P from GMEM to TMEM via QK-style MMA, then PV from TMEM. -Step 1: QK MMA writes FP32 S to TMEM (we know this works) -Step 2: Softmax packing writes BF16 P to TMEM (test this) -Step 3: PV MMA reads BF16 P from TMEM and V from SMEM, produces O - -But to isolate the bug, let me test just the PV MMA in isolation. -I'll write known BF16 values to TMEM using the softmax packing path, -then immediately read them back using the PV A-fragment path, -and compare. - -Actually, the simplest isolation test: -1. Do QK MMA to get S in TMEM (cosine 0.999999 verified) -2. Do softmax packing: S → P in TMEM (at offset 32) -3. Skip PV entirely — read P from TMEM using the C-fragment composition LOAD path -4. Output P to GMEM and compare against S.to(BF16) - -This tests whether the softmax packing writes P correctly to the same TMEM -that the PV would read from. - -But we can't easily read P from TMEM using the standard epilogue path -because the epilogue expects FP32 accumulator data. - -Alternative: Use the PV MMA with V=I (identity). If P is correct, -then P @ I = P. But V needs to be MN-major and (128, 128), not (128, 64). -The output would be (128, 128) which doesn't match our (128, 64) c tensor. - -Let me use V that selects the first 64 columns: V[k, n] = delta(k, n) for k in [0,63]. -This gives P @ V = P[:, :64], and the output is (128, 64). -But V is (128, 128) in the MMA K,N dims. V[k, n] for k in [0,127], n in [0,63]. -Hmm, this is getting complicated. Let me just do the identity approach with a (128, 128) output. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class Test128x16Tiler: - """QK + softmax packing + PV with V=I to isolate PV MMA correctness. - Output should be P = S.to(BF16), i.e. (Q@K^T).bfloat16() - With V=I, O = P @ I = P. - But V is (K=128, N=128) in the MMA. We need a 128x128 identity in MN-major. - Output tensor is (128, 128). - """ - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False # needed by epilogue_tma_store - self.epilog_sync_bar_id = 1 # needed by epilogue_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = int(cute.size(qk_mma.shape_mnk, mode=[2])) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - # PV with V=I: output is (128, 128), same as QK - self.pv_mma_tiler = (self.qk_mma_tiler[0], qk_inst_k, self.qk_mma_tiler[1]) - # pv_mma_tiler = (128, 128, 128) since V is 128x128 - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - (self.pv_mma_tiler[0], self.pv_mma_tiler[1], self.pv_mma_tiler[2]), False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - # ⛔⛔⛔ CRITICAL: num_tma_load_bytes MUST include ALL TMA-loaded tensors (Q + K + V). Missing V → DEADLOCK. See FOOTGUN #0 in README. - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - # PV with 128x128 output (V=I) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, (128, 16), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_128_16_qkread.py b/tests/archive/test_128_16_qkread.py deleted file mode 100644 index e526c3e1..00000000 --- a/tests/archive/test_128_16_qkread.py +++ /dev/null @@ -1,383 +0,0 @@ -""" -Minimal PV-only test: Load P from GMEM to TMEM via QK-style MMA, then PV from TMEM. -Step 1: QK MMA writes FP32 S to TMEM (we know this works) -Step 2: Softmax packing writes BF16 P to TMEM (test this) -Step 3: PV MMA reads BF16 P from TMEM and V from SMEM, produces O - -But to isolate the bug, let me test just the PV MMA in isolation. -I'll write known BF16 values to TMEM using the softmax packing path, -then immediately read them back using the PV A-fragment path, -and compare. - -Actually, the simplest isolation test: -1. Do QK MMA to get S in TMEM (cosine 0.999999 verified) -2. Do softmax packing: S → P in TMEM (at offset 32) -3. Skip PV entirely — read P from TMEM using the C-fragment composition LOAD path -4. Output P to GMEM and compare against S.to(BF16) - -This tests whether the softmax packing writes P correctly to the same TMEM -that the PV would read from. - -But we can't easily read P from TMEM using the standard epilogue path -because the epilogue expects FP32 accumulator data. - -Alternative: Use the PV MMA with V=I (identity). If P is correct, -then P @ I = P. But V needs to be MN-major and (128, 128), not (128, 64). -The output would be (128, 128) which doesn't match our (128, 64) c tensor. - -Let me use V that selects the first 64 columns: V[k, n] = delta(k, n) for k in [0,63]. -This gives P @ V = P[:, :64], and the output is (128, 64). -But V is (128, 128) in the MMA K,N dims. V[k, n] for k in [0,127], n in [0,63]. -Hmm, this is getting complicated. Let me just do the identity approach with a (128, 128) output. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class Test128x16Tiler: - """QK + softmax packing + PV with V=I to isolate PV MMA correctness. - Output should be P = S.to(BF16), i.e. (Q@K^T).bfloat16() - With V=I, O = P @ I = P. - But V is (K=128, N=128) in the MMA. We need a 128x128 identity in MN-major. - Output tensor is (128, 128). - """ - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False # needed by epilogue_tma_store - self.epilog_sync_bar_id = 1 # needed by epilogue_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = int(cute.size(qk_mma.shape_mnk, mode=[2])) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - # PV with V=I: output is (128, 128), same as QK - self.pv_mma_tiler = (self.qk_mma_tiler[0], qk_inst_k, self.qk_mma_tiler[1]) - # pv_mma_tiler = (128, 128, 128) since V is 128x128 - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - (self.pv_mma_tiler[0], self.pv_mma_tiler[1], self.pv_mma_tiler[2]), False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - # ⛔⛔⛔ CRITICAL: num_tma_load_bytes MUST include ALL TMA-loaded tensors (Q + K + V). Missing V → DEADLOCK. See FOOTGUN #0 in README. - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - # PV with 128x128 output (V=I) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, (128, 16), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, tStS.layout) # Use QK layout so PV reads match softmax writes - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_128_16_smem.py b/tests/archive/test_128_16_smem.py deleted file mode 100644 index 7dd0be5c..00000000 --- a/tests/archive/test_128_16_smem.py +++ /dev/null @@ -1,382 +0,0 @@ -""" -Minimal PV-only test: Load P from GMEM to TMEM via QK-style MMA, then PV from TMEM. -Step 1: QK MMA writes FP32 S to TMEM (we know this works) -Step 2: Softmax packing writes BF16 P to TMEM (test this) -Step 3: PV MMA reads BF16 P from TMEM and V from SMEM, produces O - -But to isolate the bug, let me test just the PV MMA in isolation. -I'll write known BF16 values to TMEM using the softmax packing path, -then immediately read them back using the PV A-fragment path, -and compare. - -Actually, the simplest isolation test: -1. Do QK MMA to get S in TMEM (cosine 0.999999 verified) -2. Do softmax packing: S → P in TMEM (at offset 32) -3. Skip PV entirely — read P from TMEM using the C-fragment composition LOAD path -4. Output P to GMEM and compare against S.to(BF16) - -This tests whether the softmax packing writes P correctly to the same TMEM -that the PV would read from. - -But we can't easily read P from TMEM using the standard epilogue path -because the epilogue expects FP32 accumulator data. - -Alternative: Use the PV MMA with V=I (identity). If P is correct, -then P @ I = P. But V needs to be MN-major and (128, 128), not (128, 64). -The output would be (128, 128) which doesn't match our (128, 64) c tensor. - -Let me use V that selects the first 64 columns: V[k, n] = delta(k, n) for k in [0,63]. -This gives P @ V = P[:, :64], and the output is (128, 64). -But V is (128, 128) in the MMA K,N dims. V[k, n] for k in [0,127], n in [0,63]. -Hmm, this is getting complicated. Let me just do the identity approach with a (128, 128) output. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class SMEMPVKernel: - """QK + softmax packing + PV with V=I to isolate PV MMA correctness. - Output should be P = S.to(BF16), i.e. (Q@K^T).bfloat16() - With V=I, O = P @ I = P. - But V is (K=128, N=128) in the MMA. We need a 128x128 identity in MN-major. - Output tensor is (128, 128). - """ - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False # needed by epilogue_tma_store - self.epilog_sync_bar_id = 1 # needed by epilogue_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = int(cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - # PV with V=I: output is (128, 128), same as QK - self.pv_mma_tiler = (self.qk_mma_tiler[0], qk_inst_k, self.qk_mma_tiler[1]) - # pv_mma_tiler = (128, 128, 128) since V is 128x128 - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - (self.pv_mma_tiler[0], self.pv_mma_tiler[1], self.pv_mma_tiler[2]), False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - # ⛔⛔⛔ CRITICAL: num_tma_load_bytes MUST include ALL TMA-loaded tensors (Q + K + V). Missing V → DEADLOCK. See FOOTGUN #0 in README. - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - # PV with 128x128 output (V=I) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, (128, 16), tcgen05.OperandSource.SMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_128_16_stepA.py b/tests/archive/test_128_16_stepA.py deleted file mode 100644 index a44ebf71..00000000 --- a/tests/archive/test_128_16_stepA.py +++ /dev/null @@ -1,383 +0,0 @@ -""" -Minimal PV-only test: Load P from GMEM to TMEM via QK-style MMA, then PV from TMEM. -Step 1: QK MMA writes FP32 S to TMEM (we know this works) -Step 2: Softmax packing writes BF16 P to TMEM (test this) -Step 3: PV MMA reads BF16 P from TMEM and V from SMEM, produces O - -But to isolate the bug, let me test just the PV MMA in isolation. -I'll write known BF16 values to TMEM using the softmax packing path, -then immediately read them back using the PV A-fragment path, -and compare. - -Actually, the simplest isolation test: -1. Do QK MMA to get S in TMEM (cosine 0.999999 verified) -2. Do softmax packing: S → P in TMEM (at offset 32) -3. Skip PV entirely — read P from TMEM using the C-fragment composition LOAD path -4. Output P to GMEM and compare against S.to(BF16) - -This tests whether the softmax packing writes P correctly to the same TMEM -that the PV would read from. - -But we can't easily read P from TMEM using the standard epilogue path -because the epilogue expects FP32 accumulator data. - -Alternative: Use the PV MMA with V=I (identity). If P is correct, -then P @ I = P. But V needs to be MN-major and (128, 128), not (128, 64). -The output would be (128, 128) which doesn't match our (128, 64) c tensor. - -Let me use V that selects the first 64 columns: V[k, n] = delta(k, n) for k in [0,63]. -This gives P @ V = P[:, :64], and the output is (128, 64). -But V is (128, 128) in the MMA K,N dims. V[k, n] for k in [0,127], n in [0,63]. -Hmm, this is getting complicated. Let me just do the identity approach with a (128, 128) output. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class Test128x16Tiler: - """QK + softmax packing + PV with V=I to isolate PV MMA correctness. - Output should be P = S.to(BF16), i.e. (Q@K^T).bfloat16() - With V=I, O = P @ I = P. - But V is (K=128, N=128) in the MMA. We need a 128x128 identity in MN-major. - Output tensor is (128, 128). - """ - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False # needed by epilogue_tma_store - self.epilog_sync_bar_id = 1 # needed by epilogue_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = int(cute.size(qk_mma.shape_mnk, mode=[2])) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - # PV with V=I: output is (128, 128), same as QK - self.pv_mma_tiler = (self.qk_mma_tiler[0], qk_inst_k, self.qk_mma_tiler[1]) - # pv_mma_tiler = (128, 128, 128) since V is 128x128 - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - # ⛔⛔⛔ CRITICAL: num_tma_load_bytes MUST include ALL TMA-loaded tensors (Q + K + V). Missing V → DEADLOCK. See FOOTGUN #0 in README. - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - # PV with 128x128 output (V=I) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, (128, 16), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_128_16_stepB.py b/tests/archive/test_128_16_stepB.py deleted file mode 100644 index e4527eb2..00000000 --- a/tests/archive/test_128_16_stepB.py +++ /dev/null @@ -1,383 +0,0 @@ -""" -Minimal PV-only test: Load P from GMEM to TMEM via QK-style MMA, then PV from TMEM. -Step 1: QK MMA writes FP32 S to TMEM (we know this works) -Step 2: Softmax packing writes BF16 P to TMEM (test this) -Step 3: PV MMA reads BF16 P from TMEM and V from SMEM, produces O - -But to isolate the bug, let me test just the PV MMA in isolation. -I'll write known BF16 values to TMEM using the softmax packing path, -then immediately read them back using the PV A-fragment path, -and compare. - -Actually, the simplest isolation test: -1. Do QK MMA to get S in TMEM (cosine 0.999999 verified) -2. Do softmax packing: S → P in TMEM (at offset 32) -3. Skip PV entirely — read P from TMEM using the C-fragment composition LOAD path -4. Output P to GMEM and compare against S.to(BF16) - -This tests whether the softmax packing writes P correctly to the same TMEM -that the PV would read from. - -But we can't easily read P from TMEM using the standard epilogue path -because the epilogue expects FP32 accumulator data. - -Alternative: Use the PV MMA with V=I (identity). If P is correct, -then P @ I = P. But V needs to be MN-major and (128, 128), not (128, 64). -The output would be (128, 128) which doesn't match our (128, 64) c tensor. - -Let me use V that selects the first 64 columns: V[k, n] = delta(k, n) for k in [0,63]. -This gives P @ V = P[:, :64], and the output is (128, 64). -But V is (128, 128) in the MMA K,N dims. V[k, n] for k in [0,127], n in [0,63]. -Hmm, this is getting complicated. Let me just do the identity approach with a (128, 128) output. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class Test128x16Tiler: - """QK + softmax packing + PV with V=I to isolate PV MMA correctness. - Output should be P = S.to(BF16), i.e. (Q@K^T).bfloat16() - With V=I, O = P @ I = P. - But V is (K=128, N=128) in the MMA. We need a 128x128 identity in MN-major. - Output tensor is (128, 128). - """ - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False # needed by epilogue_tma_store - self.epilog_sync_bar_id = 1 # needed by epilogue_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = int(cute.size(qk_mma.shape_mnk, mode=[2])) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - # PV with V=I: output is (128, 128), same as QK - self.pv_mma_tiler = (self.qk_mma_tiler[0], qk_inst_k, self.qk_mma_tiler[1]) - # pv_mma_tiler = (128, 128, 128) since V is 128x128 - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - (self.pv_mma_tiler[0], self.pv_mma_tiler[1], self.pv_mma_tiler[2]), False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - # ⛔⛔⛔ CRITICAL: num_tma_load_bytes MUST include ALL TMA-loaded tensors (Q + K + V). Missing V → DEADLOCK. See FOOTGUN #0 in README. - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - # PV with 128x128 output (V=I) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, (128, 16), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_128_16_stepC.py b/tests/archive/test_128_16_stepC.py deleted file mode 100644 index 409893b6..00000000 --- a/tests/archive/test_128_16_stepC.py +++ /dev/null @@ -1,383 +0,0 @@ -""" -Minimal PV-only test: Load P from GMEM to TMEM via QK-style MMA, then PV from TMEM. -Step 1: QK MMA writes FP32 S to TMEM (we know this works) -Step 2: Softmax packing writes BF16 P to TMEM (test this) -Step 3: PV MMA reads BF16 P from TMEM and V from SMEM, produces O - -But to isolate the bug, let me test just the PV MMA in isolation. -I'll write known BF16 values to TMEM using the softmax packing path, -then immediately read them back using the PV A-fragment path, -and compare. - -Actually, the simplest isolation test: -1. Do QK MMA to get S in TMEM (cosine 0.999999 verified) -2. Do softmax packing: S → P in TMEM (at offset 32) -3. Skip PV entirely — read P from TMEM using the C-fragment composition LOAD path -4. Output P to GMEM and compare against S.to(BF16) - -This tests whether the softmax packing writes P correctly to the same TMEM -that the PV would read from. - -But we can't easily read P from TMEM using the standard epilogue path -because the epilogue expects FP32 accumulator data. - -Alternative: Use the PV MMA with V=I (identity). If P is correct, -then P @ I = P. But V needs to be MN-major and (128, 128), not (128, 64). -The output would be (128, 128) which doesn't match our (128, 64) c tensor. - -Let me use V that selects the first 64 columns: V[k, n] = delta(k, n) for k in [0,63]. -This gives P @ V = P[:, :64], and the output is (128, 64). -But V is (128, 128) in the MMA K,N dims. V[k, n] for k in [0,127], n in [0,63]. -Hmm, this is getting complicated. Let me just do the identity approach with a (128, 128) output. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class Test128x16Tiler: - """QK + softmax packing + PV with V=I to isolate PV MMA correctness. - Output should be P = S.to(BF16), i.e. (Q@K^T).bfloat16() - With V=I, O = P @ I = P. - But V is (K=128, N=128) in the MMA. We need a 128x128 identity in MN-major. - Output tensor is (128, 128). - """ - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False # needed by epilogue_tma_store - self.epilog_sync_bar_id = 1 # needed by epilogue_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = int(cute.size(qk_mma.shape_mnk, mode=[2])) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - # PV with V=I: output is (128, 128), same as QK - self.pv_mma_tiler = (self.qk_mma_tiler[0], qk_inst_k, self.qk_mma_tiler[1]) - # pv_mma_tiler = (128, 128, 128) since V is 128x128 - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - (self.pv_mma_tiler[0], self.pv_mma_tiler[1], self.pv_mma_tiler[2]), False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - # ⛔⛔⛔ CRITICAL: num_tma_load_bytes MUST include ALL TMA-loaded tensors (Q + K + V). Missing V → DEADLOCK. See FOOTGUN #0 in README. - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - # PV with 128x128 output (V=I) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, (128, 16), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.pv_mma_tiler, (None,0,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_128_16_stepD.py b/tests/archive/test_128_16_stepD.py deleted file mode 100644 index 57d4b577..00000000 --- a/tests/archive/test_128_16_stepD.py +++ /dev/null @@ -1,386 +0,0 @@ -""" -Minimal PV-only test: Load P from GMEM to TMEM via QK-style MMA, then PV from TMEM. -Step 1: QK MMA writes FP32 S to TMEM (we know this works) -Step 2: Softmax packing writes BF16 P to TMEM (test this) -Step 3: PV MMA reads BF16 P from TMEM and V from SMEM, produces O - -But to isolate the bug, let me test just the PV MMA in isolation. -I'll write known BF16 values to TMEM using the softmax packing path, -then immediately read them back using the PV A-fragment path, -and compare. - -Actually, the simplest isolation test: -1. Do QK MMA to get S in TMEM (cosine 0.999999 verified) -2. Do softmax packing: S → P in TMEM (at offset 32) -3. Skip PV entirely — read P from TMEM using the C-fragment composition LOAD path -4. Output P to GMEM and compare against S.to(BF16) - -This tests whether the softmax packing writes P correctly to the same TMEM -that the PV would read from. - -But we can't easily read P from TMEM using the standard epilogue path -because the epilogue expects FP32 accumulator data. - -Alternative: Use the PV MMA with V=I (identity). If P is correct, -then P @ I = P. But V needs to be MN-major and (128, 128), not (128, 64). -The output would be (128, 128) which doesn't match our (128, 64) c tensor. - -Let me use V that selects the first 64 columns: V[k, n] = delta(k, n) for k in [0,63]. -This gives P @ V = P[:, :64], and the output is (128, 64). -But V is (128, 128) in the MMA K,N dims. V[k, n] for k in [0,127], n in [0,63]. -Hmm, this is getting complicated. Let me just do the identity approach with a (128, 128) output. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class Test128x16Tiler: - """QK + softmax packing + PV with V=I to isolate PV MMA correctness. - Output should be P = S.to(BF16), i.e. (Q@K^T).bfloat16() - With V=I, O = P @ I = P. - But V is (K=128, N=128) in the MMA. We need a 128x128 identity in MN-major. - Output tensor is (128, 128). - """ - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False # needed by epilogue_tma_store - self.epilog_sync_bar_id = 1 # needed by epilogue_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = int(cute.size(qk_mma.shape_mnk, mode=[2])) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - # PV with V=I: output is (128, 128), same as QK - self.pv_mma_tiler = (self.qk_mma_tiler[0], qk_inst_k, self.qk_mma_tiler[1]) - # pv_mma_tiler = (128, 128, 128) since V is 128x128 - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - (self.pv_mma_tiler[0], self.pv_mma_tiler[1], self.pv_mma_tiler[2]), False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - # ⛔⛔⛔ CRITICAL: num_tma_load_bytes MUST include ALL TMA-loaded tensors (Q + K + V). Missing V → DEADLOCK. See FOOTGUN #0 in README. - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - # PV with 128x128 output (V=I) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, (128, 16), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.pv_mma_tiler, (None,0,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_128_16_stepE.py b/tests/archive/test_128_16_stepE.py deleted file mode 100644 index d97125f3..00000000 --- a/tests/archive/test_128_16_stepE.py +++ /dev/null @@ -1,385 +0,0 @@ -""" -Minimal PV-only test: Load P from GMEM to TMEM via QK-style MMA, then PV from TMEM. -Step 1: QK MMA writes FP32 S to TMEM (we know this works) -Step 2: Softmax packing writes BF16 P to TMEM (test this) -Step 3: PV MMA reads BF16 P from TMEM and V from SMEM, produces O - -But to isolate the bug, let me test just the PV MMA in isolation. -I'll write known BF16 values to TMEM using the softmax packing path, -then immediately read them back using the PV A-fragment path, -and compare. - -Actually, the simplest isolation test: -1. Do QK MMA to get S in TMEM (cosine 0.999999 verified) -2. Do softmax packing: S → P in TMEM (at offset 32) -3. Skip PV entirely — read P from TMEM using the C-fragment composition LOAD path -4. Output P to GMEM and compare against S.to(BF16) - -This tests whether the softmax packing writes P correctly to the same TMEM -that the PV would read from. - -But we can't easily read P from TMEM using the standard epilogue path -because the epilogue expects FP32 accumulator data. - -Alternative: Use the PV MMA with V=I (identity). If P is correct, -then P @ I = P. But V needs to be MN-major and (128, 128), not (128, 64). -The output would be (128, 128) which doesn't match our (128, 64) c tensor. - -Let me use V that selects the first 64 columns: V[k, n] = delta(k, n) for k in [0,63]. -This gives P @ V = P[:, :64], and the output is (128, 64). -But V is (128, 128) in the MMA K,N dims. V[k, n] for k in [0,127], n in [0,63]. -Hmm, this is getting complicated. Let me just do the identity approach with a (128, 128) output. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class Test128x16Tiler: - """QK + softmax packing + PV with V=I to isolate PV MMA correctness. - Output should be P = S.to(BF16), i.e. (Q@K^T).bfloat16() - With V=I, O = P @ I = P. - But V is (K=128, N=128) in the MMA. We need a 128x128 identity in MN-major. - Output tensor is (128, 128). - """ - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False # needed by epilogue_tma_store - self.epilog_sync_bar_id = 1 # needed by epilogue_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = int(cute.size(qk_mma.shape_mnk, mode=[2])) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - # PV with V=I: output is (128, 128), same as QK - self.pv_mma_tiler = (self.qk_mma_tiler[0], qk_inst_k, self.qk_mma_tiler[1]) - # pv_mma_tiler = (128, 128, 128) since V is 128x128 - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - (self.pv_mma_tiler[0], self.pv_mma_tiler[1], self.pv_mma_tiler[2]), False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - # ⛔⛔⛔ CRITICAL: num_tma_load_bytes MUST include ALL TMA-loaded tensors (Q + K + V). Missing V → DEADLOCK. See FOOTGUN #0 in README. - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - # PV with 128x128 output (V=I) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, (128, 16), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.pv_mma_tiler, (None,0,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_128_16_tiler.py b/tests/archive/test_128_16_tiler.py deleted file mode 100644 index 1f72c097..00000000 --- a/tests/archive/test_128_16_tiler.py +++ /dev/null @@ -1,383 +0,0 @@ -""" -Minimal PV-only test: Load P from GMEM to TMEM via QK-style MMA, then PV from TMEM. -Step 1: QK MMA writes FP32 S to TMEM (we know this works) -Step 2: Softmax packing writes BF16 P to TMEM (test this) -Step 3: PV MMA reads BF16 P from TMEM and V from SMEM, produces O - -But to isolate the bug, let me test just the PV MMA in isolation. -I'll write known BF16 values to TMEM using the softmax packing path, -then immediately read them back using the PV A-fragment path, -and compare. - -Actually, the simplest isolation test: -1. Do QK MMA to get S in TMEM (cosine 0.999999 verified) -2. Do softmax packing: S → P in TMEM (at offset 32) -3. Skip PV entirely — read P from TMEM using the C-fragment composition LOAD path -4. Output P to GMEM and compare against S.to(BF16) - -This tests whether the softmax packing writes P correctly to the same TMEM -that the PV would read from. - -But we can't easily read P from TMEM using the standard epilogue path -because the epilogue expects FP32 accumulator data. - -Alternative: Use the PV MMA with V=I (identity). If P is correct, -then P @ I = P. But V needs to be MN-major and (128, 128), not (128, 64). -The output would be (128, 128) which doesn't match our (128, 64) c tensor. - -Let me use V that selects the first 64 columns: V[k, n] = delta(k, n) for k in [0,63]. -This gives P @ V = P[:, :64], and the output is (128, 64). -But V is (128, 128) in the MMA K,N dims. V[k, n] for k in [0,127], n in [0,63]. -Hmm, this is getting complicated. Let me just do the identity approach with a (128, 128) output. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class Test128x16Tiler: - """QK + softmax packing + PV with V=I to isolate PV MMA correctness. - Output should be P = S.to(BF16), i.e. (Q@K^T).bfloat16() - With V=I, O = P @ I = P. - But V is (K=128, N=128) in the MMA. We need a 128x128 identity in MN-major. - Output tensor is (128, 128). - """ - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False # needed by epilogue_tma_store - self.epilog_sync_bar_id = 1 # needed by epilogue_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = int(cute.size(qk_mma.shape_mnk, mode=[2])) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - # PV with V=I: output is (128, 128), same as QK - self.pv_mma_tiler = (self.qk_mma_tiler[0], qk_inst_k, self.qk_mma_tiler[1]) - # pv_mma_tiler = (128, 128, 128) since V is 128x128 - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - # ⛔⛔⛔ CRITICAL: num_tma_load_bytes MUST include ALL TMA-loaded tensors (Q + K + V). Missing V → DEADLOCK. See FOOTGUN #0 in README. - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - # PV with 128x128 output (V=I) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_128_16_v8.py b/tests/archive/test_128_16_v8.py deleted file mode 100644 index 598fc38f..00000000 --- a/tests/archive/test_128_16_v8.py +++ /dev/null @@ -1,385 +0,0 @@ -""" -Minimal PV-only test: Load P from GMEM to TMEM via QK-style MMA, then PV from TMEM. -Step 1: QK MMA writes FP32 S to TMEM (we know this works) -Step 2: Softmax packing writes BF16 P to TMEM (test this) -Step 3: PV MMA reads BF16 P from TMEM and V from SMEM, produces O - -But to isolate the bug, let me test just the PV MMA in isolation. -I'll write known BF16 values to TMEM using the softmax packing path, -then immediately read them back using the PV A-fragment path, -and compare. - -Actually, the simplest isolation test: -1. Do QK MMA to get S in TMEM (cosine 0.999999 verified) -2. Do softmax packing: S → P in TMEM (at offset 32) -3. Skip PV entirely — read P from TMEM using the C-fragment composition LOAD path -4. Output P to GMEM and compare against S.to(BF16) - -This tests whether the softmax packing writes P correctly to the same TMEM -that the PV would read from. - -But we can't easily read P from TMEM using the standard epilogue path -because the epilogue expects FP32 accumulator data. - -Alternative: Use the PV MMA with V=I (identity). If P is correct, -then P @ I = P. But V needs to be MN-major and (128, 128), not (128, 64). -The output would be (128, 128) which doesn't match our (128, 64) c tensor. - -Let me use V that selects the first 64 columns: V[k, n] = delta(k, n) for k in [0,63]. -This gives P @ V = P[:, :64], and the output is (128, 64). -But V is (128, 128) in the MMA K,N dims. V[k, n] for k in [0,127], n in [0,63]. -Hmm, this is getting complicated. Let me just do the identity approach with a (128, 128) output. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class Test128x16Tiler: - """QK + softmax packing + PV with V=I to isolate PV MMA correctness. - Output should be P = S.to(BF16), i.e. (Q@K^T).bfloat16() - With V=I, O = P @ I = P. - But V is (K=128, N=128) in the MMA. We need a 128x128 identity in MN-major. - Output tensor is (128, 128). - """ - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False # needed by epilogue_tma_store - self.epilog_sync_bar_id = 1 # needed by epilogue_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = int(cute.size(qk_mma.shape_mnk, mode=[2])) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - # PV with V=I: output is (128, 128), same as QK - self.pv_mma_tiler = (self.qk_mma_tiler[0], qk_inst_k, self.qk_mma_tiler[1]) - # pv_mma_tiler = (128, 128, 128) since V is 128x128 - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - (self.pv_mma_tiler[0], self.pv_mma_tiler[1], self.pv_mma_tiler[2]), False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - # ⛔⛔⛔ CRITICAL: num_tma_load_bytes MUST include ALL TMA-loaded tensors (Q + K + V). Missing V → DEADLOCK. See FOOTGUN #0 in README. - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - # PV with 128x128 output (V=I) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, (128, 16), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.pv_mma_tiler, (None,0,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_128_16_zeropad.py b/tests/archive/test_128_16_zeropad.py deleted file mode 100644 index ef014d9c..00000000 --- a/tests/archive/test_128_16_zeropad.py +++ /dev/null @@ -1,385 +0,0 @@ -""" -Minimal PV-only test: Load P from GMEM to TMEM via QK-style MMA, then PV from TMEM. -Step 1: QK MMA writes FP32 S to TMEM (we know this works) -Step 2: Softmax packing writes BF16 P to TMEM (test this) -Step 3: PV MMA reads BF16 P from TMEM and V from SMEM, produces O - -But to isolate the bug, let me test just the PV MMA in isolation. -I'll write known BF16 values to TMEM using the softmax packing path, -then immediately read them back using the PV A-fragment path, -and compare. - -Actually, the simplest isolation test: -1. Do QK MMA to get S in TMEM (cosine 0.999999 verified) -2. Do softmax packing: S → P in TMEM (at offset 32) -3. Skip PV entirely — read P from TMEM using the C-fragment composition LOAD path -4. Output P to GMEM and compare against S.to(BF16) - -This tests whether the softmax packing writes P correctly to the same TMEM -that the PV would read from. - -But we can't easily read P from TMEM using the standard epilogue path -because the epilogue expects FP32 accumulator data. - -Alternative: Use the PV MMA with V=I (identity). If P is correct, -then P @ I = P. But V needs to be MN-major and (128, 128), not (128, 64). -The output would be (128, 128) which doesn't match our (128, 64) c tensor. - -Let me use V that selects the first 64 columns: V[k, n] = delta(k, n) for k in [0,63]. -This gives P @ V = P[:, :64], and the output is (128, 64). -But V is (128, 128) in the MMA K,N dims. V[k, n] for k in [0,127], n in [0,63]. -Hmm, this is getting complicated. Let me just do the identity approach with a (128, 128) output. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class ZeroPadV8Kernel: - """QK + softmax packing + PV with V=I to isolate PV MMA correctness. - Output should be P = S.to(BF16), i.e. (Q@K^T).bfloat16() - With V=I, O = P @ I = P. - But V is (K=128, N=128) in the MMA. We need a 128x128 identity in MN-major. - Output tensor is (128, 128). - """ - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False # needed by epilogue_tma_store - self.epilog_sync_bar_id = 1 # needed by epilogue_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - # PV with V=I: output is (128, 128), same as QK - self.pv_mma_tiler = (self.qk_mma_tiler[0], self.qk_mma_tiler[1], self.qk_mma_tiler[1]) - # pv_mma_tiler = (128, 128, 128) since V is 128x128 - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - # ⛔⛔⛔ CRITICAL: num_tma_load_bytes MUST include ALL TMA-loaded tensors (Q + K + V). Missing V → DEADLOCK. See FOOTGUN #0 in README. - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - # PV with 128x128 output (V=I) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1 O = (Q@K^T).bf16()[:,:16], rest zeros - ref_full = (qf @ kf.T).bfloat16().float() - ref = torch.zeros_like(ref_full) - ref[:, :16] = ref_full[:, :16] - - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - kernel = ZeroPadV8Kernel(mma_tiler_mn=(128, 128)) - print('Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - print('Running...', flush=True) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - out = c[:,:,0].float() - cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - print('PV(128,128) zero-pad: cosine {:.6f} {}'.format(cos, 'PASS' if cos >= 0.99 else 'FAIL')) - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_128_32_ctafix.py b/tests/archive/test_128_32_ctafix.py deleted file mode 100644 index 0b1f2815..00000000 --- a/tests/archive/test_128_32_ctafix.py +++ /dev/null @@ -1,386 +0,0 @@ -""" -Minimal PV-only test: Load P from GMEM to TMEM via QK-style MMA, then PV from TMEM. -Step 1: QK MMA writes FP32 S to TMEM (we know this works) -Step 2: Softmax packing writes BF16 P to TMEM (test this) -Step 3: PV MMA reads BF16 P from TMEM and V from SMEM, produces O - -But to isolate the bug, let me test just the PV MMA in isolation. -I'll write known BF16 values to TMEM using the softmax packing path, -then immediately read them back using the PV A-fragment path, -and compare. - -Actually, the simplest isolation test: -1. Do QK MMA to get S in TMEM (cosine 0.999999 verified) -2. Do softmax packing: S → P in TMEM (at offset 32) -3. Skip PV entirely — read P from TMEM using the C-fragment composition LOAD path -4. Output P to GMEM and compare against S.to(BF16) - -This tests whether the softmax packing writes P correctly to the same TMEM -that the PV would read from. - -But we can't easily read P from TMEM using the standard epilogue path -because the epilogue expects FP32 accumulator data. - -Alternative: Use the PV MMA with V=I (identity). If P is correct, -then P @ I = P. But V needs to be MN-major and (128, 128), not (128, 64). -The output would be (128, 128) which doesn't match our (128, 64) c tensor. - -Let me use V that selects the first 64 columns: V[k, n] = delta(k, n) for k in [0,63]. -This gives P @ V = P[:, :64], and the output is (128, 64). -But V is (128, 128) in the MMA K,N dims. V[k, n] for k in [0,127], n in [0,63]. -Hmm, this is getting complicated. Let me just do the identity approach with a (128, 128) output. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class Native32Kernel: - """QK + softmax packing + PV with V=I to isolate PV MMA correctness. - Output should be P = S.to(BF16), i.e. (Q@K^T).bfloat16() - With V=I, O = P @ I = P. - But V is (K=128, N=128) in the MMA. We need a 128x128 identity in MN-major. - Output tensor is (128, 128). - """ - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False # needed by epilogue_tma_store - self.epilog_sync_bar_id = 1 # needed by epilogue_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = int(cute.size(qk_mma.shape_mnk, mode=[2])) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - # PV with V=I: output is (128, 128), same as QK - self.pv_mma_tiler = (self.qk_mma_tiler[0], qk_inst_k, self.qk_mma_tiler[1]) - # pv_mma_tiler = (128, 128, 128) since V is 128x128 - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - (self.pv_mma_tiler[0], self.pv_mma_tiler[1], self.pv_mma_tiler[2]), False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - # ⛔⛔⛔ CRITICAL: num_tma_load_bytes MUST include ALL TMA-loaded tensors (Q + K + V). Missing V → DEADLOCK. See FOOTGUN #0 in README. - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - # PV with 128x128 output (V=I) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, (128, 32), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1 O = P[:,:32] = (Q@K^T).bf16()[:,:32] - ref = (qf @ kf.T).bfloat16().float()[:, :32] - - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - kernel = Native32Kernel(mma_tiler_mn=(128, 128)) - print('Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - print('Running...', flush=True) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - out = c[:,:,0].float() - cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - print('PV(128,32) ctafix: cosine {:.6f} {}'.format(cos, 'PASS' if cos >= 0.99 else 'FAIL')) - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_128_32_ctafix2.py b/tests/archive/test_128_32_ctafix2.py deleted file mode 100644 index 2b6cddc1..00000000 --- a/tests/archive/test_128_32_ctafix2.py +++ /dev/null @@ -1,384 +0,0 @@ -""" -Minimal PV-only test: Load P from GMEM to TMEM via QK-style MMA, then PV from TMEM. -Step 1: QK MMA writes FP32 S to TMEM (we know this works) -Step 2: Softmax packing writes BF16 P to TMEM (test this) -Step 3: PV MMA reads BF16 P from TMEM and V from SMEM, produces O - -But to isolate the bug, let me test just the PV MMA in isolation. -I'll write known BF16 values to TMEM using the softmax packing path, -then immediately read them back using the PV A-fragment path, -and compare. - -Actually, the simplest isolation test: -1. Do QK MMA to get S in TMEM (cosine 0.999999 verified) -2. Do softmax packing: S → P in TMEM (at offset 32) -3. Skip PV entirely — read P from TMEM using the C-fragment composition LOAD path -4. Output P to GMEM and compare against S.to(BF16) - -This tests whether the softmax packing writes P correctly to the same TMEM -that the PV would read from. - -But we can't easily read P from TMEM using the standard epilogue path -because the epilogue expects FP32 accumulator data. - -Alternative: Use the PV MMA with V=I (identity). If P is correct, -then P @ I = P. But V needs to be MN-major and (128, 128), not (128, 64). -The output would be (128, 128) which doesn't match our (128, 64) c tensor. - -Let me use V that selects the first 64 columns: V[k, n] = delta(k, n) for k in [0,63]. -This gives P @ V = P[:, :64], and the output is (128, 64). -But V is (128, 128) in the MMA K,N dims. V[k, n] for k in [0,127], n in [0,63]. -Hmm, this is getting complicated. Let me just do the identity approach with a (128, 128) output. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class Native32Kernel: - """QK + softmax packing + PV with V=I to isolate PV MMA correctness. - Output should be P = S.to(BF16), i.e. (Q@K^T).bfloat16() - With V=I, O = P @ I = P. - But V is (K=128, N=128) in the MMA. We need a 128x128 identity in MN-major. - Output tensor is (128, 128). - """ - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False # needed by epilogue_tma_store - self.epilog_sync_bar_id = 1 # needed by epilogue_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - self.pv_cta_tile = (128, 16, 128) # Will be updated in _setup with int values - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = int(cute.size(qk_mma.shape_mnk, mode=[2])) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - # PV with V=I: output is (128, 128), same as QK - self.pv_mma_tiler = (self.qk_mma_tiler[0], qk_inst_k, self.qk_mma_tiler[1]) - self.pv_cta_tile = (int(self.pv_mma_tiler[0]), int(self.pv_mma_tiler[1]), int(self.pv_mma_tiler[2])) - # pv_mma_tiler = (128, 128, 128) since V is 128x128 - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - (self.pv_mma_tiler[0], self.pv_mma_tiler[1], self.pv_mma_tiler[2]), False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - # ⛔⛔⛔ CRITICAL: num_tma_load_bytes MUST include ALL TMA-loaded tensors (Q + K + V). Missing V → DEADLOCK. See FOOTGUN #0 in README. - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - # PV with 128x128 output (V=I) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, (128, 32), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1 O = P[:,:32] = (Q@K^T).bf16()[:,:32] - ref = (qf @ kf.T).bfloat16().float()[:, :32] - - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - kernel = Native32Kernel(mma_tiler_mn=(128, 128)) - print('Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - print('Running...', flush=True) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - out = c[:,:,0].float() - cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - print('PV(128,32) ctafix2: cosine {:.6f} {}'.format(cos, 'PASS' if cos >= 0.99 else 'FAIL')) - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_128_32_native.py b/tests/archive/test_128_32_native.py deleted file mode 100644 index 93c0864f..00000000 --- a/tests/archive/test_128_32_native.py +++ /dev/null @@ -1,382 +0,0 @@ -""" -Minimal PV-only test: Load P from GMEM to TMEM via QK-style MMA, then PV from TMEM. -Step 1: QK MMA writes FP32 S to TMEM (we know this works) -Step 2: Softmax packing writes BF16 P to TMEM (test this) -Step 3: PV MMA reads BF16 P from TMEM and V from SMEM, produces O - -But to isolate the bug, let me test just the PV MMA in isolation. -I'll write known BF16 values to TMEM using the softmax packing path, -then immediately read them back using the PV A-fragment path, -and compare. - -Actually, the simplest isolation test: -1. Do QK MMA to get S in TMEM (cosine 0.999999 verified) -2. Do softmax packing: S → P in TMEM (at offset 32) -3. Skip PV entirely — read P from TMEM using the C-fragment composition LOAD path -4. Output P to GMEM and compare against S.to(BF16) - -This tests whether the softmax packing writes P correctly to the same TMEM -that the PV would read from. - -But we can't easily read P from TMEM using the standard epilogue path -because the epilogue expects FP32 accumulator data. - -Alternative: Use the PV MMA with V=I (identity). If P is correct, -then P @ I = P. But V needs to be MN-major and (128, 128), not (128, 64). -The output would be (128, 128) which doesn't match our (128, 64) c tensor. - -Let me use V that selects the first 64 columns: V[k, n] = delta(k, n) for k in [0,63]. -This gives P @ V = P[:, :64], and the output is (128, 64). -But V is (128, 128) in the MMA K,N dims. V[k, n] for k in [0,127], n in [0,63]. -Hmm, this is getting complicated. Let me just do the identity approach with a (128, 128) output. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class Native32Kernel: - """QK + softmax packing + PV with V=I to isolate PV MMA correctness. - Output should be P = S.to(BF16), i.e. (Q@K^T).bfloat16() - With V=I, O = P @ I = P. - But V is (K=128, N=128) in the MMA. We need a 128x128 identity in MN-major. - Output tensor is (128, 128). - """ - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False # needed by epilogue_tma_store - self.epilog_sync_bar_id = 1 # needed by epilogue_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = int(cute.size(qk_mma.shape_mnk, mode=[2])) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - # PV with V=I: output is (128, 128), same as QK - self.pv_mma_tiler = (self.qk_mma_tiler[0], qk_inst_k, self.qk_mma_tiler[1]) - # pv_mma_tiler = (128, 128, 128) since V is 128x128 - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - (self.pv_mma_tiler[0], self.pv_mma_tiler[1], self.pv_mma_tiler[2]), False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - # ⛔⛔⛔ CRITICAL: num_tma_load_bytes MUST include ALL TMA-loaded tensors (Q + K + V). Missing V → DEADLOCK. See FOOTGUN #0 in README. - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - # PV with 128x128 output (V=I) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, (128, 32), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1 O = P[:,:32] = (Q@K^T).bf16()[:,:32] - ref = (qf @ kf.T).bfloat16().float()[:, :32] - - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - kernel = Native32Kernel(mma_tiler_mn=(128, 128)) - print('Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - print('Running...', flush=True) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - out = c[:,:,0].float() - cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - print('PV(128,32) native: cosine {:.6f} {}'.format(cos, 'PASS' if cos >= 0.99 else 'FAIL')) - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_128_32_vdiag.py b/tests/archive/test_128_32_vdiag.py deleted file mode 100644 index 19373c56..00000000 --- a/tests/archive/test_128_32_vdiag.py +++ /dev/null @@ -1,385 +0,0 @@ -""" -Minimal PV-only test: Load P from GMEM to TMEM via QK-style MMA, then PV from TMEM. -Step 1: QK MMA writes FP32 S to TMEM (we know this works) -Step 2: Softmax packing writes BF16 P to TMEM (test this) -Step 3: PV MMA reads BF16 P from TMEM and V from SMEM, produces O - -But to isolate the bug, let me test just the PV MMA in isolation. -I'll write known BF16 values to TMEM using the softmax packing path, -then immediately read them back using the PV A-fragment path, -and compare. - -Actually, the simplest isolation test: -1. Do QK MMA to get S in TMEM (cosine 0.999999 verified) -2. Do softmax packing: S → P in TMEM (at offset 32) -3. Skip PV entirely — read P from TMEM using the C-fragment composition LOAD path -4. Output P to GMEM and compare against S.to(BF16) - -This tests whether the softmax packing writes P correctly to the same TMEM -that the PV would read from. - -But we can't easily read P from TMEM using the standard epilogue path -because the epilogue expects FP32 accumulator data. - -Alternative: Use the PV MMA with V=I (identity). If P is correct, -then P @ I = P. But V needs to be MN-major and (128, 128), not (128, 64). -The output would be (128, 128) which doesn't match our (128, 64) c tensor. - -Let me use V that selects the first 64 columns: V[k, n] = delta(k, n) for k in [0,63]. -This gives P @ V = P[:, :64], and the output is (128, 64). -But V is (128, 128) in the MMA K,N dims. V[k, n] for k in [0,127], n in [0,63]. -Hmm, this is getting complicated. Let me just do the identity approach with a (128, 128) output. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class VDiag32: - """QK + softmax packing + PV with V=I to isolate PV MMA correctness. - Output should be P = S.to(BF16), i.e. (Q@K^T).bfloat16() - With V=I, O = P @ I = P. - But V is (K=128, N=128) in the MMA. We need a 128x128 identity in MN-major. - Output tensor is (128, 128). - """ - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False # needed by epilogue_tma_store - self.epilog_sync_bar_id = 1 # needed by epilogue_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = int(cute.size(qk_mma.shape_mnk, mode=[2])) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - # PV with V=I: output is (128, 128), same as QK - self.pv_mma_tiler = (self.qk_mma_tiler[0], qk_inst_k, self.qk_mma_tiler[1]) - # pv_mma_tiler = (128, 128, 128) since V is 128x128 - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - (self.pv_mma_tiler[0], self.pv_mma_tiler[1], self.pv_mma_tiler[2]), False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - # ⛔⛔⛔ CRITICAL: num_tma_load_bytes MUST include ALL TMA-loaded tensors (Q + K + V). Missing V → DEADLOCK. See FOOTGUN #0 in README. - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - # PV with 128x128 output (V=I) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, (128, 32), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_128_32_zeropad.py b/tests/archive/test_128_32_zeropad.py deleted file mode 100644 index 286a3758..00000000 --- a/tests/archive/test_128_32_zeropad.py +++ /dev/null @@ -1,385 +0,0 @@ -""" -Minimal PV-only test: Load P from GMEM to TMEM via QK-style MMA, then PV from TMEM. -Step 1: QK MMA writes FP32 S to TMEM (we know this works) -Step 2: Softmax packing writes BF16 P to TMEM (test this) -Step 3: PV MMA reads BF16 P from TMEM and V from SMEM, produces O - -But to isolate the bug, let me test just the PV MMA in isolation. -I'll write known BF16 values to TMEM using the softmax packing path, -then immediately read them back using the PV A-fragment path, -and compare. - -Actually, the simplest isolation test: -1. Do QK MMA to get S in TMEM (cosine 0.999999 verified) -2. Do softmax packing: S → P in TMEM (at offset 32) -3. Skip PV entirely — read P from TMEM using the C-fragment composition LOAD path -4. Output P to GMEM and compare against S.to(BF16) - -This tests whether the softmax packing writes P correctly to the same TMEM -that the PV would read from. - -But we can't easily read P from TMEM using the standard epilogue path -because the epilogue expects FP32 accumulator data. - -Alternative: Use the PV MMA with V=I (identity). If P is correct, -then P @ I = P. But V needs to be MN-major and (128, 128), not (128, 64). -The output would be (128, 128) which doesn't match our (128, 64) c tensor. - -Let me use V that selects the first 64 columns: V[k, n] = delta(k, n) for k in [0,63]. -This gives P @ V = P[:, :64], and the output is (128, 64). -But V is (128, 128) in the MMA K,N dims. V[k, n] for k in [0,127], n in [0,63]. -Hmm, this is getting complicated. Let me just do the identity approach with a (128, 128) output. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class ZeroPad32Kernel: - """QK + softmax packing + PV with V=I to isolate PV MMA correctness. - Output should be P = S.to(BF16), i.e. (Q@K^T).bfloat16() - With V=I, O = P @ I = P. - But V is (K=128, N=128) in the MMA. We need a 128x128 identity in MN-major. - Output tensor is (128, 128). - """ - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False # needed by epilogue_tma_store - self.epilog_sync_bar_id = 1 # needed by epilogue_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - # PV with V=I: output is (128, 128), same as QK - self.pv_mma_tiler = (self.qk_mma_tiler[0], self.qk_mma_tiler[1], self.qk_mma_tiler[1]) - # pv_mma_tiler = (128, 128, 128) since V is 128x128 - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - # ⛔⛔⛔ CRITICAL: num_tma_load_bytes MUST include ALL TMA-loaded tensors (Q + K + V). Missing V → DEADLOCK. See FOOTGUN #0 in README. - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - # PV with 128x128 output (V=I) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1 O = (Q@K^T).bf16()[:,:16], rest zeros - ref_full = (qf @ kf.T).bfloat16().float() - ref = torch.zeros_like(ref_full) - ref[:, :32] = ref_full[:, :32] - - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - kernel = ZeroPad32Kernel(mma_tiler_mn=(128, 128)) - print('Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - print('Running...', flush=True) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - out = c[:,:,0].float() - cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - print('PV(128,32) zero-pad: cosine {:.6f} {}'.format(cos, 'PASS' if cos >= 0.99 else 'FAIL')) - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_afrag_roundtrip.py b/tests/archive/test_afrag_roundtrip.py deleted file mode 100644 index 12237293..00000000 --- a/tests/archive/test_afrag_roundtrip.py +++ /dev/null @@ -1,175 +0,0 @@ -"""Test: ld FP32 from S0, st BF16 to A-fragment layout tdVrP, -ld BF16 back from tdVrP, epi the result. -If this works, the A-fragment store is correct and the issue is in the PV MMA.""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - -class AFragRoundtrip: - def __init__(self, mma_tiler_mn): - self.qk_acc_dtype = Float32; self.q_dtype = BFloat16; self.o_dtype = BFloat16 - self.c_dtype = BFloat16; self.acc_dtype = Float32 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.num_c_stage = 2; self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - pv_inst_k = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (*self.mma_tiler_mn, pv_inst_k * 4) - self.mma_tiler = self.qk_mma_tiler - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - c_layout = LayoutEnum.ROW_MAJOR; self.c_layout = c_layout - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, c_layout, self.o_dtype) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, c_layout, self.epi_tile, 2) - self.num_ab_stage = 1; self.num_acc_stage = 1 - qk_thr = qk_mma.get_slice(0); qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape); self.s_cols = find_tmem_tensor_col_offset(tStS) - self.tmem_alloc_cols = self.s_cols # Only need S region - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)); b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = (cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem)) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, LayoutEnum.from_tensor(a).mma_major_mode(), LayoutEnum.from_tensor(b).mma_major_mode(), self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, LayoutEnum.from_tensor(b).mma_major_mode(), self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)); b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - tma_a, tma_ta = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), a, a_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_b, tma_tb = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), b, b_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - self._kernel(qk_mma, pv_mma, tma_a, tma_ta, tma_b, tma_tb, tma_c, tma_tc, self.cluster_layout_vmnk, self.a_smem_s, self.b_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_a, mA, tma_b, mB, tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()); tidx, _, _ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: cpasync.prefetch_descriptor(tma_a); cpasync.prefetch_descriptor(tma_b); cpasync.prefetch_descriptor(tma_c) - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2]; mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2]; tmem_dealloc: cutlass.Int64; holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - ab_p, ab_c = pipeline.PipelineTmaUmma.create(barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True).make_participants() - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), cta_layout_vmnk=cl_vmnk, defer_sync=True).make_participants() - acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, len(self.epilogue_warp_id)), cta_layout_vmnk=cl_vmnk, defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2, num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=False, two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - sA = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sB = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - gA = cute.local_tile(mA, cute.slice_(self.mma_tiler, (None,0,None)), (None,None,None)) - gB = cute.local_tile(mB, cute.slice_(self.mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gA, mode=[3]) - qk_thr = qk_mma.get_slice(0); tCgA = qk_thr.partition_A(gA); tCgB = qk_thr.partition_B(gB); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsA, tAgA = cpasync.tma_partition(tma_a, 0, a_lay, cute.group_modes(sA,0,3), cute.group_modes(tCgA,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsB, tBgB = cpasync.tma_partition(tma_b, 0, b_lay, cute.group_modes(sB,0,3), cute.group_modes(tCgB,0,3)) - tAgA = tAgA[(None,0,None,0)]; tBgB = tBgB[(None,0,None,0)] - tCrA = qk_mma.make_fragment_A(sA); tCrB = qk_mma.make_fragment_B(sB) - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator, tStS.layout) - # A-fragment for pv_mma - pv_thr = pv_mma.get_slice(0) - tP_iter = cute.recast_ptr(tStS.iterator, dtype=self.q_dtype) - tP = cute.make_tensor(tP_iter, p_tmem_s.outer) - tOrP = pv_thr.make_fragment_A(tP)[None, None, None, 0] - tdVrP = cute.make_tensor(tOrP.iterator, tOrP.layout) - # TMEM ld (C-fragment) - tmem_ld = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_ld = tcgen05.make_tmem_copy(tmem_ld, tStS0) - sfw = tidx % (32 * len(self.epilogue_warp_id)) - thr_ld = tiled_ld.get_slice(sfw) - tLdS = thr_ld.partition_S(tStS0) - cS_id = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS_id) - tLdcS = thr_ld.partition_D(tScS) - # TMEM st (A-fragment layout, BF16) - tmem_st = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(8)), self.q_dtype) - tiled_st = tcgen05.make_tmem_copy(tmem_st, tdVrP) - thr_st = tiled_st.get_slice(sfw) - tStP = thr_st.partition_D(tdVrP) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, 1)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - # TMA - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek); cute.copy(tma_a, tAgA[(None,h.count)], tAsA[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_b, tBgB[(None,h.count)], tBsB[(None,h.index)], tma_bar_ptr=h.barrier); peek = cutlass.Boolean(1) - if h.count+1> 4) & 0x0F).long()] - out_features = packed_uint8.shape[0] - in_features = packed_uint8.shape[1] * 2 - unpacked = torch.empty(out_features, in_features, dtype=torch.float32, device=device) - unpacked[:, 0::2] = lower - unpacked[:, 1::2] = upper - block_scale = scale_e4m3.float() - block_expanded = block_scale.repeat_interleave(16, dim=1)[:out_features, :in_features] - return (unpacked * block_expanded * global_scale).to(torch.bfloat16) - - -def test_projection(name, weight, weight_sf, weight_gs, hidden_states, in_features, out_features): - """Test a single NVFP4 projection.""" - sys.path.insert(0, "/root/nvfp4-megamoe-kernel") - from dsv4.layers.linear import Nvfp4Linear - - # Convert weight to CuTeDSL format: (out, in_packed) uint8 → (in_packed, out) float4 - fp4 = [weight.view(torch.float4_e2m1fn_x2).permute(1, 0).contiguous()] - sf = [weight_sf.permute(1, 0).contiguous()] - gs = [weight_gs] - - runner = Nvfp4Linear( - in_features=in_features, - out_features=out_features, - max_num_tokens=8192, - device=DEVICE, - ) - runner.fp4 = fp4 - runner.sf = sf - runner.gs = gs - runner.finalize_weights() - - # Warmup - runner._ensure_initialized() - runner.compute_activation_global_scale(hidden_states) - - # Run CuTeDSL - with torch.no_grad(): - output = runner.run(hidden_states) - - # BF16 reference - bf16_w = dequant_nvfp4(weight, weight_sf, weight_gs) - with torch.no_grad(): - ref = hidden_states @ bf16_w.T - - # Compare - cos = F.cosine_similarity(ref.flatten().unsqueeze(0), - output.flatten().unsqueeze(0)).item() - mse = (ref - output).pow(2).mean().item() - status = "✅" if cos >= 0.98 else "❌" - print(f" {name}: cosine={cos:.6f} MSE={mse:.6e} amax_ref={ref.amax():.4f} amax_out={output.amax():.4f} {status}") - return cos - - -def main(): - torch.cuda.set_device(0) - torch.manual_seed(42) - - with open(os.path.join(MODEL_PATH, "model.safetensors.index.json")) as f: - wm = json.load(f)["weight_map"] - P = lambda key: load_tensor(key, wm, MODEL_PATH).to(DEVICE) - - prefix = f"model.layers.{LAYER_IDX}.self_attn" - - print("=== Attention Projection Tests (CuTeDSL NVFP4 Linear) ===\n") - - # Load weights and determine dimensions from shapes - projs = { - "q_a_proj": {"key": f"{prefix}.q_a_proj"}, - "q_b_proj": {"key": f"{prefix}.q_b_proj"}, - "kv_proj": {"key": f"{prefix}.kv_proj"}, - "o_b_proj": {"key": f"{prefix}.o_b_proj"}, - } - - for name, info in projs.items(): - key = info["key"] - w = P(f"{key}.weight") - sf = P(f"{key}.weight_scale") - gs = P(f"{key}.weight_scale_2").item() - out_features = w.shape[0] - in_features = w.shape[1] * 2 # unpacked - info["weight"] = w - info["sf"] = sf - info["gs"] = gs - info["in_features"] = in_features - info["out_features"] = out_features - print(f" {name}: weight={w.shape} → in={in_features} out={out_features} gs={gs:.8f}") - - print() - - # Test each projection - # q_a_proj: input is hidden_states (HIDDEN_SIZE=7168) - hidden = torch.randn(NUM_TOKENS, HIDDEN_SIZE, dtype=torch.bfloat16, device=DEVICE) * 2.0 - - cos_qa = test_projection("q_a_proj", projs["q_a_proj"]["weight"], - projs["q_a_proj"]["sf"], projs["q_a_proj"]["gs"], - hidden, projs["q_a_proj"]["in_features"], projs["q_a_proj"]["out_features"]) - - # q_b_proj: input is q_a output (1536 features) - q_a_out_features = projs["q_a_proj"]["out_features"] - q_a_out = torch.randn(NUM_TOKENS, q_a_out_features, dtype=torch.bfloat16, device=DEVICE) * 2.0 - cos_qb = test_projection("q_b_proj", projs["q_b_proj"]["weight"], - projs["q_b_proj"]["sf"], projs["q_b_proj"]["gs"], - q_a_out, projs["q_b_proj"]["in_features"], projs["q_b_proj"]["out_features"]) - - # kv_proj: input is hidden_states (7168) - cos_kv = test_projection("kv_proj", projs["kv_proj"]["weight"], - projs["kv_proj"]["sf"], projs["kv_proj"]["gs"], - hidden, projs["kv_proj"]["in_features"], projs["kv_proj"]["out_features"]) - - # o_b_proj: input is o_a output (16384 features after attention) - o_b_in_features = projs["o_b_proj"]["in_features"] - o_b_input = torch.randn(NUM_TOKENS, o_b_in_features, dtype=torch.bfloat16, device=DEVICE) * 2.0 - cos_ob = test_projection("o_b_proj", projs["o_b_proj"]["weight"], - projs["o_b_proj"]["sf"], projs["o_b_proj"]["gs"], - o_b_input, projs["o_b_proj"]["in_features"], projs["o_b_proj"]["out_features"]) - - print(f"\n=== SUMMARY ===") - results = {"q_a_proj": cos_qa, "q_b_proj": cos_qb, "kv_proj": cos_kv, "o_b_proj": cos_ob} - all_pass = True - for name, cos in results.items(): - status = "✅" if cos >= 0.98 else "❌" - if cos < 0.98: - all_pass = False - print(f" {name}: cosine={cos:.6f} {status}") - - if all_pass: - print("\n✅ ALL PASS") - else: - print("\n❌ SOME FAILED") - - -if __name__ == "__main__": - main() diff --git a/tests/archive/test_attention_path_b200.py b/tests/archive/test_attention_path_b200.py deleted file mode 100644 index 334a513c..00000000 --- a/tests/archive/test_attention_path_b200.py +++ /dev/null @@ -1,267 +0,0 @@ -#!/usr/bin/env python3 -""" -Pinpoint which vLLM attention component fails on B200. - -Runs each step of the attention forward pass individually: -1. fused_wqa_wkv (CuTeDSL) ✅ already verified -2. q_norm + kv_norm (RMS) — trivial -3. wq_b (CuTeDSL) ✅ already verified -4. RoPE (pure PyTorch reference) -5. FlashMLA attention — THE SUSPECT -6. wo_a BMM (BF16) -7. wo_b (CuTeDSL) ✅ already verified - -Then builds a FAKE attention output (random but reasonable) and runs -steps 6+7 to verify the post-attention path works. - -Usage (on B200): - source /root/nvfp4-megamoe-kernel/tests/.venv/bin/activate - python3 tests/test_attention_path_b200.py -""" - -import sys, os, json, torch, torch.nn.functional as F -from safetensors import safe_open - -REPO = "/root/nvfp4-megamoe-kernel" -sys.path.insert(0, REPO) -MODEL = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4" -DEV = "cuda:0" - -H = 7168; NH = 128; HD = 512; NOPE = 448; ROPE = 64 -QL = 1536; OL = 1024; OG = 16; HPG = NH // OG -EPS = 1e-6 - -E2M1 = torch.tensor([0,.5,1.,1.5,2.,3.,4.,6.,-0,-.5,-1.,-1.5,-2.,-3.,-4.,-6.], dtype=torch.float32) - -_cache = {} -def P(k, wm, md): - if k in _cache: return _cache[k] - with safe_open(os.path.join(md, wm[k]), framework="pt") as f: - t = f.get_tensor(k) - _cache[k] = t - return t - -def dequant(w, sf, gs): - d = w.device; lut = E2M1.to(d) - lo = lut[(w & 0xF).long()]; hi = lut[((w >> 4) & 0xF).long()] - O, I2 = w.shape; I = I2*2 - u = torch.empty(O, I, dtype=torch.float32, device=d) - u[:,0::2] = lo; u[:,1::2] = hi - bs = sf.float().repeat_interleave(16, dim=1)[:O,:I] - return (u * bs * gs).to(torch.bfloat16) - -def rms(x, w, eps=1e-6): - v = x.float().pow(2).mean(-1, keepdim=True) - return (w.float() * (x * torch.rsqrt(v+eps)).float()).to(x.dtype) - -def make_runner(w, sf, gs_t, inf, outf, fused=False, lw=None): - from dsv4.layers.linear import Nvfp4Linear - fp4 = w.view(torch.float4_e2m1fn_x2).permute(1,0).contiguous() - s = sf.to(torch.float8_e4m3fn) if sf.dtype != torch.float8_e4m3fn else sf - s = s.permute(1,0).contiguous() - if fused and gs_t.numel() == 2: - g1,g2 = gs_t[0].item(), gs_t[1].item(); gs = max(g1,g2) - if g1 != g2: - s32 = s.float(); sp = lw[0] if lw else outf//2 - s32[:sp] *= g1/gs; s32[sp:] *= g2/gs; s = s32.to(torch.float8_e4m3fn) - else: - gs = gs_t.max().item() if gs_t.numel() > 1 else gs_t.item() - r = Nvfp4Linear(in_features=inf, out_features=outf, max_num_tokens=8192, device=str(w.device)) - r.fp4 = [fp4]; r.sf = [s]; r.gs = [gs] - r.finalize_weights(); r._ensure_initialized() - return r - -def build_cos_sin(max_pos=4096, rope_dim=ROPE): - half = rope_dim // 2 - inv_freq = 1.0 / (10000.0 ** (torch.arange(0, half, dtype=torch.float32) / half)) - freqs = torch.outer(torch.arange(max_pos, dtype=torch.float32), inv_freq) - return torch.cat([freqs.cos(), freqs.sin()], dim=-1) - - -def main(): - torch.cuda.set_device(0) - torch.manual_seed(42) - - print("=" * 70) - print(" Attention Path Test: Pinpoint FlashMLA Failure") - print("=" * 70) - - with open(os.path.join(MODEL, "model.safetensors.index.json")) as f: - wm = json.load(f)["weight_map"] - G = lambda k: P(k, wm, MODEL).to(DEV) - - p = "model.layers.0"; a = f"{p}.self_attn" - emb = G("model.embed_tokens.weight") - anorm = G(f"{p}.input_layernorm.weight") - qn = G(f"{a}.q_a_norm.weight"); kvn = G(f"{a}.kv_norm.weight") - woa = G(f"{a}.o_a_proj.weight") # (16384, 4096) BF16 - - # Load weights - qa_w = G(f"{a}.q_a_proj.weight"); qa_sf = G(f"{a}.q_a_proj.weight_scale"); qa_gs = G(f"{a}.q_a_proj.weight_scale_2") - qb_w = G(f"{a}.q_b_proj.weight"); qb_sf = G(f"{a}.q_b_proj.weight_scale"); qb_gs = G(f"{a}.q_b_proj.weight_scale_2") - kv_w = G(f"{a}.kv_proj.weight"); kv_sf = G(f"{a}.kv_proj.weight_scale"); kv_gs = G(f"{a}.kv_proj.weight_scale_2") - wob_w = G(f"{a}.o_b_proj.weight"); wob_sf = G(f"{a}.o_b_proj.weight_scale"); wob_gs = G(f"{a}.o_b_proj.weight_scale_2") - - # BF16 references - qa_bf16 = dequant(qa_w, qa_sf, qa_gs.item()) - qb_bf16 = dequant(qb_w, qb_sf, qb_gs.item()) - kv_bf16 = dequant(kv_w, kv_sf, kv_gs.item()) - wob_bf16 = dequant(wob_w, wob_sf, wob_gs.item()) - - # CuTeDSL runners - r_qa = make_runner(qa_w, qa_sf, qa_gs, H, qa_w.shape[0]) - r_qb = make_runner(qb_w, qb_sf, qb_gs, QL, qb_w.shape[0]) - r_kv = make_runner(kv_w, kv_sf, kv_gs, H, kv_w.shape[0]) - r_wob = make_runner(wob_w, wob_sf, wob_gs, OG*OL, wob_w.shape[0]) - - # Input - token_ids = torch.tensor([1, 450, 8403, 315, 5413, 374], dtype=torch.long, device=DEV) - NT = len(token_ids) - with torch.no_grad(): - hidden = emb[token_ids] - normed = rms(hidden, anorm, EPS) - print(f" Input: {NT} tokens, amax={normed.amax():.4f}") - - # ── Step 1: fused_wqa_wkv ───────────────────────────────────────── - print("\n--- Step 1: fused_wqa_wkv (q_a + kv) ---") - with torch.no_grad(): - qa_cute = r_qa.run(normed) - kv_cute = r_kv.run(normed) - qa_ref = normed @ qa_bf16.T - kv_ref = normed @ kv_bf16.T - print(f" q_a CuTeDSL vs BF16: cosine={F.cosine_similarity(qa_cute.flatten().unsqueeze(0).float(), qa_ref.flatten().unsqueeze(0).float()).item():.6f} ✅") - print(f" kv CuTeDSL vs BF16: cosine={F.cosine_similarity(kv_cute.flatten().unsqueeze(0).float(), kv_ref.flatten().unsqueeze(0).float()).item():.6f} ✅") - - # ── Step 2: q_norm + kv_norm ────────────────────────────────────── - print("\n--- Step 2: RMS norm (q_a_norm, kv_norm) ---") - with torch.no_grad(): - qa_normed = rms(qa_cute, qn, EPS) - kv_normed = rms(kv_cute, kvn, EPS) - print(f" q_a normed: amax={qa_normed.amax():.4f} NaN={torch.isnan(qa_normed).any()}") - print(f" kv normed: amax={kv_normed.amax():.4f} NaN={torch.isnan(kv_normed).any()}") - - # ── Step 3: wq_b ────────────────────────────────────────────────── - print("\n--- Step 3: wq_b (q_a → full q) ---") - with torch.no_grad(): - q_cute = r_qb.run(qa_normed) - q_ref = qa_normed @ qb_bf16.T - c = F.cosine_similarity(q_cute.flatten().unsqueeze(0).float(), q_ref.flatten().unsqueeze(0).float()).item() - print(f" q_b CuTeDSL vs BF16: cosine={c:.6f} ✅") - print(f" q shape: {q_cute.shape} → ({NT}, {NH}, {HD})") - - q_3d = q_cute.view(NT, NH, HD) - print(f" q_3d amax: {q_3d.amax():.4f}") - - # ── Step 4: RoPE (reference, GPT-J style) ───────────────────────── - print("\n--- Step 4: RoPE (GPT-J style reference) ---") - cos_sin = build_cos_sin().to(DEV) - positions = torch.arange(NT, dtype=torch.int64, device=DEV) - half_rot = ROPE // 2 - cos_q = cos_sin[positions, :half_rot].unsqueeze(1) # (NT, 1, 32) - sin_q = cos_sin[positions, half_rot:].unsqueeze(1) - - q_nope = q_3d[:, :, :NOPE].clone() - q_rope = q_3d[:, :, NOPE:].clone() - # GPT-J style: interleave even/odd, not split halves - q_even = q_rope[:, :, 0::2].clone() - q_odd = q_rope[:, :, 1::2].clone() - cos_f = cos_q.to(q_3d.dtype) - sin_f = sin_q.to(q_3d.dtype) - q_even_rot = q_even * cos_f - q_odd * sin_f - q_odd_rot = q_even * sin_f + q_odd * cos_f - q_rope_rot = torch.stack([q_even_rot, q_odd_rot], dim=-1).flatten(-2) - q_with_rope = torch.cat([q_nope, q_rope_rot], dim=-1) - print(f" q with RoPE: amax={q_with_rope.amax():.4f}") - - # ── Step 5: Attention (SKIP — use reference) ────────────────────── - print("\n--- Step 5: Attention output ---") - print(" ⚠️ FlashMLA cannot run standalone — using reference implementation") - print(" ⚠️ Running naive scaled dot-product attention in BF16") - - # Naive attention: q @ k.T / sqrt(d) @ v - # We need K, V. K comes from kv (after RoPE), V is kv (nope part) - # Actually in MLA, kv is the latent, not full K/V - # For this test, just use the kv latent directly as a proxy - # and do a simplified attention to get a reasonable output - kv_rope = kv_normed # (NT, HD) — latent representation - print(f" kv latent: shape={kv_normed.shape} amax={kv_normed.amax():.4f}") - - # Simplified: treat q as (NT, NH, HD) and kv as K=V=(NT, HD) - # This isn't the real MLA attention but gives us a non-garbage output - # to test the post-attention path (wo_a, wo_b) - k_simple = kv_normed.unsqueeze(1).expand(-1, NH, -1) # (NT, NH, HD) - v_simple = kv_normed.unsqueeze(1).expand(-1, NH, -1) # (NT, NH, HD) - scale = HD ** -0.5 - attn_weights = torch.matmul(q_with_rope, k_simple.transpose(-1, -2)) * scale - attn_weights = F.softmax(attn_weights.float(), dim=-1).to(torch.bfloat16) - o_ref = torch.matmul(attn_weights, v_simple) - print(f" Naive attention output: amax={o_ref.amax():.4f} NaN={torch.isnan(o_ref).any()}") - - # ── Step 6: wo_a (inverse RoPE + BMM) ───────────────────────────── - print("\n--- Step 6: wo_a (inverse RoPE + BMM) ---") - # Inverse RoPE: same as RoPE but sin -> -sin - o_nope = o_ref[:, :, :NOPE].clone() - o_rope = o_ref[:, :, NOPE:].clone() - o_even = o_rope[:, :, 0::2].clone() - o_odd = o_rope[:, :, 1::2].clone() - o_even_inv = o_even * cos_f + o_odd * sin_f - o_odd_inv = -o_even * sin_f + o_odd * cos_f - o_rope_inv = torch.stack([o_even_inv, o_odd_inv], dim=-1).flatten(-2) - o_inv = torch.cat([o_nope, o_rope_inv], dim=-1) - - # BMM: (OG, NT, HPG*HD) @ (OG, HPG*HD, OL) → (OG, NT, OL) - o_grouped = o_inv.view(NT, OG, HPG * HD).permute(1, 0, 2) - woa_3d = woa.view(OG, OL, HPG * HD) - z = torch.bmm(o_grouped, woa_3d.transpose(1, 2)).permute(1, 0, 2).reshape(NT, OG * OL) - print(f" wo_a z: amax={z.amax():.4f} NaN={torch.isnan(z).any()}") - - # ── Step 7: wo_b (CuTeDSL) ──────────────────────────────────────── - print("\n--- Step 7: wo_b (CuTeDSL vs BF16) ---") - with torch.no_grad(): - wob_cute = r_wob.run(z) - wob_ref = z @ wob_bf16.T - c = F.cosine_similarity(wob_cute.flatten().unsqueeze(0).float(), wob_ref.flatten().unsqueeze(0).float()).item() - print(f" wo_b CuTeDSL vs BF16: cosine={c:.6f} {'✅' if c >= 0.98 else '❌'}") - - # ── Final: attention output through LM head ─────────────────────── - print("\n--- Final: attn output → residual → norm → LM head ---") - fnorm_w = G("model.norm.weight") - lm_head = G("lm_head.weight") - with torch.no_grad(): - x = hidden + wob_cute - x_normed = rms(x, fnorm_w, EPS) - logits = x_normed @ lm_head.T - print(f" logits: amax={logits.amax():.4f} NaN={torch.isnan(logits).any()}") - top5 = torch.topk(logits[-1], 5) - print(f" top5 IDs: {top5.indices.tolist()}") - log_std = logits[-1].float().std().item() - print(f" logit std: {log_std:.4f} {'✅' if 0.5 < log_std < 50 else '❌'}") - - # ── KEY DIAGNOSTIC: What does vLLM's FlashMLA actually do? ──────── - print("\n" + "=" * 70) - print(" KEY DIAGNOSTIC: Check FlashMLA availability on B200") - print("=" * 70) - - try: - import flash_mla - print(f" flash_mla imported: version={getattr(flash_mla, '__version__', 'unknown')}") - # Check if it supports SM100 - cap = torch.cuda.get_device_capability() - print(f" GPU capability: {cap}") - if cap.major >= 10: - print(f" ⚠️ SM{cap.major}{cap.minor} (Blackwell) — FlashMLA may not support this!") - except ImportError: - print(" flash_mla NOT available in this venv (expected — it's in the container)") - - # Check what the vLLM MLA attention calls - print("\n vLLM attention path on B200:") - print(" 1. torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert") - print(" → C++ CUDA kernel for RoPE + KV cache insert") - print(" 2. self.mla_attn(q, kv, positions, output=out)") - print(" → FlashMLA sparse attention") - print(" Both are compiled CUDA kernels that may NOT work on SM100.") - print(" If either returns garbage, the model outputs EOS immediately.") - - -if __name__ == "__main__": - main() diff --git a/tests/archive/test_b_afrag2.py b/tests/archive/test_b_afrag2.py deleted file mode 100644 index 8778e3a9..00000000 --- a/tests/archive/test_b_afrag2.py +++ /dev/null @@ -1,216 +0,0 @@ -"""Stage B: Store P via A-fragment layout with recast C-fragment iterator. - -Matching the backward FMHA pattern exactly: -1. tOrP = pv_thr.make_fragment_A(tP)[None,None,None,0] (A-fragment layout) -2. tdVrP_iter = cute.recast_ptr(tStS.iterator, dtype=BF16) (C-fragment base, recast to BF16) -3. tdVrP = cute.make_tensor(tdVrP_iter + offset, tOrP.layout) -4. make_tmem_copy(St32x32bOp(Repetition(8)), BF16, tdVrP) -5. Store BF16 registers to tdVrP -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - -class StageBAfrag2: - def __init__(self, mma_tiler_mn): - self.qk_acc_dtype = Float32; self.q_dtype = BFloat16; self.o_dtype = BFloat16 - self.c_dtype = BFloat16; self.acc_dtype = Float32 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.num_c_stage = 2; self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - pv_inst_k = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (*self.mma_tiler_mn, pv_inst_k * 4) - self.mma_tiler = self.qk_mma_tiler - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - c_layout = LayoutEnum.ROW_MAJOR; self.c_layout = c_layout - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, c_layout, self.o_dtype) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, c_layout, self.epi_tile, 2) - self.num_ab_stage = 1; self.num_acc_stage = 1 - qk_thr = qk_mma.get_slice(0); qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape); self.s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0); pv_acc_shape = pv_thr.partition_shape_C(self.mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape); self.o_cols = find_tmem_tensor_col_offset(tOtO) - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 0 - self.tmem_o0_offset = self.s_cols * 2 - self.tmem_alloc_cols = 512 - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, 1)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, 1)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)); b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = (cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem)) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, LayoutEnum.from_tensor(a).mma_major_mode(), LayoutEnum.from_tensor(b).mma_major_mode(), self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, LayoutEnum.from_tensor(b).mma_major_mode(), self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)); b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - tma_a, tma_ta = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), a, a_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_b, tma_tb = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), b, b_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - self._kernel(qk_mma, pv_mma, tma_a, tma_ta, tma_b, tma_tb, tma_c, tma_tc, self.cluster_layout_vmnk, self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_a, mA, tma_b, mB, tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()); tidx, _, _ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: cpasync.prefetch_descriptor(tma_a); cpasync.prefetch_descriptor(tma_b); cpasync.prefetch_descriptor(tma_c) - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2]; mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2]; tmem_dealloc: cutlass.Int64; holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - ab_p, ab_c = pipeline.PipelineTmaUmma.create(barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True).make_participants() - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), cta_layout_vmnk=cl_vmnk, defer_sync=True).make_participants() - acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, len(self.epilogue_warp_id)), cta_layout_vmnk=cl_vmnk, defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2, num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=False, two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - sA = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sB = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV_ptr = cute.recast_ptr(sB.iterator, v_smem_s.inner); sV = cute.make_tensor(sV_ptr, v_smem_s.outer) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - gA = cute.local_tile(mA, cute.slice_(self.mma_tiler, (None,0,None)), (None,None,None)) - gB = cute.local_tile(mB, cute.slice_(self.mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gA, mode=[3]) - qk_thr = qk_mma.get_slice(0); tCgA = qk_thr.partition_A(gA); tCgB = qk_thr.partition_B(gB); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsA, tAgA = cpasync.tma_partition(tma_a, 0, a_lay, cute.group_modes(sA,0,3), cute.group_modes(tCgA,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsB, tBgB = cpasync.tma_partition(tma_b, 0, b_lay, cute.group_modes(sB,0,3), cute.group_modes(tCgB,0,3)) - tAgA = tAgA[(None,0,None,0)]; tBgB = tBgB[(None,0,None,0)] - tCrA = qk_mma.make_fragment_A(sA); tCrB = qk_mma.make_fragment_B(sB) - tCrV = pv_mma.make_fragment_B(sV) - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - pv_thr = pv_mma.get_slice(0); pv_acc_shape = pv_thr.partition_shape_C(self.mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - # ── P A-fragment (backward FMHA pattern) ── - # 1. Get A-fragment layout from pv_mma - tP_iter = cute.recast_ptr(tStS.iterator, dtype=self.q_dtype) - tP = cute.make_tensor(tP_iter, p_tmem_s.outer) - tOrP = pv_thr.make_fragment_A(tP)[None, None, None, 0] - # 2. Recast C-fragment iterator to BF16 (matching backward FMHA line 962) - tdVrP_iter = cute.recast_ptr(tStS.iterator, dtype=self.q_dtype) - # 3. Create store target with A-fragment layout + recast iterator - # The offset for P within TMEM: qk_acc_dtype.width / q_dtype.width * tmem_p0_offset - # But since we recast to BF16, the offset should be in BF16 units - tdVrP = cute.make_tensor( - tdVrP_iter + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - # PV MMA's A-fragment (for reading) - tOrP0 = cute.make_tensor(tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.s_cols, tOrP.layout) - # ── TMEM LOAD from C-fragment ── - tmem_ld = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_ld = tcgen05.make_tmem_copy(tmem_ld, tStS0) - sfw = tidx % (32 * len(self.epilogue_warp_id)) - thr_ld = tiled_ld.get_slice(sfw) - tLdS = thr_ld.partition_S(tStS0) - cS_id = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS_id) - tLdcS = thr_ld.partition_D(tScS) - # ── TMEM STORE via A-fragment layout (backward FMHA pattern) ── - tmem_st = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(8)), self.q_dtype) - tiled_st = tcgen05.make_tmem_copy(tmem_st, tdVrP) - thr_st = tiled_st.get_slice(sfw) - tStP = thr_st.partition_D(tdVrP) - # Source identity for store (A-fragment shape) - cS_P = cute.make_identity_tensor((self.qk_mma_tiler[0], self.pv_mma_tiler[2])) - tScS_P = pv_thr.partition_A(cS_P) - tStcS = thr_st.partition_S(tScS_P) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, 1)) - print(f'[A2] tdVrP.layout: {tdVrP.layout}') - print(f'[A2] tOrP0.layout: {tOrP0.layout}') - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - # TMA - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek); cute.copy(tma_a, tAgA[(None,h.count)], tAsA[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_b, tBgB[(None,h.count)], tBsB[(None,h.index)], tma_bar_ptr=h.barrier); peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_b_layout.py b/tests/archive/test_b_layout.py deleted file mode 100644 index 613eff7d..00000000 --- a/tests/archive/test_b_layout.py +++ /dev/null @@ -1,49 +0,0 @@ -"""Test: verify B matrix weight layout with larger dimensions. -Use M=1, N=128, K=256 — big enough for CUTLASS tiles. -Fill B columns with distinct patterns and check if GEMM output -matches the expected column sums. -""" -import torch, sys -sys.path.insert(0, 'src') -from nvfp4_megamoe_kernel.cutlass_nvfp4_gemm.kernel import cutlass_nvfp4_blockscaled_gemm -from nvfp4_megamoe_kernel.nvfp4_mega_moe import _quantize_to_e2m1 - -torch.manual_seed(123) -device = "cuda" - -M, N, K = 1, 128, 256 - -# Create weight with column-dependent pattern -# Column j has value (j % 7) * 0.5 + 0.5 to get distinct E2M1 values -w_bf16 = torch.zeros(K, N, dtype=torch.bfloat16, device=device) -for j in range(N): - w_bf16[:, j] = (j % 7) * 0.5 + 0.5 - -# Quantize -w_fp4, w_sf = _quantize_to_e2m1(w_bf16.T.float()) -w_fp4 = w_fp4.T.contiguous() -w_sf = w_sf.T.contiguous() - -# All-ones A (sum all K elements) -x_bf16 = torch.ones(M, K, dtype=torch.bfloat16, device=device) -x_fp4, x_sf = _quantize_to_e2m1(x_bf16.float()) - -out = cutlass_nvfp4_blockscaled_gemm(x_fp4, x_sf, w_fp4, w_sf, M, N, K, alpha=1.0) - -# Each column j has the same value repeated K times -# So output[j] should be proportional to K * column_value -# Column values cycle: 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 0.5, ... -# So columns with same j%7 should have the same output -print(f"Output first 14: {[f'{v:.2f}' for v in out[0, :14].tolist()]}") -print(f"Expected pattern: columns 0,7 should match; 1,8 should match; etc") - -# Check: columns with same j%7 should be close -for mod_val in range(7): - cols = [j for j in range(N) if j % 7 == mod_val] - vals = out[0, cols] - if len(cols) > 1: - spread = (vals.max() - vals.min()).item() - if spread > 0.5: - print(f"WARNING: j%7={mod_val} spread={spread:.4f} — columns with same weight have different outputs!") - else: - print(f"j%7={mod_val}: mean={vals.mean():.4f} spread={spread:.4f} OK") diff --git a/tests/archive/test_bf16_elemwise.py b/tests/archive/test_bf16_elemwise.py deleted file mode 100644 index 8899958c..00000000 --- a/tests/archive/test_bf16_elemwise.py +++ /dev/null @@ -1,237 +0,0 @@ -"""Absolute minimal: ld FP32 from S0, st FP32 to S1, epi reads S1. -No recast, no BF16, no packing. Pure FP32 copy between TMEM regions.""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - -class BF16Elemwise: - def __init__(self, mma_tiler_mn): - self.qk_acc_dtype = Float32; self.q_dtype = BFloat16; self.o_dtype = BFloat16 - self.c_dtype = BFloat16; self.acc_dtype = Float32 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2; self.use_2cta_instrs = False - self.epilog_sync_bar_id = 1 - - def _setup(self, qk_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - self.mma_tiler = self.qk_mma_tiler - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - c_layout = LayoutEnum.ROW_MAJOR; self.c_layout = c_layout - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, False, c_layout, self.o_dtype) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, c_layout, self.epi_tile, 2) - self.num_ab_stage = 1; self.num_acc_stage = 1 - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - self.s_cols = find_tmem_tensor_col_offset(tStS) - self.tmem_alloc_cols = self.s_cols * 2 - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, - LayoutEnum.from_tensor(a).mma_major_mode(), - LayoutEnum.from_tensor(b).mma_major_mode(), - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, - tcgen05.OperandSource.SMEM) - self._setup(qk_mma) - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - tma_a, tma_ta = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - a, a_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_b, tma_tb = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - b, b_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - self._kernel(qk_mma, tma_a, tma_ta, tma_b, tma_tb, tma_c, tma_tc, - self.cluster_layout_vmnk, self.a_smem_s, self.b_smem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, tma_a, mA, tma_b, mB, tma_c, mC, cl_vmnk, - a_smem_s, b_smem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_a); cpasync.prefetch_descriptor(tma_b); cpasync.prefetch_descriptor(tma_c) - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, len(self.epilogue_warp_id)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=False, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - sA = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sB = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - gA = cute.local_tile(mA, cute.slice_(self.mma_tiler, (None,0,None)), (None,None,None)) - gB = cute.local_tile(mB, cute.slice_(self.mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gA, mode=[3]) - qk_thr = qk_mma.get_slice(0) - tCgA = qk_thr.partition_A(gA); tCgB = qk_thr.partition_B(gB); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsA, tAgA = cpasync.tma_partition(tma_a, 0, a_lay, cute.group_modes(sA,0,3), cute.group_modes(tCgA,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsB, tBgB = cpasync.tma_partition(tma_b, 0, b_lay, cute.group_modes(sB,0,3), cute.group_modes(tCgB,0,3)) - tAgA = tAgA[(None,0,None,0)]; tBgB = tBgB[(None,0,None,0)] - tCrA = qk_mma.make_fragment_A(sA); tCrB = qk_mma.make_fragment_B(sB) - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator, tStS.layout) - tStS1 = cute.make_tensor(tStS.iterator + self.s_cols, tStS.layout) - - # LD and ST on same layout - tmem_ld = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tmem_st = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_ld = tcgen05.make_tmem_copy(tmem_ld, tStS0) - tiled_st = tcgen05.make_tmem_copy(tmem_st, tStS1) - sfw = tidx % (32 * len(self.epilogue_warp_id)) - thr_ld = tiled_ld.get_slice(sfw) - thr_st = tiled_st.get_slice(sfw) - tLdS = thr_ld.partition_S(tStS0) - tStS = thr_st.partition_D(tStS1) - cS_id = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS_id) - tLdcS = thr_ld.partition_D(tScS) - tStcS = thr_st.partition_S(tScS) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, 1)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_a, tAgA[(None,h.count)], tAsA[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_b, tBgB[(None,h.count)], tBsB[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_bf16_pack.py b/tests/archive/test_bf16_pack.py deleted file mode 100644 index e12596ac..00000000 --- a/tests/archive/test_bf16_pack.py +++ /dev/null @@ -1,237 +0,0 @@ -"""Absolute minimal: ld FP32 from S0, st FP32 to S1, epi reads S1. -No recast, no BF16, no packing. Pure FP32 copy between TMEM regions.""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - -class BF16PackTest: - def __init__(self, mma_tiler_mn): - self.qk_acc_dtype = Float32; self.q_dtype = BFloat16; self.o_dtype = BFloat16 - self.c_dtype = BFloat16; self.acc_dtype = Float32 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2; self.use_2cta_instrs = False - self.epilog_sync_bar_id = 1 - - def _setup(self, qk_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - self.mma_tiler = self.qk_mma_tiler - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - c_layout = LayoutEnum.ROW_MAJOR; self.c_layout = c_layout - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, False, c_layout, self.o_dtype) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, c_layout, self.epi_tile, 2) - self.num_ab_stage = 1; self.num_acc_stage = 1 - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - self.s_cols = find_tmem_tensor_col_offset(tStS) - self.tmem_alloc_cols = self.s_cols * 2 - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, - LayoutEnum.from_tensor(a).mma_major_mode(), - LayoutEnum.from_tensor(b).mma_major_mode(), - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, - tcgen05.OperandSource.SMEM) - self._setup(qk_mma) - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - tma_a, tma_ta = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - a, a_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_b, tma_tb = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - b, b_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - self._kernel(qk_mma, tma_a, tma_ta, tma_b, tma_tb, tma_c, tma_tc, - self.cluster_layout_vmnk, self.a_smem_s, self.b_smem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, tma_a, mA, tma_b, mB, tma_c, mC, cl_vmnk, - a_smem_s, b_smem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_a); cpasync.prefetch_descriptor(tma_b); cpasync.prefetch_descriptor(tma_c) - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, len(self.epilogue_warp_id)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=False, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - sA = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sB = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - gA = cute.local_tile(mA, cute.slice_(self.mma_tiler, (None,0,None)), (None,None,None)) - gB = cute.local_tile(mB, cute.slice_(self.mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gA, mode=[3]) - qk_thr = qk_mma.get_slice(0) - tCgA = qk_thr.partition_A(gA); tCgB = qk_thr.partition_B(gB); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsA, tAgA = cpasync.tma_partition(tma_a, 0, a_lay, cute.group_modes(sA,0,3), cute.group_modes(tCgA,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsB, tBgB = cpasync.tma_partition(tma_b, 0, b_lay, cute.group_modes(sB,0,3), cute.group_modes(tCgB,0,3)) - tAgA = tAgA[(None,0,None,0)]; tBgB = tBgB[(None,0,None,0)] - tCrA = qk_mma.make_fragment_A(sA); tCrB = qk_mma.make_fragment_B(sB) - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator, tStS.layout) - tStS1 = cute.make_tensor(tStS.iterator + self.s_cols, tStS.layout) - - # LD and ST on same layout - tmem_ld = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tmem_st = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_ld = tcgen05.make_tmem_copy(tmem_ld, tStS0) - tiled_st = tcgen05.make_tmem_copy(tmem_st, tStS1) - sfw = tidx % (32 * len(self.epilogue_warp_id)) - thr_ld = tiled_ld.get_slice(sfw) - thr_st = tiled_st.get_slice(sfw) - tLdS = thr_ld.partition_S(tStS0) - tStS = thr_st.partition_D(tStS1) - cS_id = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS_id) - tLdcS = thr_ld.partition_D(tScS) - tStcS = thr_st.partition_S(tScS) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, 1)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_a, tAgA[(None,h.count)], tAsA[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_b, tBgB[(None,h.count)], tBsB[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_bf16_recast_full.py b/tests/archive/test_bf16_recast_full.py deleted file mode 100644 index 1ca6ce8e..00000000 --- a/tests/archive/test_bf16_recast_full.py +++ /dev/null @@ -1,242 +0,0 @@ -"""Test BF16 recast pattern with FULL C-fragment layout (no subview). -ld from S0 (full 128x128), recast BF16, st to S1 (full 128x128 at offset 128). -Since both ld and st use the same layout, the recast should work (shapes match). -Then epi reads S1. If this works, the recast pattern IS correct for same-layout cases. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - -class BF16RecastFull: - def __init__(self, mma_tiler_mn): - self.qk_acc_dtype = Float32; self.q_dtype = BFloat16; self.o_dtype = BFloat16 - self.c_dtype = BFloat16; self.acc_dtype = Float32 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2; self.use_2cta_instrs = False - self.epilog_sync_bar_id = 1 - - def _setup(self, qk_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - self.mma_tiler = self.qk_mma_tiler - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - c_layout = LayoutEnum.ROW_MAJOR; self.c_layout = c_layout - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, False, c_layout, self.o_dtype) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, c_layout, self.epi_tile, 2) - self.num_ab_stage = 1; self.num_acc_stage = 1 - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - self.s_cols = find_tmem_tensor_col_offset(tStS) - self.tmem_alloc_cols = self.s_cols * 2 - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, - LayoutEnum.from_tensor(a).mma_major_mode(), - LayoutEnum.from_tensor(b).mma_major_mode(), - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, - tcgen05.OperandSource.SMEM) - self._setup(qk_mma) - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - tma_a, tma_ta = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - a, a_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_b, tma_tb = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - b, b_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - self._kernel(qk_mma, tma_a, tma_ta, tma_b, tma_tb, tma_c, tma_tc, - self.cluster_layout_vmnk, self.a_smem_s, self.b_smem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, tma_a, mA, tma_b, mB, tma_c, mC, cl_vmnk, - a_smem_s, b_smem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_a); cpasync.prefetch_descriptor(tma_b); cpasync.prefetch_descriptor(tma_c) - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, len(self.epilogue_warp_id)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=False, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - sA = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sB = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - gA = cute.local_tile(mA, cute.slice_(self.mma_tiler, (None,0,None)), (None,None,None)) - gB = cute.local_tile(mB, cute.slice_(self.mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gA, mode=[3]) - qk_thr = qk_mma.get_slice(0) - tCgA = qk_thr.partition_A(gA); tCgB = qk_thr.partition_B(gB); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsA, tAgA = cpasync.tma_partition(tma_a, 0, a_lay, cute.group_modes(sA,0,3), cute.group_modes(tCgA,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsB, tBgB = cpasync.tma_partition(tma_b, 0, b_lay, cute.group_modes(sB,0,3), cute.group_modes(tCgB,0,3)) - tAgA = tAgA[(None,0,None,0)]; tBgB = tBgB[(None,0,None,0)] - tCrA = qk_mma.make_fragment_A(sA); tCrB = qk_mma.make_fragment_B(sB) - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator, tStS.layout) - tStS1 = cute.make_tensor(tStS.iterator + self.s_cols, tStS.layout) - - tmem_ld = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tmem_st = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_ld = tcgen05.make_tmem_copy(tmem_ld, tStS0) - tiled_st = tcgen05.make_tmem_copy(tmem_st, tStS1) - sfw = tidx % (32 * len(self.epilogue_warp_id)) - thr_ld = tiled_ld.get_slice(sfw) - thr_st = tiled_st.get_slice(sfw) - tLdS = thr_ld.partition_S(tStS0) - tStS1_dst = thr_st.partition_D(tStS1) - cS_id = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS_id) - tLdcS = thr_ld.partition_D(tScS) - tStcS = thr_st.partition_S(tScS) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, 1)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_a, tAgA[(None,h.count)], tAsA[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_b, tBgB[(None,h.count)], tBsB[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_bf16_recast_simple.py b/tests/archive/test_bf16_recast_simple.py deleted file mode 100644 index b6588c30..00000000 --- a/tests/archive/test_bf16_recast_simple.py +++ /dev/null @@ -1,239 +0,0 @@ -"""Simplest BF16 recast test: ld FP32 from S0, use recast to convert, -st back to S0, then epi reads S0. Single region, no subview.""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - -class SimpleBF16Recast: - def __init__(self, mma_tiler_mn): - self.qk_acc_dtype = Float32; self.q_dtype = BFloat16; self.o_dtype = BFloat16 - self.c_dtype = BFloat16; self.acc_dtype = Float32 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2; self.use_2cta_instrs = False - self.epilog_sync_bar_id = 1 - - def _setup(self, qk_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - self.mma_tiler = self.qk_mma_tiler - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - c_layout = LayoutEnum.ROW_MAJOR; self.c_layout = c_layout - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, False, c_layout, self.o_dtype) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, c_layout, self.epi_tile, 2) - self.num_ab_stage = 1; self.num_acc_stage = 1 - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - self.s_cols = find_tmem_tensor_col_offset(tStS) - self.tmem_alloc_cols = self.s_cols - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, - LayoutEnum.from_tensor(a).mma_major_mode(), - LayoutEnum.from_tensor(b).mma_major_mode(), - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, - tcgen05.OperandSource.SMEM) - self._setup(qk_mma) - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - tma_a, tma_ta = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - a, a_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_b, tma_tb = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - b, b_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - self._kernel(qk_mma, tma_a, tma_ta, tma_b, tma_tb, tma_c, tma_tc, - self.cluster_layout_vmnk, self.a_smem_s, self.b_smem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, tma_a, mA, tma_b, mB, tma_c, mC, cl_vmnk, - a_smem_s, b_smem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_a); cpasync.prefetch_descriptor(tma_b); cpasync.prefetch_descriptor(tma_c) - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, len(self.epilogue_warp_id)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=False, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - sA = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sB = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - gA = cute.local_tile(mA, cute.slice_(self.mma_tiler, (None,0,None)), (None,None,None)) - gB = cute.local_tile(mB, cute.slice_(self.mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gA, mode=[3]) - qk_thr = qk_mma.get_slice(0) - tCgA = qk_thr.partition_A(gA); tCgB = qk_thr.partition_B(gB); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsA, tAgA = cpasync.tma_partition(tma_a, 0, a_lay, cute.group_modes(sA,0,3), cute.group_modes(tCgA,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsB, tBgB = cpasync.tma_partition(tma_b, 0, b_lay, cute.group_modes(sB,0,3), cute.group_modes(tCgB,0,3)) - tAgA = tAgA[(None,0,None,0)]; tBgB = tBgB[(None,0,None,0)] - tCrA = qk_mma.make_fragment_A(sA); tCrB = qk_mma.make_fragment_B(sB) - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator, tStS.layout) - - # ld and st on the SAME tensor (S0), same layout - tmem_ld = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tmem_st = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_ld = tcgen05.make_tmem_copy(tmem_ld, tStS0) - tiled_st = tcgen05.make_tmem_copy(tmem_st, tStS0) - sfw = tidx % (32 * len(self.epilogue_warp_id)) - thr_ld = tiled_ld.get_slice(sfw) - thr_st = tiled_st.get_slice(sfw) - tLdS = thr_ld.partition_S(tStS0) - tStS_dst = thr_st.partition_D(tStS0) - cS_id = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS_id) - tLdcS = thr_ld.partition_D(tScS) - tStcS = thr_st.partition_S(tScS) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, 1)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_a, tAgA[(None,h.count)], tAsA[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_b, tBgB[(None,h.count)], tBsB[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_blackwell_attn_b200.py b/tests/archive/test_blackwell_attn_b200.py deleted file mode 100644 index 41394760..00000000 --- a/tests/archive/test_blackwell_attn_b200.py +++ /dev/null @@ -1,318 +0,0 @@ -#!/usr/bin/env python3 -""" -DeepSeek-V4 Blackwell Attention — Full Pipeline Test - -Tests the cutedsl.blackwell_attention module with real weights: -1. Prefill: process N tokens, write KV to paged cache -2. Decode: process 1 new token, read ALL cached KV, attend -3. Verify decode output matches BF16 reference - -This is the core of the fix for the vLLM Blackwell garbage output bug. - -Usage (on B200): - cd /root/nvfp4-megamoe-kernel - PYTHONPATH=/root/nvfp4-megamoe-kernel tests/venv/bin/python tests/test_blackwell_attn_b200.py -""" - -import sys, os, json, torch, torch.nn.functional as F, math -from safetensors import safe_open - -REPO = "/root/nvfp4-megamoe-kernel" -sys.path.insert(0, REPO) -MODEL = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4" -DEV = "cuda:0" - -H = 7168; NH = 128; HD = 512; NOPE = 448; ROPE = 64 -QL = 1536; OL = 1024; OG = 16; HPG = NH // OG -EPS = 1e-6; WINDOW = 128; SCALE = HD ** -0.5 - -E2M1 = torch.tensor([0,.5,1.,1.5,2.,3.,4.,6.,-0,-.5,-1.,-1.5,-2.,-3.,-4.,-6.], dtype=torch.float32) - -_cache = {} -def P(k, wm, md): - if k in _cache: return _cache[k] - with safe_open(os.path.join(md, wm[k]), framework="pt") as f: - t = f.get_tensor(k) - _cache[k] = t - return t - -def dequant(w, sf, gs): - d = w.device; lut = E2M1.to(d) - lo = lut[(w & 0xF).long()]; hi = lut[((w >> 4) & 0xF).long()] - O, I2 = w.shape; I = I2*2 - u = torch.empty(O, I, dtype=torch.float32, device=d) - u[:,0::2] = lo; u[:,1::2] = hi - bs = sf.float().repeat_interleave(16, dim=1)[:O,:I] - return (u * bs * gs).to(torch.bfloat16) - -def rms(x, w, eps=1e-6): - v = x.float().pow(2).mean(-1, keepdim=True) - return (w.float() * (x * torch.rsqrt(v+eps)).float()).to(x.dtype) - -def make_runner(w, sf, gs_t, inf, outf, fused=False, lw=None): - from dsv4.layers.linear import Nvfp4Linear - fp4 = w.view(torch.float4_e2m1fn_x2).permute(1,0).contiguous() - s = sf.to(torch.float8_e4m3fn) if sf.dtype != torch.float8_e4m3fn else sf - s = s.permute(1,0).contiguous() - if fused and gs_t.numel() == 2: - g1,g2 = gs_t[0].item(), gs_t[1].item(); gs = max(g1,g2) - if g1 != g2: - s32 = s.float(); sp = lw[0] if lw else outf//2 - s32[:sp] *= g1/gs; s32[sp:] *= g2/gs; s = s32.to(torch.float8_e4m3fn) - else: - gs = gs_t.max().item() if gs_t.numel() > 1 else gs_t.item() - r = Nvfp4Linear(in_features=inf, out_features=outf, max_num_tokens=8192, device=str(w.device)) - r.fp4 = [fp4]; r.sf = [s]; r.gs = [gs] - r.finalize_weights(); r._ensure_initialized() - return r - -def build_cos_sin(max_pos=4096, rope_dim=ROPE): - half = rope_dim // 2 - inv_freq = 1.0 / (10000.0 ** (torch.arange(0, half, dtype=torch.float32) / half)) - freqs = torch.outer(torch.arange(max_pos, dtype=torch.float32), inv_freq) - return torch.cat([freqs.cos(), freqs.sin(), freqs.cos(), freqs.sin()], dim=-1) # extra for safety - -# Only use the first rope_dim cols -def build_cos_sin(max_pos=4096, rope_dim=ROPE): - half = rope_dim // 2 - inv_freq = 1.0 / (10000.0 ** (torch.arange(0, half, dtype=torch.float32) / half)) - freqs = torch.outer(torch.arange(max_pos, dtype=torch.float32), inv_freq) - return torch.cat([freqs.cos(), freqs.sin()], dim=-1) - - -def test_blackwell_attention(layer_id, compress_ratio): - """Test the full blackwell attention pipeline for a specific layer.""" - from dsv4.reference.attention import ( - apply_gptj_rope, apply_inv_gptj_rope, - blackwell_attention_forward, - kv_quantize_fp8, kv_dequantize_fp8, - paged_kv_write, paged_kv_read, - causal_prefill_attention, decode_attention, - ) - - torch.cuda.set_device(0) - torch.manual_seed(42) - torch.cuda.empty_cache() - - with open(os.path.join(MODEL, "model.safetensors.index.json")) as f: - wm = json.load(f)["weight_map"] - G = lambda k: P(k, wm, MODEL).to(DEV) - - p = f"model.layers.{layer_id}"; a = f"{p}.self_attn" - layer_type = "SWA" if compress_ratio <= 1 else f"CSA(c={compress_ratio})" - - print(f"\n{'='*70}") - print(f" Layer {layer_id} — {layer_type} — Blackwell Attention Test") - print(f"{'='*70}") - - # Load weights - emb = G("model.embed_tokens.weight") - anorm = G(f"{p}.input_layernorm.weight") - qn = G(f"{a}.q_a_norm.weight"); kvn = G(f"{a}.kv_norm.weight") - woa = G(f"{a}.o_a_proj.weight") - sinks = G(f"{a}.sinks") - - qa_w = G(f"{a}.q_a_proj.weight"); qa_sf = G(f"{a}.q_a_proj.weight_scale"); qa_gs = G(f"{a}.q_a_proj.weight_scale_2") - qb_w = G(f"{a}.q_b_proj.weight"); qb_sf = G(f"{a}.q_b_proj.weight_scale"); qb_gs = G(f"{a}.q_b_proj.weight_scale_2") - kv_w = G(f"{a}.kv_proj.weight"); kv_sf = G(f"{a}.kv_proj.weight_scale"); kv_gs = G(f"{a}.kv_proj.weight_scale_2") - wob_w = G(f"{a}.o_b_proj.weight"); wob_sf = G(f"{a}.o_b_proj.weight_scale"); wob_gs = G(f"{a}.o_b_proj.weight_scale_2") - - r_qa = make_runner(qa_w, qa_sf, qa_gs, H, qa_w.shape[0]) - r_qb = make_runner(qb_w, qb_sf, qb_gs, QL, qb_w.shape[0]) - r_kv = make_runner(kv_w, kv_sf, kv_gs, H, kv_w.shape[0]) - r_wob = make_runner(wob_w, wob_sf, wob_gs, OG*OL, wob_w.shape[0]) - - cos_sin = build_cos_sin(max_pos=4096).to(DEV) - - # ── Test 1: Prefill-only attention ──────────────────────────────── - print(f"\n --- Test 1: Prefill attention (8 tokens) ---") - N = 8 - token_ids = torch.tensor([1, 450, 8403, 315, 5413, 374, 2198, 643], dtype=torch.long, device=DEV) - positions = torch.arange(N, dtype=torch.int64, device=DEV) - - with torch.no_grad(): - hidden = emb[token_ids] - normed = rms(hidden, anorm, EPS) - - qa = r_qa.run(normed) - kv = r_kv.run(normed) - qa_n = rms(qa, qn, EPS) - kv_n = rms(kv, kvn, EPS) - q = r_qb.run(qa_n).view(N, NH, HD) - - q_rope = apply_gptj_rope(q, positions, cos_sin, NOPE, ROPE) - kv_rope = apply_gptj_rope(kv_n.unsqueeze(1), positions, cos_sin, NOPE, ROPE).squeeze(1) - - # Causal attention - o_prefill = causal_prefill_attention(q_rope, kv_rope, SCALE) - print(f" Prefill attention output: amax={o_prefill.amax():.4f} NaN={torch.isnan(o_prefill).any()}") - - # BF16 reference (same computation, different path) - q_t = q_rope.permute(1, 0, 2) - kv_exp = kv_rope.unsqueeze(0).expand(NH, -1, -1) - o_ref = F.scaled_dot_product_attention(q_t, kv_exp, kv_exp, is_causal=True, scale=SCALE).permute(1, 0, 2) - c = F.cosine_similarity(o_prefill.flatten().unsqueeze(0).float(), o_ref.flatten().unsqueeze(0).float()).item() - print(f" Prefill vs SDPA reference cosine: {c:.6f} {'✅' if c>=0.999 else '❌'}") - - # ── Test 2: Decode attention with KV cache ──────────────────────── - print(f"\n --- Test 2: Decode attention (1 token, 8 cached) ---") - - block_size = 256 - num_blocks = 64 - kv_cache_fp8 = torch.zeros(num_blocks, block_size, HD, dtype=torch.float8_e4m3fn, device=DEV) - inv_scale_cache = torch.zeros(num_blocks * block_size, 1, dtype=torch.bfloat16, device=DEV) - - with torch.no_grad(): - # Write prefill KV to cache - kv_fp8, inv_s = kv_quantize_fp8(kv_rope) - prefill_slots = positions - paged_kv_write(kv_fp8, prefill_slots, kv_cache_fp8, block_size) - for t in range(N): - inv_scale_cache[prefill_slots[t]] = inv_s[t] - - # Decode: token at position 8 - decode_id = torch.tensor([991], dtype=torch.long, device=DEV) - decode_pos = torch.tensor([N], dtype=torch.int64, device=DEV) - - hidden_d = emb[decode_id] - normed_d = rms(hidden_d, anorm, EPS) - qa_d = r_qa.run(normed_d) - kv_d = r_kv.run(normed_d) - qa_n_d = rms(qa_d, qn, EPS) - kv_n_d = rms(kv_d, kvn, EPS) - q_d = r_qb.run(qa_n_d).view(1, NH, HD) - q_rope_d = apply_gptj_rope(q_d, decode_pos, cos_sin, NOPE, ROPE) - kv_rope_d = apply_gptj_rope(kv_n_d.unsqueeze(1), decode_pos, cos_sin, NOPE, ROPE).squeeze(1) - - # Write decode KV to cache - kv_fp8_d, inv_s_d = kv_quantize_fp8(kv_rope_d) - paged_kv_write(kv_fp8_d, decode_pos, kv_cache_fp8, block_size) - inv_scale_cache[decode_pos[0]] = inv_s_d[0] - - # Read ALL 9 tokens from cache - all_slots = torch.arange(N + 1, dtype=torch.int64, device=DEV) - kv_cached_fp8 = paged_kv_read(all_slots, kv_cache_fp8, block_size, N + 1, HD) - kv_cached = kv_dequantize_fp8(kv_cached_fp8, inv_scale_cache[all_slots]) - - # Decode attention: 1 query vs 9 cached KVs - o_decode = decode_attention(q_rope_d, kv_cached, SCALE) - print(f" Decode attention output: amax={o_decode.amax():.4f} NaN={torch.isnan(o_decode).any()}") - - # BF16 reference: process all 9 tokens at once - all_ids = torch.cat([token_ids, decode_id]) - all_pos = torch.arange(N + 1, dtype=torch.int64, device=DEV) - hidden_all = emb[all_ids] - normed_all = rms(hidden_all, anorm, EPS) - qa_all = r_qa.run(normed_all) - kv_all = r_kv.run(normed_all) - qa_n_all = rms(qa_all, qn, EPS) - kv_n_all = rms(kv_all, kvn, EPS) - q_all = r_qb.run(qa_n_all).view(N + 1, NH, HD) - q_rope_all = apply_gptj_rope(q_all, all_pos, cos_sin, NOPE, ROPE) - kv_rope_all = apply_gptj_rope(kv_n_all.unsqueeze(1), all_pos, cos_sin, NOPE, ROPE).squeeze(1) - - o_ref_all = causal_prefill_attention(q_rope_all, kv_rope_all, SCALE) - o_ref_decode = o_ref_all[N:] # Only the decode token - - c = F.cosine_similarity(o_decode.flatten().unsqueeze(0).float(), o_ref_decode.flatten().unsqueeze(0).float()).item() - print(f" Decode vs BF16 reference cosine: {c:.6f} {'✅' if c>=0.98 else '❌'}") - - # ── Test 3: Full output pipeline (inverse RoPE + o_a + o_b) ────── - print(f"\n --- Test 3: Full output pipeline ---") - with torch.no_grad(): - # Using decode attention output - o_inv = apply_inv_gptj_rope(o_decode, decode_pos, cos_sin, NOPE, ROPE) - o_grouped = o_inv.view(1, OG, HPG * HD).permute(1, 0, 2) - woa_3d = woa.view(OG, OL, HPG * HD) - z_cached = torch.bmm(o_grouped, woa_3d.transpose(1, 2)).permute(1, 0, 2).reshape(1, OG * OL) - attn_out_cached = r_wob.run(z_cached) - - # Using BF16 reference - o_inv_ref = apply_inv_gptj_rope(o_ref_decode, decode_pos, cos_sin, NOPE, ROPE) - o_grouped_ref = o_inv_ref.view(1, OG, HPG * HD).permute(1, 0, 2) - z_ref = torch.bmm(o_grouped_ref, woa_3d.transpose(1, 2)).permute(1, 0, 2).reshape(1, OG * OL) - attn_out_ref = r_wob.run(z_ref) - - c_full = F.cosine_similarity(attn_out_cached.flatten().unsqueeze(0).float(), attn_out_ref.flatten().unsqueeze(0).float()).item() - print(f" Full pipeline cosine: {c_full:.6f} {'✅' if c_full>=0.98 else '❌'}") - print(f" Output amax: cached={attn_out_cached.amax():.4f} ref={attn_out_ref.amax():.4f}") - - # ── Test 4: Multi-step decode (3 decode steps) ─────────────────── - print(f"\n --- Test 4: Multi-step decode (3 steps) ---") - decode_ids = torch.tensor([991, 1502, 4200], dtype=torch.long, device=DEV) - - with torch.no_grad(): - cosines = [] - for step in range(3): - pos = N + step - dpos = torch.tensor([pos], dtype=torch.int64, device=DEV) - d_id = decode_ids[step:step+1] - - hidden_s = emb[d_id] - normed_s = rms(hidden_s, anorm, EPS) - qa_s = r_qa.run(normed_s) - kv_s = r_kv.run(normed_s) - qa_n_s = rms(qa_s, qn, EPS) - kv_n_s = rms(kv_s, kvn, EPS) - q_s = r_qb.run(qa_n_s).view(1, NH, HD) - q_rope_s = apply_gptj_rope(q_s, dpos, cos_sin, NOPE, ROPE) - kv_rope_s = apply_gptj_rope(kv_n_s.unsqueeze(1), dpos, cos_sin, NOPE, ROPE).squeeze(1) - - # Write to cache - kv_fp8_s, inv_s_s = kv_quantize_fp8(kv_rope_s) - paged_kv_write(kv_fp8_s, dpos, kv_cache_fp8, block_size) - inv_scale_cache[dpos[0]] = inv_s_s[0] - - # Read all cached KV - all_s = torch.arange(pos + 1, dtype=torch.int64, device=DEV) - kv_all_fp8 = paged_kv_read(all_s, kv_cache_fp8, block_size, pos + 1, HD) - kv_all_dequant = kv_dequantize_fp8(kv_all_fp8, inv_scale_cache[all_s]) - - # Decode attention - o_s = decode_attention(q_rope_s, kv_all_dequant, SCALE) - - # BF16 reference - all_ids_ref = torch.cat([token_ids, decode_ids[:step+1]]) - all_pos_ref = torch.arange(pos + 1, dtype=torch.int64, device=DEV) - hidden_ref = emb[all_ids_ref] - normed_ref = rms(hidden_ref, anorm, EPS) - qa_ref = r_qa.run(normed_ref) - kv_ref = r_kv.run(normed_ref) - qa_n_ref = rms(qa_ref, qn, EPS) - kv_n_ref = rms(kv_ref, kvn, EPS) - q_ref = r_qb.run(qa_n_ref).view(pos + 1, NH, HD) - q_rope_ref = apply_gptj_rope(q_ref, all_pos_ref, cos_sin, NOPE, ROPE) - kv_rope_ref = apply_gptj_rope(kv_n_ref.unsqueeze(1), all_pos_ref, cos_sin, NOPE, ROPE).squeeze(1) - o_ref_full = causal_prefill_attention(q_rope_ref, kv_rope_ref, SCALE) - o_ref_last = o_ref_full[-1:] - - c = F.cosine_similarity(o_s.flatten().unsqueeze(0).float(), o_ref_last.flatten().unsqueeze(0).float()).item() - cosines.append(c) - print(f" Step {step} (pos={pos}, {pos+1} cached): cosine = {c:.6f} {'✅' if c>=0.98 else '❌'}") - - # Cleanup - del r_qa, r_qb, r_kv, r_wob - torch.cuda.empty_cache() - - return c_full, cosines - - -def main(): - print("=" * 70) - print(" DeepSeek-V4 Blackwell Attention Pipeline Test") - print(" Tests cutedsl.blackwell_attention with real weights") - print("=" * 70) - - # Test SWA layer (layer 60, compress_ratio=0) - c_swa, cosines_swa = test_blackwell_attention(60, 0) - - print(f"\n{'='*70}") - print(f" SUMMARY") - print(f" Layer 60 (SWA):") - print(f" Full pipeline cosine: {c_swa:.6f}") - print(f" Multi-step decode: {', '.join(f'{c:.6f}' for c in cosines_swa)}") - print(f"{'='*70}") - - -if __name__ == "__main__": - main() diff --git a/tests/archive/test_cache.py b/tests/archive/test_cache.py deleted file mode 100644 index 56170e47..00000000 --- a/tests/archive/test_cache.py +++ /dev/null @@ -1,252 +0,0 @@ -"""Tests for KV cache: schema, allocator, pools, manager lifecycle.""" - -import torch -import pytest - -from dsv4.model.config import DSV4Config -from dsv4.model.layer_schedule import build_schedule, AttentionType, RouterMode, LayerSpec -from dsv4.cache.schema import build_schema, compute_block_budget, BLOCK_SIZE_ORIGINAL_TOKENS -from dsv4.cache.allocator import BlockAllocator -from dsv4.cache.paged_cache import PagedKVPool -from dsv4.cache.state_cache import StateCachePool -from dsv4.cache.manager import KVCacheManager - - -# ---- Schema tests ---- - -def test_csa_schema(): - config = DSV4Config.pro() - spec = LayerSpec(layer_idx=2, attn=AttentionType.CSA, - ffn=__import__('dsv4.model.layer_schedule', fromlist=['FFNType']).FFNType.MOE, - router_mode=RouterMode.HASH) - schema = build_schema(config, spec) - assert schema.entries_per_block == 32 # 128 / 4 - assert schema.indexer_entries_per_block == 32 - assert schema.tail_buffer_size == 3 # m - 1 - assert schema.swa_window_size == 128 - - -def test_hca_schema(): - config = DSV4Config.pro() - spec = LayerSpec(layer_idx=3, attn=AttentionType.HCA, - ffn=__import__('dsv4.model.layer_schedule', fromlist=['FFNType']).FFNType.MOE, - router_mode=RouterMode.DENSE) - schema = build_schema(config, spec) - assert schema.entries_per_block == 1 # 128 / 128 - assert schema.indexer_entries_per_block == 0 - assert schema.tail_buffer_size == 127 # m' - 1 - - -def test_swa_schema(): - config = DSV4Config.flash() - spec = LayerSpec(layer_idx=0, attn=AttentionType.SWA, - ffn=__import__('dsv4.model.layer_schedule', fromlist=['FFNType']).FFNType.MOE, - router_mode=RouterMode.HASH) - schema = build_schema(config, spec) - assert schema.entries_per_block == 0 - assert schema.indexer_entries_per_block == 0 - assert schema.tail_buffer_size == 0 - assert schema.swa_window_size == 128 - - -def test_schema_from_schedule(): - """Every layer in a full schedule produces a valid schema.""" - config = DSV4Config.flash() - schedule = build_schedule(config) - for spec in schedule: - schema = build_schema(config, spec) - assert schema.swa_window_size > 0 - assert schema.entry_head_dim == config.head_dim - if spec.attn == AttentionType.SWA: - assert schema.entries_per_block == 0 - else: - assert schema.entries_per_block > 0 - - -def test_block_budget(): - config = DSV4Config.pro() - schedule = build_schedule(config) - budget = compute_block_budget(config, schedule, 1_000_000, 16) - assert "csa" in budget - assert "hca" in budget - assert budget["csa"] > budget["hca"] # CSA uses more blocks per request - - -# ---- Allocator tests ---- - -def test_acquire_release_roundtrip(): - alloc = BlockAllocator(num_total_blocks=1024) - a = alloc.acquire(10) - assert alloc.num_free == 1014 - b = alloc.acquire(5) - assert alloc.num_free == 1009 - alloc.release(a) - assert alloc.num_free == 1019 - alloc.release(b) - assert alloc.num_free == 1024 - # Re-acquire works after release. - c = alloc.acquire(20) - assert alloc.num_free == 1004 - - -def test_oom_raises(): - alloc = BlockAllocator(num_total_blocks=4) - alloc.acquire(4) - with pytest.raises(RuntimeError, match="OOM"): - alloc.acquire(1) - - -def test_acquire_returns_unique_ids(): - alloc = BlockAllocator(num_total_blocks=100) - a = alloc.acquire(50) - b = alloc.acquire(50) - assert len(torch.intersect1d(a, b)) == 0 - - -# ---- Pool shape tests ---- - -def test_paged_pool_shapes_csa(): - from dsv4.model.layer_schedule import FFNType - config = DSV4Config.pro() - spec = LayerSpec(layer_idx=2, attn=AttentionType.CSA, - ffn=FFNType.MOE, router_mode=RouterMode.HASH) - schema = build_schema(config, spec) - pool = PagedKVPool(schema, num_blocks=16) - assert pool.entries_fp8.shape == (16, 32, config.head_dim - config.rope_dim) - assert pool.entries_rope.shape == (16, 32, config.rope_dim) - assert pool.inv_scale.shape == (16, 32) - assert pool.indexer_keys_fp4 is not None - assert pool.indexer_keys_fp4.shape[1] == 32 - - -def test_paged_pool_shapes_hca(): - from dsv4.model.layer_schedule import FFNType - config = DSV4Config.pro() - spec = LayerSpec(layer_idx=3, attn=AttentionType.HCA, - ffn=FFNType.MOE, router_mode=RouterMode.DENSE) - schema = build_schema(config, spec) - pool = PagedKVPool(schema, num_blocks=256) - assert pool.entries_fp8.shape == (256, 1, config.head_dim - config.rope_dim) - assert pool.entries_rope.shape == (256, 1, config.rope_dim) - assert pool.indexer_keys_fp4 is None - - -def test_state_pool_shapes_csa(): - from dsv4.model.layer_schedule import FFNType - config = DSV4Config.pro() - spec = LayerSpec(layer_idx=2, attn=AttentionType.CSA, - ffn=FFNType.MOE, router_mode=RouterMode.HASH) - schema = build_schema(config, spec) - pool = StateCachePool(schema, max_requests=8) - assert pool.swa_fp8.shape == (8, 128, config.head_dim - config.rope_dim) - assert pool.swa_rope.shape == (8, 128, config.rope_dim) - assert pool.tail_ka is not None - assert pool.tail_kb is not None # CSA has both streams - assert pool.tail_len is not None - - -def test_state_pool_shapes_hca(): - from dsv4.model.layer_schedule import FFNType - config = DSV4Config.pro() - spec = LayerSpec(layer_idx=3, attn=AttentionType.HCA, - ffn=FFNType.MOE, router_mode=RouterMode.DENSE) - schema = build_schema(config, spec) - pool = StateCachePool(schema, max_requests=8) - assert pool.tail_ka is not None - assert pool.tail_kb is None # HCA only one stream - assert pool.tail_za is not None - assert pool.tail_len is not None - - -def test_state_pool_shapes_swa(): - from dsv4.model.layer_schedule import FFNType - config = DSV4Config.flash() - spec = LayerSpec(layer_idx=0, attn=AttentionType.SWA, - ffn=FFNType.MOE, router_mode=RouterMode.HASH) - schema = build_schema(config, spec) - pool = StateCachePool(schema, max_requests=8) - assert pool.swa_fp8.shape == (8, 128, config.head_dim - config.rope_dim) - assert pool.tail_ka is None # No tail for SWA-only - assert pool.tail_len is None - - -# ---- Manager lifecycle tests ---- - -def test_admit_release_recycles_slot(): - config = DSV4Config.flash() - schedule = build_schedule(config) - mgr = KVCacheManager(config, schedule, max_concurrent_requests=4, - num_blocks_per_csa_layer=64, num_blocks_per_hca_layer=64) - s1 = mgr.admit_request() - mgr.release_request(s1) - s2 = mgr.admit_request() - assert s1 == s2 # slot was recycled - - -def test_admit_exhaustion(): - config = DSV4Config.flash() - schedule = build_schedule(config) - mgr = KVCacheManager(config, schedule, max_concurrent_requests=2, - num_blocks_per_csa_layer=64, num_blocks_per_hca_layer=64) - mgr.admit_request() - mgr.admit_request() - with pytest.raises(RuntimeError, match="concurrent"): - mgr.admit_request() - - -def test_handle_construction_no_alloc(): - """acquire() should not allocate GPU memory — critical for cudagraph.""" - config = DSV4Config.flash() - schedule = build_schedule(config) - mgr = KVCacheManager(config, schedule, max_concurrent_requests=4, - num_blocks_per_csa_layer=64, num_blocks_per_hca_layer=64) - slot = mgr.admit_request() - torch.cuda.synchronize() - before = torch.cuda.memory_allocated() - handle = mgr.acquire( - layer_idx=0, - request_slots=torch.tensor([slot], dtype=torch.int32, device="cuda"), - positions=torch.tensor([0], dtype=torch.int32, device="cuda"), - request_ids=torch.tensor([0], dtype=torch.int32, device="cuda"), - ) - torch.cuda.synchronize() - after = torch.cuda.memory_allocated() - # Allow small variance from tensor view creation, but no large alloc - assert after - before < 1024, f"acquire() allocated {after - before} bytes — breaks cudagraph" - assert handle.paged is None # layer 0 is SWA - mgr.release_request(slot) - - -def test_manager_memory_tracking(): - config = DSV4Config.flash() - schedule = build_schedule(config) - mgr = KVCacheManager(config, schedule, max_concurrent_requests=4, - num_blocks_per_csa_layer=64, num_blocks_per_hca_layer=64) - total = mgr.memory_bytes() - assert total > 0 - # Rough sanity: should be in the MB range for this config - assert total > 1_000_000 # at least 1 MB - assert total < 10_000_000_000 # less than 10 GB - - -def test_full_flash_stack_construction(): - """Construct manager with all 43 Flash layers — pools for every layer.""" - config = DSV4Config.flash() - schedule = build_schedule(config) - mgr = KVCacheManager(config, schedule, max_concurrent_requests=4, - num_blocks_per_csa_layer=64, num_blocks_per_hca_layer=64) - # 2 SWA layers (no paged pool) + 41 compressed layers - assert len(mgr.paged_pools) == 43 - assert mgr.paged_pools[0] is None # SWA - assert mgr.paged_pools[1] is None # SWA - assert mgr.paged_pools[2] is not None # CSA - assert mgr.paged_pools[3] is not None # HCA - - # All state pools present - assert len(mgr.state_pools) == 43 - for i in range(43): - assert mgr.state_pools[i] is not None - - -if __name__ == "__main__": - pytest.main([__file__, "-v"]) diff --git a/tests/archive/test_compile_custom_op.py b/tests/archive/test_compile_custom_op.py deleted file mode 100644 index 46b033d1..00000000 --- a/tests/archive/test_compile_custom_op.py +++ /dev/null @@ -1,189 +0,0 @@ -#!/usr/bin/env python3 -"""Test torch.compile + CuTeDSL NVFP4 runner via custom_op. - -Critical test: does torch.compile (fullgraph mode) accept the -nvfp4::moe_gemm custom op and produce a working compiled graph? - -Run on the B200: - docker run --rm --gpus all --entrypoint python3 \ - -v /root/nvfp4-megamoe-kernel:/root/nvfp4-megamoe-kernel \ - -v /root/nvidia-meeting:/root/nvidia-meeting:ro \ - nvfp4-megamoe-kernel-vllm:latest \ - /root/nvfp4-megamoe-kernel/tests/test_compile_custom_op.py -""" -import os -import sys -import json -import glob -import torch -from safetensors import safe_open - -REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) -sys.path.insert(0, REPO_ROOT) - -from dsv4.layers.moe import Nvfp4MoE -from dsv4.ops.custom_ops import register_runner, nvfp4_moe_gemm - -NVFP4_MODEL_DIR = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4" -DEVICE = "cuda" - - -def find_shards(model_dir): - index_path = os.path.join(model_dir, "model.safetensors.index.json") - key_to_shard = {} - if os.path.exists(index_path): - with open(index_path) as f: - index = json.load(f) - for key, shard in index["weight_map"].items(): - key_to_shard[key] = os.path.join(model_dir, shard) - else: - for sf in glob.glob(os.path.join(model_dir, "*.safetensors")): - with safe_open(sf, framework="pt") as f: - for key in f.keys(): - key_to_shard[key] = sf - return key_to_shard - - -def load_layer_tensors(model_dir, layer_idx): - key_to_shard = find_shards(model_dir) - layer_prefix = f"layers.{layer_idx}." - shard_to_keys = {} - for key, shard in key_to_shard.items(): - norm_key = key.removeprefix("model.") - if not norm_key.startswith(layer_prefix): - continue - shard_to_keys.setdefault(shard, []).append((key, norm_key)) - tensors = {} - for shard, keys in shard_to_keys.items(): - with safe_open(shard, framework="pt") as f: - for orig_key, norm_key in keys: - tensors[norm_key] = f.get_tensor(orig_key) - return tensors - - -def prepare_nvfp4_weights_direct(nvfp4_tensors, layer_idx, expert_indices, intermediate_size): -from dsv4.ops.quantize import ( - quantize_activation_nvfp4, - quantize_weight_to_nvfp4, -) - l1_fp4, l1_sf, l1_gs = [], [], [] - l2_fp4, l2_sf, l2_gs = [], [], [] - - for e in expert_indices: - gate_w = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.gate_proj.weight"].to(DEVICE) - up_w = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.up_proj.weight"].to(DEVICE) - gate_sf = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.gate_proj.weight_scale"].to(DEVICE) - up_sf = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.up_proj.weight_scale"].to(DEVICE) - gate_gs = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.gate_proj.weight_scale_2"].item() - up_gs = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.up_proj.weight_scale_2"].item() - - fused_w = torch.cat([gate_w, up_w], dim=0) - fused_w_fp4 = fused_w.view(torch.float4_e2m1fn_x2).permute(1, 0).contiguous() - fused_sf = torch.cat([gate_sf, up_sf], dim=0).permute(1, 0).contiguous() - l1_max_gs = max(gate_gs, up_gs) - if gate_gs != up_gs: - fused_sf_f32 = fused_sf.float() - fused_sf_f32[:, :intermediate_size] *= (gate_gs / l1_max_gs) - fused_sf_f32[:, intermediate_size:] *= (up_gs / l1_max_gs) - fused_sf = fused_sf_f32.to(torch.float8_e4m3fn) - l1_fp4.append(fused_w_fp4) - l1_sf.append(fused_sf) - l1_gs.append(l1_max_gs) - - down_key = f"layers.{layer_idx}.mlp.experts.{e}.down_proj.weight" - if down_key in nvfp4_tensors: - down_w = nvfp4_tensors[down_key].to(DEVICE) - down_sf = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.down_proj.weight_scale"].to(DEVICE) - down_gs = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.down_proj.weight_scale_2"].item() - l2_fp4.append(down_w.view(torch.float4_e2m1fn_x2).permute(1, 0).contiguous()) - l2_sf.append(down_sf.permute(1, 0).contiguous()) - l2_gs.append(down_gs) - - return { - 'l1_fp4': l1_fp4, 'l1_sf': l1_sf, 'l1_gs': l1_gs, - 'l2_fp4': l2_fp4, 'l2_sf': l2_sf, 'l2_gs': l2_gs, - } - - -def main(): - torch.manual_seed(42) - expert_indices = [0, 1, 2] - hidden_size = 7168 - intermediate_size = 3072 - - print("=" * 70) - print(" torch.compile + CuTeDSL Custom Op Test") - print("=" * 70) - - # Load weights - nvfp4_tensors = load_layer_tensors(NVFP4_MODEL_DIR, 0) - weights = prepare_nvfp4_weights_direct(nvfp4_tensors, 0, expert_indices, intermediate_size) - - # Create runner - runner = Nvfp4MoE( - num_experts=len(expert_indices), - hidden_size=hidden_size, - intermediate_size=intermediate_size, - max_num_tokens=8, - top_k=2, - device="cuda", - ) - runner.prepare_weights_direct( - weights['l1_fp4'], weights['l1_sf'], weights['l1_gs'], - weights['l2_fp4'], weights['l2_sf'], weights['l2_gs'], - ) - runner_id = register_runner(runner) - - # Test input - hidden_states = torch.randn(4, hidden_size, dtype=torch.bfloat16, device=DEVICE) * 2.0 - topk_ids = torch.tensor([[0, 1]] * 4, dtype=torch.int32, device=DEVICE) - topk_weights = torch.tensor([[0.6, 0.4]] * 4, dtype=torch.float32, device=DEVICE) - - # 1. Warmup: compute activation global scales - print("\n[0] Computing activation global scales (warmup)...") - runner.compute_activation_global_scales(hidden_states, topk_weights, topk_ids) - print(f" L1 gs: {runner._l1_activation_global_scale:.6f}") - print(f" L2 gs: {runner._l2_activation_global_scale:.6f}") - - # 1. Eager mode (baseline) - print("\n[1/2] Running eager mode (baseline)...") - runner._ensure_stacked() - eager_out = nvfp4_moe_gemm(hidden_states, topk_weights, topk_ids, runner_id, hidden_size) - print(f" Eager output: amax={eager_out.abs().max():.4f} mean={eager_out.float().mean():.6f}") - - # 2. torch.compile fullgraph - print("\n[2/2] Running torch.compile(fullgraph=True)...") - try: - @torch.compile(fullgraph=True) - def compiled_fn(hs, tw, ti): - return nvfp4_moe_gemm(hs, tw, ti, runner_id, hidden_size) - - compiled_out = compiled_fn(hidden_states, topk_weights, topk_ids) - print(f" Compiled output: amax={compiled_out.abs().max():.4f} mean={compiled_out.float().mean():.6f}") - - # Compare - if eager_out.shape == compiled_out.shape: - cos = torch.nn.functional.cosine_similarity( - eager_out.flatten().unsqueeze(0).float(), - compiled_out.flatten().unsqueeze(0).float(), - ).item() - print(f"\n Eager vs Compiled: cosine={cos:.6f}") - if cos > 0.99: - print(" ✅ torch.compile produces matching output!") - else: - print(f" ⚠️ Cosine {cos:.4f} < 0.99 — check for numerical issues") - else: - print(f" ❌ Shape mismatch: eager={eager_out.shape} compiled={compiled_out.shape}") - except Exception as e: - print(f" ❌ torch.compile FAILED: {type(e).__name__}: {e}") - import traceback - traceback.print_exc() - sys.exit(1) - - print("\n" + "=" * 70) - print(" Test complete ✅") - print("=" * 70) - - -if __name__ == "__main__": - main() diff --git a/tests/archive/test_csa_attention_b200.py b/tests/archive/test_csa_attention_b200.py deleted file mode 100644 index 6f4f6935..00000000 --- a/tests/archive/test_csa_attention_b200.py +++ /dev/null @@ -1,251 +0,0 @@ -#!/usr/bin/env python3 -""" -Test CSA/HCA attention kernel with real model weights. - -Runs the full attention path for layer 0 (C128A): -1. q_a_proj, kv_proj (CuTeDSL NVFP4) -2. q_norm, kv_norm (RMS) -3. q_b_proj (CuTeDSL NVFP4) -4. RoPE (BF16 reference) -5. CSA sparse attention (our kernel using PyTorch SDPA) -6. wo_a BMM + wo_b (BF16 + CuTeDSL NVFP4) -7. Compare against full BF16 reference - -Usage (on B200): - source /root/nvfp4-megamoe-kernel/tests/.venv/bin/activate - python3 tests/test_csa_attention_b200.py -""" - -import sys, os, json, torch, torch.nn.functional as F -from safetensors import safe_open - -REPO = "/root/nvfp4-megamoe-kernel" -sys.path.insert(0, REPO) -MODEL = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4" -DEV = "cuda:0" - -H = 7168; NH = 128; HD = 512; NOPE = 448; ROPE = 64 -QL = 1536; OL = 1024; OG = 16; HPG = NH // OG -EPS = 1e-6; WINDOW = 8192; SCALE = HD ** -0.5 - -E2M1 = torch.tensor([0,.5,1.,1.5,2.,3.,4.,6.,-0,-.5,-1.,-1.5,-2.,-3.,-4.,-6.], dtype=torch.float32) - -_cache = {} -def P(k, wm, md): - if k in _cache: return _cache[k] - with safe_open(os.path.join(md, wm[k]), framework="pt") as f: - t = f.get_tensor(k) - _cache[k] = t - return t - -def dequant(w, sf, gs): - d = w.device; lut = E2M1.to(d) - lo = lut[(w & 0xF).long()]; hi = lut[((w >> 4) & 0xF).long()] - O, I2 = w.shape; I = I2*2 - u = torch.empty(O, I, dtype=torch.float32, device=d) - u[:,0::2] = lo; u[:,1::2] = hi - bs = sf.float().repeat_interleave(16, dim=1)[:O,:I] - return (u * bs * gs).to(torch.bfloat16) - -def rms(x, w, eps=1e-6): - v = x.float().pow(2).mean(-1, keepdim=True) - return (w.float() * (x * torch.rsqrt(v+eps)).float()).to(x.dtype) - -def make_runner(w, sf, gs_t, inf, outf, fused=False, lw=None): - from dsv4.layers.linear import Nvfp4Linear - fp4 = w.view(torch.float4_e2m1fn_x2).permute(1,0).contiguous() - s = sf.to(torch.float8_e4m3fn) if sf.dtype != torch.float8_e4m3fn else sf - s = s.permute(1,0).contiguous() - if fused and gs_t.numel() == 2: - g1,g2 = gs_t[0].item(), gs_t[1].item(); gs = max(g1,g2) - if g1 != g2: - s32 = s.float(); sp = lw[0] if lw else outf//2 - s32[:sp] *= g1/gs; s32[sp:] *= g2/gs; s = s32.to(torch.float8_e4m3fn) - else: - gs = gs_t.max().item() if gs_t.numel() > 1 else gs_t.item() - r = Nvfp4Linear(in_features=inf, out_features=outf, max_num_tokens=8192, device=str(w.device)) - r.fp4 = [fp4]; r.sf = [s]; r.gs = [gs] - r.finalize_weights(); r._ensure_initialized() - return r - - -def apply_gptj_rope(x, positions, cos_sin, nope, rope): - """GPT-J style RoPE (interleaved). Applied to last `rope` dims of x.""" - if rope == 0 or x.numel() == 0: - return x - half = rope // 2 - cos = cos_sin[positions, :half].to(x.dtype) # (T, half) or (T, 1, half) - sin = cos_sin[positions, half:].to(x.dtype) - - if x.dim() == 3: - cos = cos.unsqueeze(1) # (T, 1, half) - sin = sin.unsqueeze(1) - x_rope = x[..., nope:].clone() - even = x_rope[..., 0::2] - odd = x_rope[..., 1::2] - out = x.clone() - out[..., nope:][..., 0::2] = even * cos - odd * sin - out[..., nope:][..., 1::2] = even * sin + odd * cos - return out - - -def build_cos_sin(max_pos=4096, rope_dim=ROPE): - half = rope_dim // 2 - inv_freq = 1.0 / (10000.0 ** (torch.arange(0, half, dtype=torch.float32) / half)) - freqs = torch.outer(torch.arange(max_pos, dtype=torch.float32), inv_freq) - return torch.cat([freqs.cos(), freqs.sin()], dim=-1) - - -def main(): - torch.cuda.set_device(0) - torch.manual_seed(42) - - print("=" * 70) - print(" CSA/HCA Attention Kernel Test (Layer 0, C128A)") - print("=" * 70) - - with open(os.path.join(MODEL, "model.safetensors.index.json")) as f: - wm = json.load(f)["weight_map"] - G = lambda k: P(k, wm, MODEL).to(DEV) - - p = "model.layers.0"; a = f"{p}.self_attn" - - # Load weights - emb = G("model.embed_tokens.weight") - anorm = G(f"{p}.input_layernorm.weight") - qn = G(f"{a}.q_a_norm.weight"); kvn = G(f"{a}.kv_norm.weight") - woa = G(f"{a}.o_a_proj.weight") # (16384, 4096) BF16 - - qa_w = G(f"{a}.q_a_proj.weight"); qa_sf = G(f"{a}.q_a_proj.weight_scale"); qa_gs = G(f"{a}.q_a_proj.weight_scale_2") - qb_w = G(f"{a}.q_b_proj.weight"); qb_sf = G(f"{a}.q_b_proj.weight_scale"); qb_gs = G(f"{a}.q_b_proj.weight_scale_2") - kv_w = G(f"{a}.kv_proj.weight"); kv_sf = G(f"{a}.kv_proj.weight_scale"); kv_gs = G(f"{a}.kv_proj.weight_scale_2") - wob_w = G(f"{a}.o_b_proj.weight"); wob_sf = G(f"{a}.o_b_proj.weight_scale"); wob_gs = G(f"{a}.o_b_proj.weight_scale_2") - sinks = G(f"{a}.sinks") - - # BF16 references - qa_bf16 = dequant(qa_w, qa_sf, qa_gs.item()) - qb_bf16 = dequant(qb_w, qb_sf, qb_gs.item()) - kv_bf16 = dequant(kv_w, kv_sf, kv_gs.item()) - wob_bf16 = dequant(wob_w, wob_sf, wob_gs.item()) - - # CuTeDSL runners - r_qa = make_runner(qa_w, qa_sf, qa_gs, H, qa_w.shape[0]) - r_qb = make_runner(qb_w, qb_sf, qb_gs, QL, qb_w.shape[0]) - r_kv = make_runner(kv_w, kv_sf, kv_gs, H, kv_w.shape[0]) - r_wob = make_runner(wob_w, wob_sf, wob_gs, OG*OL, wob_w.shape[0]) - - # Input - token_ids = torch.tensor([1, 450, 8403, 315, 5413, 374], dtype=torch.long, device=DEV) - NT = len(token_ids) - cos_sin = build_cos_sin(max_pos=WINDOW + 256).to(DEV) - positions = torch.arange(NT, dtype=torch.int64, device=DEV) - - print(f" Input: {NT} tokens") - print(f" attn_sink: shape={sinks.shape} values={sinks.flatten()[:8].tolist()}") - - with torch.no_grad(): - hidden = emb[token_ids] - normed = rms(hidden, anorm, EPS) - - # ── Step 1: q_a + kv projections ────────────────────────────── - qa_cute = r_qa.run(normed) - kv_cute = r_kv.run(normed) - qa_ref = normed @ qa_bf16.T - kv_ref = normed @ kv_bf16.T - - # ── Step 2: RMS norm ────────────────────────────────────────── - qa_n = rms(qa_cute, qn, EPS) - kv_n = rms(kv_cute, kvn, EPS) - - # ── Step 3: q_b ─────────────────────────────────────────────── - q_cute = r_qb.run(qa_n).view(NT, NH, HD) - - # ── Step 4: RoPE on Q ───────────────────────────────────────── - q_rope = apply_gptj_rope(q_cute, positions, cos_sin, NOPE, ROPE) - - # ── Step 5: KV insert (simulated — just keep kv_n) ──────────── - # In production, kv_n would be written to the SWA KV cache (FP8) - # and the compressor would write to the state cache - # For this test, we use kv_n directly as the KV for attention - - # ── Step 6: FULL ATTENTION (PyTorch SDPA, works on Blackwell) ── - from dsv4.reference.csa_attention import full_attention_reference - - o_attn = full_attention_reference(q_rope, kv_n, scale=SCALE) - print(f" Attention output: amax={o_attn.amax():.4f} NaN={torch.isnan(o_attn).any()}") - - # ── Step 7: wo_a (inverse RoPE + BMM) ───────────────────────── - # Inverse RoPE: same as forward RoPE but sin → -sin - o_inv = apply_gptj_rope(o_attn, positions, cos_sin, NOPE, ROPE) - # Actually inverse RoPE negates sin, so: - # Let me re-do with correct inverse - half = ROPE // 2 - cos_f = cos_sin[positions, :half].unsqueeze(1).to(o_attn.dtype) - sin_f = cos_sin[positions, half:].unsqueeze(1).to(o_attn.dtype) - o_nope = o_attn[:, :, :NOPE].clone() - o_rope = o_attn[:, :, NOPE:].clone() - o_even = o_rope[:, :, 0::2].clone() - o_odd = o_rope[:, :, 1::2].clone() - # Inverse: even' = even*cos + odd*sin, odd' = -even*sin + odd*cos - o_even_inv = o_even * cos_f + o_odd * sin_f - o_odd_inv = -o_even * sin_f + o_odd * cos_f - o_inv = torch.cat([o_nope, torch.stack([o_even_inv, o_odd_inv], -1).flatten(-2)], dim=-1) - - # BMM - o_grouped = o_inv.view(NT, OG, HPG * HD).permute(1, 0, 2) - woa_3d = woa.view(OG, OL, HPG * HD) - z = torch.bmm(o_grouped, woa_3d.transpose(1, 2)).permute(1, 0, 2).reshape(NT, OG * OL) - - # ── Step 8: wo_b ────────────────────────────────────────────── - attn_out = r_wob.run(z) - attn_ref = z @ wob_bf16.T - c = F.cosine_similarity(attn_out.flatten().unsqueeze(0).float(), attn_ref.flatten().unsqueeze(0).float()).item() - print(f" wo_b cosine: {c:.6f} {'✅' if c>=0.98 else '❌'}") - - # ── Full forward: attention output → residual → LM head ─────────── - print("\n--- Full forward: attn → residual → norm → LM head ---") - fnorm_w = G("model.norm.weight") - lm_head = G("lm_head.weight") - with torch.no_grad(): - x = hidden + attn_out - x_normed = rms(x, fnorm_w, EPS) - logits = x_normed @ lm_head.T - print(f" logits: amax={logits.amax():.4f}") - top5 = torch.topk(logits[-1], 5) - print(f" top5 IDs: {top5.indices.tolist()}") - log_std = logits[-1].float().std().item() - print(f" logit std: {log_std:.4f} {'✅' if 0.5 < log_std < 50 else '❌'}") - - # ── Compare: BF16 full path vs CuTeDSL + SDPA ──────────────────── - print("\n--- Compare: Full BF16 path vs CuTeDSL + SDPA ---") - with torch.no_grad(): - qa_bf = normed @ qa_bf16.T - kv_bf = normed @ kv_bf16.T - qa_n_bf = rms(qa_bf, qn, EPS) - kv_n_bf = rms(kv_bf, kvn, EPS) - q_bf = (qa_n_bf @ qb_bf16.T).view(NT, NH, HD) - q_rope_bf = apply_gptj_rope(q_bf, positions, cos_sin, NOPE, ROPE) - o_bf = full_attention_reference(q_rope_bf, kv_n_bf, scale=SCALE) - # wo_a BMM - o_nope_bf = o_bf[:, :, :NOPE].clone() - o_rope_bf = o_bf[:, :, NOPE:].clone() - o_even_bf = o_rope_bf[:, :, 0::2].clone() - o_odd_bf = o_rope_bf[:, :, 1::2].clone() - o_even_inv_bf = o_even_bf * cos_f + o_odd_bf * sin_f - o_odd_inv_bf = -o_even_bf * sin_f + o_odd_bf * cos_f - o_inv_bf = torch.cat([o_nope_bf, torch.stack([o_even_inv_bf, o_odd_inv_bf], -1).flatten(-2)], dim=-1) - o_grouped_bf = o_inv_bf.view(NT, OG, HPG * HD).permute(1, 0, 2) - z_bf = torch.bmm(o_grouped_bf, woa_3d.transpose(1, 2)).permute(1, 0, 2).reshape(NT, OG * OL) - attn_bf = z_bf @ wob_bf16.T - - c = F.cosine_similarity(attn_out.flatten().unsqueeze(0).float(), attn_bf.flatten().unsqueeze(0).float()).item() - print(f" Full path CuTeDSL vs BF16 cosine: {c:.6f} {'✅' if c>=0.95 else '❌'}") - - print("\n" + "=" * 70) - print(" SUMMARY: All attention components work with PyTorch SDPA.") - print(" Next: integrate into vLLM to replace broken FlashMLA kernel.") - print("=" * 70) - - -if __name__ == "__main__": - main() diff --git a/tests/archive/test_csa_sparse_attn_b200.py b/tests/archive/test_csa_sparse_attn_b200.py deleted file mode 100644 index d7afd83f..00000000 --- a/tests/archive/test_csa_sparse_attn_b200.py +++ /dev/null @@ -1,399 +0,0 @@ -#!/usr/bin/env python3 -""" -CSA Sparse Attention Test - -Tests the csa_sparse_attention_batched function with simulated compressor output: -1. Create compressed KV cache (simulating compressor output) -2. Create topk_indices (simulating indexer output) -3. Do sparse attention on compressed KV at topk positions -4. Do SWA attention on the window -5. Merge with sink weights -6. Compare against full attention reference - -Usage (on B200): - cd /root/nvfp4-megamoe-kernel - PYTHONPATH=/root/nvfp4-megamoe-kernel tests/venv/bin/python tests/test_csa_sparse_attn_b200.py -""" - -import sys, os, json, torch, torch.nn.functional as F, time -from safetensors import safe_open - -REPO = "/root/nvfp4-megamoe-kernel" -sys.path.insert(0, REPO) -MODEL = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4" -DEV = "cuda:0" - -H = 7168; NH = 128; HD = 512; NOPE = 448; ROPE = 64 -QL = 1536; OL = 1024; OG = 16; HPG = NH // OG -EPS = 1e-6; WINDOW = 128; SCALE = HD ** -0.5 - -E2M1 = torch.tensor([0,.5,1.,1.5,2.,3.,4.,6.,-0,-.5,-1.,-1.5,-2.,-3.,-4.,6.], dtype=torch.float32) - -_cache = {} -def P(k, wm, md): - if k in _cache: return _cache[k] - with safe_open(os.path.join(md, wm[k]), framework="pt") as f: - t = f.get_tensor(k) - _cache[k] = t - return t - -def rms(x, w, eps=1e-6): - v = x.float().pow(2).mean(-1, keepdim=True) - return (w.float() * (x * torch.rsqrt(v+eps)).float()).to(x.dtype) - -def make_runner(w, sf, gs_t, inf, outf, fused=False, lw=None): - from dsv4.layers.linear import Nvfp4Linear - fp4 = w.view(torch.float4_e2m1fn_x2).permute(1,0).contiguous() - s = sf.to(torch.float8_e4m3fn) if sf.dtype != torch.float8_e4m3fn else sf - s = s.permute(1,0).contiguous() - if fused and gs_t.numel() == 2: - g1,g2 = gs_t[0].item(), gs_t[1].item(); gs = max(g1,g2) - if g1 != g2: - s32 = s.float(); sp = lw[0] if lw else outf//2 - s32[:sp] *= g1/gs; s32[sp:] *= g2/gs; s = s32.to(torch.float8_e4m3fn) - else: - gs = gs_t.max().item() if gs_t.numel() > 1 else gs_t.item() - r = Nvfp4Linear(in_features=inf, out_features=outf, max_num_tokens=8192, device=str(w.device)) - r.fp4 = [fp4]; r.sf = [s]; r.gs = [gs] - r.finalize_weights(); r._ensure_initialized() - return r - -def build_cos_sin(max_pos=4096, rope_dim=ROPE): - half = rope_dim // 2 - inv_freq = 1.0 / (10000.0 ** (torch.arange(0, half, dtype=torch.float32) / half)) - freqs = torch.outer(torch.arange(max_pos, dtype=torch.float32), inv_freq) - return torch.cat([freqs.cos(), freqs.sin()], dim=-1) - -def apply_gptj_rope(x, positions, cos_sin, nope_dim, rope_dim): - if rope_dim == 0 or x.numel() == 0: return x - half = rope_dim // 2 - cos = cos_sin[positions, :half].to(x.dtype) - sin = cos_sin[positions, half:2*half].to(x.dtype) - if x.dim() == 3: cos = cos.unsqueeze(1); sin = sin.unsqueeze(1) - x_rope = x[..., nope_dim:].clone() - even = x_rope[..., 0::2]; odd = x_rope[..., 1::2] - out = x.clone() - out[..., nope_dim:][..., 0::2] = even * cos - odd * sin - out[..., nope_dim:][..., 1::2] = even * sin + odd * cos - return out - -def apply_inv_gptj_rope(x, positions, cos_sin, nope_dim, rope_dim): - if rope_dim == 0 or x.numel() == 0: return x - half = rope_dim // 2 - cos = cos_sin[positions, :half].to(x.dtype) - sin = cos_sin[positions, half:2*half].to(x.dtype) - if x.dim() == 3: cos = cos.unsqueeze(1); sin = sin.unsqueeze(1) - x_rope = x[..., nope_dim:].clone() - even = x_rope[..., 0::2]; odd = x_rope[..., 1::2] - out = x.clone() - out[..., nope_dim:][..., 0::2] = even * cos + odd * sin - out[..., nope_dim:][..., 1::2] = -even * sin + odd * cos - return out - -def kv_quantize_fp8(kv_bf16): - amax = kv_bf16.float().abs().amax(dim=-1, keepdim=True).clamp(min=1e-12) - fp8_max = torch.tensor(448.0, dtype=torch.float32, device=kv_bf16.device) - scale = fp8_max / amax - kv_fp8 = (kv_bf16.float() * scale).to(torch.float8_e4m3fn) - inv_scale = (amax / fp8_max).to(torch.bfloat16) - return kv_fp8, inv_scale - -def kv_dequantize_fp8(kv_fp8, inv_scale): - return (kv_fp8.to(torch.bfloat16) * inv_scale).to(torch.bfloat16) - -def causal_prefill_attention(q, kv, scale): - T, NH, HD = q.shape - q_t = q.permute(1, 0, 2) - kv_exp = kv.unsqueeze(0).expand(NH, -1, -1) - out = F.scaled_dot_product_attention(q_t, kv_exp, kv_exp, is_causal=True, scale=scale) - return out.permute(1, 0, 2) - - -def csa_sparse_gather_attention(q, compressed_kv, topk_indices, topk_lens, scale, cos_sin, nope_dim, rope_dim): - """CSA sparse attention: gather compressed KV at topk positions, attend. - - q: (T, NH, HD) with RoPE already applied - compressed_kv: (num_compressed, HD) — all compressed KV vectors - topk_indices: (T, num_topk) — which compressed positions to attend to - topk_lens: (T,) — how many of the topk_indices are valid - """ - T, NH, HD = q.shape - device = q.device - num_topk = topk_indices.shape[-1] - - # Gather compressed KV at topk positions - # Clamp to valid range - safe_idx = topk_indices.clamp(min=0, max=compressed_kv.shape[0] - 1) - # (T, num_topk, HD) - k_gathered = compressed_kv[safe_idx] - - # Mask invalid positions (set to 0) - valid_mask = torch.arange(num_topk, device=device).unsqueeze(0) < topk_lens.unsqueeze(1) - k_gathered = k_gathered * valid_mask.unsqueeze(-1).to(k_gathered.dtype) - - # Apply RoPE to gathered K at their original (compressed) positions - if rope_dim > 0 and cos_sin is not None: - kv_positions = safe_idx # The positions in the compressed cache - # BUT: compressed position i represents the i-th group of compress_ratio tokens - # The "position" for RoPE should be the original token position, not the compressed index - # For now, use the compressed index as a proxy (this is a simplification) - # In the real pipeline, the compressor stores KV with RoPE already applied - pass # Skip RoPE for now — the compressor already applies it - - # Multi-head attention: expand K for all heads - # k_gathered: (T, num_topk, HD) → (T, NH, num_topk, HD) - k_heads = k_gathered.unsqueeze(1).expand(-1, NH, -1, -1) - v_heads = k_heads.clone() - - # Q: (T, NH, HD) → (T*NH, 1, HD) - q_2d = q.reshape(T * NH, 1, HD) - k_2d = k_heads.reshape(T * NH, num_topk, HD) - v_2d = v_heads.reshape(T * NH, num_topk, HD) - - # Attention mask: (T, num_topk) → (T*NH, 1, num_topk) - attn_mask = valid_mask.unsqueeze(1).expand(-1, NH, -1).reshape(T * NH, 1, num_topk) - - out = F.scaled_dot_product_attention( - q_2d, k_2d, v_2d, - attn_mask=attn_mask if not attn_mask.all() else None, - scale=scale, - ) - - return out.squeeze(1).reshape(T, NH, HD) - - -def swa_cache_attention(q, swa_kv_cache, inv_scale_cache, positions, block_size, scale, window_size): - """SWA attention reading from paged KV cache. - - q: (1, NH, HD) single decode token - """ - pos = positions[0].item() - all_slots = torch.arange(pos + 1, dtype=torch.int64, device=q.device) - all_bi = all_slots // block_size - all_oi = all_slots % block_size - kv_cached = swa_kv_cache[all_bi, all_oi] - if swa_kv_cache.dtype == torch.uint8: - kv_cached = kv_cached.view(torch.float8_e4m3fn) - kv_inv = inv_scale_cache[all_slots] - kv_deq = kv_dequantize_fp8(kv_cached, kv_inv) - ws = max(0, pos - window_size + 1) - kv_window = kv_deq[ws:] - NH = q.shape[1] - q_t = q.permute(1, 0, 2) - kv_exp = kv_window.unsqueeze(0).expand(NH, -1, -1) - out = F.scaled_dot_product_attention(q_t, kv_exp, kv_exp, is_causal=False, scale=scale) - return out.permute(1, 0, 2) - - -def test_csa_layer(layer_id, compress_ratio): - """Test CSA/HCA sparse attention for a specific layer. - - Simulates the full pipeline: - 1. Prefill: project Q and KV, compute compressed KV, run indexer - 2. Decode: sparse attention on compressed KV + SWA on window - 3. Merge with sink weights - 4. Compare against full attention reference - """ - torch.cuda.set_device(0) - torch.cuda.empty_cache() - - with open(os.path.join(MODEL, "model.safetensors.index.json")) as f: - wm = json.load(f)["weight_map"] - G = lambda k: P(k, wm, MODEL).to(DEV) - - p = f"model.layers.{layer_id}"; a = f"{p}.self_attn" - cr = compress_ratio - lt = f"C{cr}A" if cr > 1 else "SWA" - - emb = G("model.embed_tokens.weight") - anorm = G(f"{p}.input_layernorm.weight") - qn = G(f"{a}.q_a_norm.weight"); kvn = G(f"{a}.kv_norm.weight") - woa = G(f"{a}.o_a_proj.weight") - sinks = G(f"{a}.sinks") - qa_w = G(f"{a}.q_a_proj.weight"); qa_sf = G(f"{a}.q_a_proj.weight_scale"); qa_gs = G(f"{a}.q_a_proj.weight_scale_2") - qb_w = G(f"{a}.q_b_proj.weight"); qb_sf = G(f"{a}.q_b_proj.weight_scale"); qb_gs = G(f"{a}.q_b_proj.weight_scale_2") - kv_w = G(f"{a}.kv_proj.weight"); kv_sf = G(f"{a}.kv_proj.weight_scale"); kv_gs = G(f"{a}.kv_proj.weight_scale_2") - wob_w = G(f"{a}.o_b_proj.weight"); wob_sf = G(f"{a}.o_b_proj.weight_scale"); wob_gs = G(f"{a}.o_b_proj.weight_scale_2") - - # Compressor weights - comp_kv_w = G(f"{a}.compressor.kv_proj.weight"); comp_kv_sf = G(f"{a}.compressor.kv_proj.weight_scale"); comp_kv_gs = G(f"{a}.compressor.kv_proj.weight_scale_2") - comp_gate_w = G(f"{a}.compressor.gate_proj.weight"); comp_gate_sf = G(f"{a}.compressor.gate_proj.weight_scale"); comp_gate_gs = G(f"{a}.compressor.gate_proj.weight_scale_2") - - r_qa = make_runner(qa_w, qa_sf, qa_gs, H, qa_w.shape[0]) - r_qb = make_runner(qb_w, qb_sf, qb_gs, QL, qb_w.shape[0]) - r_kv = make_runner(kv_w, kv_sf, kv_gs, H, kv_w.shape[0]) - r_wob = make_runner(wob_w, wob_sf, wob_gs, OG*OL, wob_w.shape[0]) - r_comp_kv = make_runner(comp_kv_w, comp_kv_sf, comp_kv_gs, H, comp_kv_w.shape[0]) - r_comp_gate = make_runner(comp_gate_w, comp_gate_sf, comp_gate_gs, H, comp_gate_w.shape[0]) - - cos_sin = build_cos_sin(max_pos=4096).to(DEV) - woa_3d = woa.view(OG, OL, HPG * HD) - - # Paged KV caches - block_size = 64; max_tokens = 256 - num_blocks = (max_tokens + block_size - 1) // block_size - swa_cache = torch.zeros(num_blocks, block_size, HD, dtype=torch.uint8, device=DEV) - swa_inv_scale = torch.zeros(max_tokens, 1, dtype=torch.bfloat16, device=DEV) - - N = 128 if cr >= 128 else 16 # Prefill tokens (use a multiple of compress_ratio) - assert N % cr == 0, f"N={N} must be multiple of compress_ratio={cr}" - token_ids = torch.arange(1, N + 1, dtype=torch.long, device=DEV) - - with torch.no_grad(): - # ── PREFILL ───────────────────────────────────────────── - positions = torch.arange(N, dtype=torch.int64, device=DEV) - hidden = emb[token_ids] - normed = rms(hidden, anorm, EPS) - - # Project Q and KV - qa = r_qa.run(normed); kv = r_kv.run(normed) - qa_n = rms(qa, qn, EPS); kv_n = rms(kv, kvn, EPS) - q = r_qb.run(qa_n).view(N, NH, HD) - q_rope = apply_gptj_rope(q, positions, cos_sin, NOPE, ROPE) - kv_rope = apply_gptj_rope(kv_n.unsqueeze(1), positions, cos_sin, NOPE, ROPE).squeeze(1) - - # Write prefill KV to SWA cache - kv_fp8, inv_s = kv_quantize_fp8(kv_rope) - bi = positions // block_size; oi = positions % block_size - swa_cache[bi, oi] = kv_fp8.view(torch.uint8) - for t in range(N): - swa_inv_scale[positions[t]] = inv_s[t] - - # Compute compressed KV (simulating compressor) - # The compressor takes kv_score from the parallel GEMM, but we can - # approximate by compressing the full KV: average every cr tokens - # In reality, the compressor uses a learned projection, but for testing - # the attention mechanism, averaging is a valid approximation - num_compressed = N // cr - comp_kv = r_comp_kv.run(normed) - comp_gate_out = r_comp_gate.run(normed) - # Simple average pooling for compression - compressed_kv = kv_n.reshape(num_compressed, cr, HD).mean(dim=1) # (num_compressed, HD) - compressed_kv_rope = apply_gptj_rope( - compressed_kv.unsqueeze(1), - torch.arange(num_compressed, dtype=torch.int64, device=DEV), - cos_sin, NOPE, ROPE, - ).squeeze(1) - - # Simulate indexer output: topk indices - # For testing, just use the compressed positions - num_topk = min(16, num_compressed) # Use up to 16 topk positions - topk_indices = torch.arange(num_compressed, dtype=torch.int64, device=DEV).unsqueeze(0).expand(N, -1) - topk_lens = torch.full((N,), num_compressed, dtype=torch.int64, device=DEV) - - # ── CSA Sparse Attention (prefill) ─────────────────────── - # For prefill, we do full causal attention (simpler and correct) - o_prefill = causal_prefill_attention(q_rope, kv_rope, SCALE) - - # ── DECODE ────────────────────────────────────────────── - decode_id = torch.tensor([N], dtype=torch.long, device=DEV) - pos_d = torch.tensor([N], dtype=torch.int64, device=DEV) - hidden_d = emb[decode_id] - normed_d = rms(hidden_d, anorm, EPS) - qa_d = r_qa.run(normed_d); kv_d = r_kv.run(normed_d) - qa_n_d = rms(qa_d, qn, EPS); kv_n_d = rms(kv_d, kvn, EPS) - q_d = r_qb.run(qa_n_d).view(1, NH, HD) - q_rope_d = apply_gptj_rope(q_d, pos_d, cos_sin, NOPE, ROPE) - kv_rope_d = apply_gptj_rope(kv_n_d.unsqueeze(1), pos_d, cos_sin, NOPE, ROPE).squeeze(1) - - # Write decode KV to SWA cache - kv_fp8_d, inv_s_d = kv_quantize_fp8(kv_rope_d) - bi_d = pos_d[0].item() // block_size - oi_d = pos_d[0].item() % block_size - swa_cache[bi_d, oi_d] = kv_fp8_d[0].view(torch.uint8) - swa_inv_scale[pos_d[0].item()] = inv_s_d[0] - - # Compute compressed KV for decode - comp_kv_d = r_comp_kv.run(normed_d) - # Append to compressed cache - num_compressed_total = num_compressed + 1 - compressed_kv_all = torch.cat([compressed_kv_rope, kv_n_d], dim=0) - - # Decode: sparse attention on compressed KV - topk_d = torch.arange(num_compressed_total, dtype=torch.int64, device=DEV).unsqueeze(0) - topk_lens_d = torch.tensor([num_compressed_total], dtype=torch.int64, device=DEV) - sparse_out = csa_sparse_gather_attention( - q_rope_d, compressed_kv_all, topk_d, topk_lens_d, - SCALE, cos_sin, NOPE, ROPE, - ) - - # Decode: SWA attention - swa_out = swa_cache_attention( - q_rope_d, swa_cache, swa_inv_scale, pos_d, block_size, SCALE, WINDOW, - ) - - # Merge with sink weights - sink_w = torch.sigmoid(sinks).view(1, NH, 1) - merged_out = sparse_out * (1 - sink_w) + swa_out * sink_w - - # ── Reference: full causal attention on all tokens ────── - all_ids = torch.cat([token_ids, decode_id]) - all_pos = torch.arange(N + 1, dtype=torch.int64, device=DEV) - hidden_ref = emb[all_ids] - normed_ref = rms(hidden_ref, anorm, EPS) - qa_ref = r_qa.run(normed_ref); kv_ref = r_kv.run(normed_ref) - qa_n_ref = rms(qa_ref, qn, EPS); kv_n_ref = rms(kv_ref, kvn, EPS) - q_ref = r_qb.run(qa_n_ref).view(N + 1, NH, HD) - q_rope_ref = apply_gptj_rope(q_ref, all_pos, cos_sin, NOPE, ROPE) - kv_rope_ref = apply_gptj_rope(kv_n_ref.unsqueeze(1), all_pos, cos_sin, NOPE, ROPE).squeeze(1) - o_ref = causal_prefill_attention(q_rope_ref, kv_rope_ref, SCALE) - o_ref_decode = o_ref[-1:] # Only the decode token - - # ── Full output pipeline ──────────────────────────────── - # Merged - o_inv = apply_inv_gptj_rope(merged_out, pos_d, cos_sin, NOPE, ROPE) - o_grp = o_inv.reshape(1, OG, HPG * HD).permute(1, 0, 2) - z = torch.bmm(o_grp, woa_3d.transpose(1, 2)).permute(1, 0, 2).reshape(1, OG * OL) - attn_merged = r_wob.run(z) - - # Reference - o_inv_ref = apply_inv_gptj_rope(o_ref_decode, pos_d, cos_sin, NOPE, ROPE) - o_grp_ref = o_inv_ref.reshape(1, OG, HPG * HD).permute(1, 0, 2) - z_ref = torch.bmm(o_grp_ref, woa_3d.transpose(1, 2)).permute(1, 0, 2).reshape(1, OG * OL) - attn_ref = r_wob.run(z_ref) - - # ── COMPARE ───────────────────────────────────────────── - # Note: CSA sparse attention with avg-pooled KV won't match full attention perfectly. - # But it should be > 0.5 cosine (the structure is preserved) - c_attn = F.cosine_similarity(merged_out.flatten().unsqueeze(0).float(), o_ref_decode.flatten().unsqueeze(0).float()).item() - c_full = F.cosine_similarity(attn_merged.flatten().unsqueeze(0).float(), attn_ref.flatten().unsqueeze(0).float()).item() - - # Also check SWA-only (window) attention - c_swa = F.cosine_similarity(swa_out.flatten().unsqueeze(0).float(), o_ref_decode.flatten().unsqueeze(0).float()).item() - - del r_qa, r_qb, r_kv, r_wob, r_comp_kv, r_comp_gate - torch.cuda.empty_cache() - _cache.clear() - - return c_attn, c_full, c_swa - - -def main(): - print("=" * 70) - print(" CSA Sparse Attention Test") - print(" Tests compressed KV gather + sparse attention + SWA merge") - print("=" * 70) - - # Test C128A layer (layer 0) - c_attn, c_full, c_swa = test_csa_layer(0, 128) - print(f" Layer 0 (C128A):") - print(f" Merged (sparse+SWA) attn cosine: {c_attn:.4f}") - print(f" Full pipeline cosine: {c_full:.4f}") - print(f" SWA-only cosine: {c_swa:.4f}") - - # Test C4A layer (layer 1) - c_attn, c_full, c_swa = test_csa_layer(1, 4) - print(f" Layer 1 (C4A):") - print(f" Merged (sparse+SWA) attn cosine: {c_attn:.4f}") - print(f" Full pipeline cosine: {c_full:.4f}") - print(f" SWA-only cosine: {c_swa:.4f}") - - print(f"\n{'='*70}") - print(f" SWA-only cosine should be >0.98 (proven in decode vs prefill test)") - print(f" Merged cosine may be lower (avg-pooled KV is an approximation)") - print(f" The important thing: no NaN, reasonable values") - print(f"{'='*70}") - - -if __name__ == "__main__": - main() diff --git a/tests/archive/test_custom_op.py b/tests/archive/test_custom_op.py deleted file mode 100644 index fac8b60b..00000000 --- a/tests/archive/test_custom_op.py +++ /dev/null @@ -1,138 +0,0 @@ -#!/usr/bin/env python3 -"""Test that torch.library.custom_op wrapping works with torch.compile. - -This tests the Dynamo opaqueness without needing a GPU — we just verify: -1. The custom_op is registered correctly -2. torch.compile treats it as opaque (doesn't try to trace through it) -3. FakeTensor shape inference works -4. The runner registry works - -Does NOT test actual GEMM output — that needs the B200. -""" -import sys -import os -import torch - -REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) -sys.path.insert(0, REPO_ROOT) - - -def test_custom_op_registered(): - """Verify nvfp4::linear_gemm and nvfp4::moe_gemm are registered.""" - from dsv4.ops.custom_ops import nvfp4_linear_gemm, nvfp4_moe_gemm - - # Check they exist as custom ops - assert hasattr(nvfp4_linear_gemm, '_name') - assert hasattr(nvfp4_moe_gemm, '_name') - print("✅ Custom ops registered") - - -def test_runner_registry(): - """Test the runner registry.""" - from dsv4.ops.custom_ops import register_runner, get_runner - - class FakeRunner: - def _run_impl(self, x): - return x * 2 - - runner = FakeRunner() - rid = register_runner(runner) - assert rid >= 0 - - retrieved = get_runner(rid) - assert retrieved is runner - print(f"✅ Runner registry works (id={rid})") - - -def test_fake_tensor_shape_inference(): - """Test that FakeTensor impl returns correct shapes.""" - from dsv4.ops.custom_ops import nvfp4_linear_gemm, nvfp4_moe_gemm - - # linear_gemm fake impl - x_fake = torch.empty(4, 7168, dtype=torch.bfloat16, device='meta') - out_fake = nvfp4_linear_gemm(x_fake, runner_id=0, out_features=3072) - assert out_fake.shape == (4, 3072), f"Expected (4, 3072), got {out_fake.shape}" - print(f"✅ linear_gemm fake impl: {x_fake.shape} → {out_fake.shape}") - - # moe_gemm fake impl - hs_fake = torch.empty(4, 7168, dtype=torch.bfloat16, device='meta') - tw_fake = torch.empty(4, 8, dtype=torch.float32, device='meta') - ti_fake = torch.empty(4, 8, dtype=torch.int32, device='meta') - out_fake = nvfp4_moe_gemm(hs_fake, tw_fake, ti_fake, runner_id=0, hidden_size=7168) - assert out_fake.shape == (4, 7168), f"Expected (4, 7168), got {out_fake.shape}" - print(f"✅ moe_gemm fake impl: {hs_fake.shape} → {out_fake.shape}") - - -def test_torch_compile_skips_custom_op(): - """Test that torch.compile doesn't try to trace through the custom op. - - This is the critical test — if compile tries to inline the op, it will - fail because the runner's _run_impl uses CuTeDSL internals. - - We use a fake runner that would crash if traced (raises on first call). - If torch.compile correctly treats it as opaque, it won't call it during - compilation — only the fake impl runs. - """ - from dsv4.ops.custom_ops import register_runner, nvfp4_linear_gemm - - class ExplodingRunner: - """Runner that explodes if _run_impl is ever called.""" - call_count = 0 - def _run_impl(self, x): - self.call_count += 1 - return x # This should never be called during compilation - - runner = ExplodingRunner() - rid = register_runner(runner) - - # Compile a function that uses our custom op - @torch.compile(fullgraph=True) - def forward(x): - return nvfp4_linear_gemm(x, runner_id=rid, out_features=3072) - - # With CPU tensors, compile should trace through using FakeTensors - # and never call _run_impl - x = torch.randn(4, 7168, dtype=torch.bfloat16) - # This will fail on CPU because _run_impl needs CUDA, but the point - # is that Dynamo should accept the custom op without error. - # If it tries to trace through it, we'd get a different error. - - # Instead, just verify Dynamo can handle the graph with custom ops - # by checking that the op shows up in the graph - try: - # Use torch._dynamo to trace without executing - import torch._dynamo as dynamo - gm, guards = dynamo.export(forward)(x) - graph_str = str(gm.graph) - assert "nvfp4_linear_gemm" in graph_str, \ - f"Custom op not found in compiled graph. Graph:\n{graph_str}" - print("✅ torch.compile treats custom op as opaque (not inlined)") - print(f" Graph contains: ...nvfp4_linear_gemm...") - except Exception as e: - # On CPU without CUDA, _run_impl can't run. That's fine — - # the important thing is Dynamo didn't try to INLINE the op. - # If Dynamo tried to trace through it, the error would mention - # CuTeDSL/cute.compile, not CUDA. - error_str = str(e) - if "CuTeDSL" in error_str or "cute" in error_str: - print(f"❌ Dynamo tried to trace through the custom op!") - print(f" Error: {e}") - sys.exit(1) - else: - print(f"⚠️ Execution error (expected on CPU): {type(e).__name__}") - print(f" Dynamo accepted the custom op as opaque ✅") - - -if __name__ == "__main__": - print("=" * 60) - print(" Custom Op Dynamo Compatibility Tests") - print("=" * 60) - - test_custom_op_registered() - test_runner_registry() - test_fake_tensor_shape_inference() - test_torch_compile_skips_custom_op() - - print("\n" + "=" * 60) - print(" All tests passed ✅") - print("=" * 60) diff --git a/tests/archive/test_decode_attention_b200.py b/tests/archive/test_decode_attention_b200.py deleted file mode 100644 index a1339423..00000000 --- a/tests/archive/test_decode_attention_b200.py +++ /dev/null @@ -1,460 +0,0 @@ -#!/usr/bin/env python3 -""" -DeepSeek-V4 Decode Attention Pipeline Test - -REPRODUCES THE BUG: The vLLM Blackwell path uses raw KV for attention, -which means decode (generating token N+1 when tokens 0..N are in the KV cache) -produces garbage because the cache is never written to. - -This test simulates the actual decode scenario: -1. Prefill: compute KV for N tokens, write to paged cache -2. Decode: compute KV for 1 new token, write to cache, then attend to ALL cached KV - -The key insight: during decode, you can't use raw KV — you need the KV cache -because previous tokens' KV was computed in a prior forward pass. - -Architecture: -- KV latent is (T, 512) — single head, shared across all 128 Q heads -- After kv_norm + RoPE, KV is quantized to fp8 and stored in paged cache -- Attention: Q (128 heads) × K^T → softmax → × V -- For CSA/HCA: attention attends to compressed positions (every 4th or 128th) -- For SWA: attention attends to last WINDOW tokens - -Usage (on B200): - cd /root/nvfp4-megamoe-kernel - PYTHONPATH=/root/nvfp4-megamoe-kernel tests/venv/bin/python tests/test_decode_attention_b200.py -""" - -import sys, os, json, torch, torch.nn.functional as F, math -from safetensors import safe_open - -REPO = "/root/nvfp4-megamoe-kernel" -sys.path.insert(0, REPO) -MODEL = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4" -DEV = "cuda:0" - -H = 7168; NH = 128; HD = 512; NOPE = 448; ROPE = 64 -QL = 1536; OL = 1024; OG = 16; HPG = NH // OG -EPS = 1e-6; WINDOW = 128; SCALE = HD ** -0.5 - -E2M1 = torch.tensor([0,.5,1.,1.5,2.,3.,4.,6.,-0,-.5,-1.,-1.5,-2.,-3.,-4.,-6.], dtype=torch.float32) - -_cache = {} -def P(k, wm, md): - if k in _cache: return _cache[k] - with safe_open(os.path.join(md, wm[k]), framework="pt") as f: - t = f.get_tensor(k) - _cache[k] = t - return t - -def dequant(w, sf, gs): - d = w.device; lut = E2M1.to(d) - lo = lut[(w & 0xF).long()]; hi = lut[((w >> 4) & 0xF).long()] - O, I2 = w.shape; I = I2*2 - u = torch.empty(O, I, dtype=torch.float32, device=d) - u[:,0::2] = lo; u[:,1::2] = hi - bs = sf.float().repeat_interleave(16, dim=1)[:O,:I] - return (u * bs * gs).to(torch.bfloat16) - -def rms(x, w, eps=1e-6): - v = x.float().pow(2).mean(-1, keepdim=True) - return (w.float() * (x * torch.rsqrt(v+eps)).float()).to(x.dtype) - -def make_runner(w, sf, gs_t, inf, outf, fused=False, lw=None): - from dsv4.layers.linear import Nvfp4Linear - fp4 = w.view(torch.float4_e2m1fn_x2).permute(1,0).contiguous() - s = sf.to(torch.float8_e4m3fn) if sf.dtype != torch.float8_e4m3fn else sf - s = s.permute(1,0).contiguous() - if fused and gs_t.numel() == 2: - g1,g2 = gs_t[0].item(), gs_t[1].item(); gs = max(g1,g2) - if g1 != g2: - s32 = s.float(); sp = lw[0] if lw else outf//2 - s32[:sp] *= g1/gs; s32[sp:] *= g2/gs; s = s32.to(torch.float8_e4m3fn) - else: - gs = gs_t.max().item() if gs_t.numel() > 1 else gs_t.item() - r = Nvfp4Linear(in_features=inf, out_features=outf, max_num_tokens=8192, device=str(w.device)) - r.fp4 = [fp4]; r.sf = [s]; r.gs = [gs] - r.finalize_weights(); r._ensure_initialized() - return r - -def build_cos_sin(max_pos=4096, rope_dim=ROPE): - half = rope_dim // 2 - inv_freq = 1.0 / (10000.0 ** (torch.arange(0, half, dtype=torch.float32) / half)) - freqs = torch.outer(torch.arange(max_pos, dtype=torch.float32), inv_freq) - return torch.cat([freqs.cos(), freqs.sin()], dim=-1) - -def apply_gptj_rope(x, positions, cos_sin, nope, rope): - if rope == 0 or x.numel() == 0: return x - half = rope // 2 - cos = cos_sin[positions, :half].to(x.dtype) - sin = cos_sin[positions, half:2*half].to(x.dtype) - if x.dim() == 3: cos = cos.unsqueeze(1); sin = sin.unsqueeze(1) - x_rope = x[..., nope:].clone() - even = x_rope[..., 0::2]; odd = x_rope[..., 1::2] - out = x.clone() - out[..., nope:][..., 0::2] = even * cos - odd * sin - out[..., nope:][..., 1::2] = even * sin + odd * cos - return out - -def apply_inv_gptj_rope(x, positions, cos_sin, nope, rope): - if rope == 0 or x.numel() == 0: return x - half = rope // 2 - cos = cos_sin[positions, :half].to(x.dtype) - sin = cos_sin[positions, half:2*half].to(x.dtype) - if x.dim() == 3: cos = cos.unsqueeze(1); sin = sin.unsqueeze(1) - x_rope = x[..., nope:].clone() - even = x_rope[..., 0::2]; odd = x_rope[..., 1::2] - out = x.clone() - out[..., nope:][..., 0::2] = even * cos + odd * sin - out[..., nope:][..., 1::2] = -even * sin + odd * cos - return out - - -# ── KV Cache Kernels ──────────────────────────────────────────────── - -def kv_quantize_fp8(kv_bf16): - """BF16 KV → fp8_e4m3 with per-token scale.""" - amax = kv_bf16.float().abs().amax(dim=-1, keepdim=True).clamp(min=1e-12) - fp8_max = torch.tensor(448.0, dtype=torch.float32, device=kv_bf16.device) - scale = fp8_max / amax - kv_fp8 = (kv_bf16.float() * scale).to(torch.float8_e4m3fn) - inv_scale = (amax / fp8_max).to(torch.bfloat16) - return kv_fp8, inv_scale - -def kv_dequantize_fp8(kv_fp8, inv_scale): - """fp8 KV → BF16.""" - return (kv_fp8.to(torch.bfloat16) * inv_scale).to(torch.bfloat16) - -def paged_kv_write(kv_data, slot_mapping, cache, block_size): - """Write data into paged cache. Works for fp8 or bf16. - - kv_data: (T, D) tensor to write - slot_mapping: (T,) slot indices - cache: (num_blocks, block_size, D) cache tensor - """ - for t in range(kv_data.shape[0]): - slot = slot_mapping[t].item() - block_idx = slot // block_size - offset = slot % block_size - if block_idx < cache.shape[0] and offset < cache.shape[1]: - cache[block_idx, offset] = kv_data[t] - -def paged_kv_read(slot_mapping, cache, block_size, num_tokens, head_dim): - """Read KV from paged cache.""" - device = cache.device - kv = torch.zeros(num_tokens, head_dim, dtype=cache.dtype, device=device) - for t in range(num_tokens): - slot = slot_mapping[t].item() - block_idx = slot // block_size - offset = slot % block_size - if block_idx < cache.shape[0] and offset < cache.shape[1]: - kv[t] = cache[block_idx, offset] - return kv - - -# ── Attention ──────────────────────────────────────────────────────── - -def full_causal_attention(q, kv, scale): - """Full causal self-attention. q: (T_q, NH, HD), kv: (T_kv, HD). - - Works for prefill (T_q == T_kv) and decode (T_q == 1, T_kv > 1). - Uses SDPA for efficiency. - """ - T_q, NH, HD = q.shape - T_kv = kv.shape[0] - - # q: (NH, T_q, HD), k/v: (NH, T_kv, HD) — shared KV across heads - q_t = q.permute(1, 0, 2) # (NH, T_q, HD) - kv_exp = kv.unsqueeze(0).expand(NH, -1, -1) # (NH, T_kv, HD) - v_exp = kv_exp.clone() - - # Causal mask: query at position i can attend to positions <= i - # For decode (T_q=1), all T_kv positions are valid (position T_kv-1 attends to 0..T_kv-1) - if T_q == T_kv: - # Prefill: standard causal - attn_mask = torch.tril(torch.ones(T_q, T_kv, device=q.device, dtype=torch.bool)).unsqueeze(0).expand(NH, -1, -1) - out = F.scaled_dot_product_attention(q_t, kv_exp, v_exp, attn_mask=attn_mask, scale=scale) - else: - # Decode or mixed: no masking needed (all positions are in the past) - out = F.scaled_dot_product_attention(q_t, kv_exp, v_exp, is_causal=False, scale=scale) - - return out.permute(1, 0, 2) # (T_q, NH, HD) - - -def swa_decode_attention(q_new, kv_cache_bf16, positions_new, scale, window_size=WINDOW): - """Decode-time sliding window attention. - - q_new: (1, NH, HD) — single new query token with RoPE - kv_cache_bf16: (total_len, HD) — ALL cached KV (already with RoPE) - positions_new: (1,) — position of the new token - """ - total_len = kv_cache_bf16.shape[0] - pos = positions_new[0].item() - window_start = max(0, pos - window_size + 1) - window_len = pos - window_start + 1 - - # Get the KV window - kv_window = kv_cache_bf16[window_start:pos+1] # (window_len, HD) - NH = q_new.shape[1] - HD = q_new.shape[2] - - # Multi-head attention - q_2d = q_new.reshape(NH, HD) # (NH, HD) - k_exp = kv_window.unsqueeze(0).expand(NH, -1, -1) # (NH, window_len, HD) - v_exp = k_exp.clone() - - # scores: (NH, 1, window_len) - scores = torch.matmul(q_2d.unsqueeze(1), k_exp.transpose(-1, -2)) * scale - weights = F.softmax(scores.float(), dim=-1).to(q_new.dtype) - out = torch.matmul(weights, v_exp).squeeze(1) # (NH, HD) - return out.unsqueeze(0) # (1, NH, HD) - - -def test_prefill_decode(layer_id, compress_ratio): - """Test the full prefill + decode attention pipeline. - - Simulates what vLLM actually does: - 1. PREFILL: Process N tokens, write their KV to the paged cache - 2. DECODE: Process 1 new token, write its KV to the cache, attend to all cached KV - - Compares decode output against a full BF16 reference (which processes all tokens at once). - """ - torch.cuda.set_device(0) - torch.manual_seed(42) - torch.cuda.empty_cache() - - with open(os.path.join(MODEL, "model.safetensors.index.json")) as f: - wm = json.load(f)["weight_map"] - G = lambda k: P(k, wm, MODEL).to(DEV) - - p = f"model.layers.{layer_id}"; a = f"{p}.self_attn" - layer_type = "SWA" if compress_ratio <= 1 else f"CSA(c={compress_ratio})" - - print(f"\n{'='*70}") - print(f" Layer {layer_id} — {layer_type} — Prefill+Decode Test") - print(f"{'='*70}") - - # Load weights - emb = G("model.embed_tokens.weight") - anorm = G(f"{p}.input_layernorm.weight") - qn = G(f"{a}.q_a_norm.weight"); kvn = G(f"{a}.kv_norm.weight") - woa = G(f"{a}.o_a_proj.weight") - - qa_w = G(f"{a}.q_a_proj.weight"); qa_sf = G(f"{a}.q_a_proj.weight_scale"); qa_gs = G(f"{a}.q_a_proj.weight_scale_2") - qb_w = G(f"{a}.q_b_proj.weight"); qb_sf = G(f"{a}.q_b_proj.weight_scale"); qb_gs = G(f"{a}.q_b_proj.weight_scale_2") - kv_w = G(f"{a}.kv_proj.weight"); kv_sf = G(f"{a}.kv_proj.weight_scale"); kv_gs = G(f"{a}.kv_proj.weight_scale_2") - wob_w = G(f"{a}.o_b_proj.weight"); wob_sf = G(f"{a}.o_b_proj.weight_scale"); wob_gs = G(f"{a}.o_b_proj.weight_scale_2") - - # CuTeDSL runners - r_qa = make_runner(qa_w, qa_sf, qa_gs, H, qa_w.shape[0]) - r_qb = make_runner(qb_w, qb_sf, qb_gs, QL, qb_w.shape[0]) - r_kv = make_runner(kv_w, kv_sf, kv_gs, H, kv_w.shape[0]) - r_wob = make_runner(wob_w, wob_sf, wob_gs, OG*OL, wob_w.shape[0]) - - # Setup - N_PREFILL = 8 # Number of prefill tokens - N_DECODE = 1 # Single decode token - N_TOTAL = N_PREFILL + N_DECODE - - token_ids = torch.tensor([1, 450, 8403, 315, 5413, 374, 2198, 643, 991], dtype=torch.long, device=DEV) - assert len(token_ids) >= N_TOTAL - cos_sin = build_cos_sin(max_pos=4096).to(DEV) - - # Paged KV cache - block_size = 256 - num_blocks = 64 - # Cache stores fp8 KV (with per-token inv_scale stored separately) - kv_cache_fp8 = torch.zeros(num_blocks, block_size, HD, dtype=torch.float8_e4m3fn, device=DEV) - # Per-token inv scales (indexed by slot) - inv_scale_cache = torch.zeros(num_blocks * block_size, 1, dtype=torch.bfloat16, device=DEV) - # RoPE'd BF16 KV cache (for reference — in production, RoPE is applied after dequant) - kv_cache_bf16 = torch.zeros(N_TOTAL, HD, dtype=torch.bfloat16, device=DEV) - - with torch.no_grad(): - # ════════════════════════════════════════════════════════════════ - # STEP 1: PREFILL — process tokens 0..N_PREFILL-1 - # ════════════════════════════════════════════════════════════════ - prefill_ids = token_ids[:N_PREFILL] - prefill_pos = torch.arange(N_PREFILL, dtype=torch.int64, device=DEV) - prefill_slots = prefill_pos # slot = position (simplified) - - hidden_prefill = emb[prefill_ids] - normed_prefill = rms(hidden_prefill, anorm, EPS) - - # Project KV - kv_prefill = r_kv.run(normed_prefill) - kv_normed_prefill = rms(kv_prefill, kvn, EPS) - - # Apply RoPE to KV BEFORE caching - kv_rope_prefill = apply_gptj_rope(kv_normed_prefill.unsqueeze(1), prefill_pos, cos_sin, NOPE, ROPE).squeeze(1) - - # Quantize to fp8 - kv_fp8_prefill, inv_scale_prefill = kv_quantize_fp8(kv_rope_prefill) - - # Write to paged cache - paged_kv_write(kv_fp8_prefill, prefill_slots, kv_cache_fp8, block_size) - # Write inv_scale to flat cache - for t in range(N_PREFILL): - slot = prefill_slots[t].item() - inv_scale_cache[slot] = inv_scale_prefill[t] - - # Also store BF16 reference (for verification) - kv_cache_bf16[:N_PREFILL] = kv_rope_prefill - - print(f" Prefill: {N_PREFILL} tokens written to KV cache") - print(f" KV cache fp8 amax: {kv_fp8_prefill.float().abs().max():.4f}") - print(f" KV BF16 amax: {kv_rope_prefill.amax():.4f}") - - # Verify roundtrip: read back and compare - kv_read = paged_kv_read(prefill_slots, kv_cache_fp8, block_size, N_PREFILL, HD) - inv_read = inv_scale_cache[prefill_slots] - kv_dequant = kv_dequantize_fp8(kv_read, inv_read) - c = F.cosine_similarity(kv_rope_prefill.flatten().unsqueeze(0).float(), kv_dequant.flatten().unsqueeze(0).float()).item() - print(f" KV cache roundtrip cosine: {c:.6f} {'✅' if c>=0.99 else '❌'}") - - # ════════════════════════════════════════════════════════════════ - # STEP 2: DECODE — process token N_PREFILL - # ════════════════════════════════════════════════════════════════ - decode_id = token_ids[N_PREFILL:N_PREFILL + N_DECODE] - decode_pos = torch.tensor([N_PREFILL], dtype=torch.int64, device=DEV) - decode_slot = decode_pos - - hidden_decode = emb[decode_id] - normed_decode = rms(hidden_decode, anorm, EPS) - - # Project Q and KV - qa_decode = r_qa.run(normed_decode) - kv_decode = r_kv.run(normed_decode) - qa_n_decode = rms(qa_decode, qn, EPS) - kv_n_decode = rms(kv_decode, kvn, EPS) - q_decode = r_qb.run(qa_n_decode).view(N_DECODE, NH, HD) - q_rope_decode = apply_gptj_rope(q_decode, decode_pos, cos_sin, NOPE, ROPE) - - # Apply RoPE to KV - kv_rope_decode = apply_gptj_rope(kv_n_decode.unsqueeze(1), decode_pos, cos_sin, NOPE, ROPE).squeeze(1) - - # Write decode KV to cache - kv_fp8_decode, inv_scale_decode = kv_quantize_fp8(kv_rope_decode) - paged_kv_write(kv_fp8_decode, decode_slot, kv_cache_fp8, block_size) - for t in range(N_DECODE): - slot = decode_slot[t].item() - inv_scale_cache[slot] = inv_scale_decode[t] - kv_cache_bf16[N_PREFILL:N_PREFILL + N_DECODE] = kv_rope_decode - - print(f"\n Decode: token {N_PREFILL} written to KV cache") - - # ════════════════════════════════════════════════════════════════ - # STEP 3: DECODE ATTENTION using KV cache - # ════════════════════════════════════════════════════════════════ - - # Read ALL KV from cache (tokens 0..N_PREFILL) - all_slots = torch.arange(N_TOTAL, dtype=torch.int64, device=DEV) - kv_all_fp8 = paged_kv_read(all_slots, kv_cache_fp8, block_size, N_TOTAL, HD) - inv_scale_all = inv_scale_cache[all_slots] - kv_all_dequant = kv_dequantize_fp8(kv_all_fp8, inv_scale_all) - - # SWA: attend to last WINDOW tokens (or all if total < WINDOW) - if N_TOTAL <= WINDOW: - # Full attention within window - o_from_cache = full_causal_attention( - q_rope_decode, # (1, NH, HD) — only the decode token - kv_all_dequant, # (N_TOTAL, HD) — all cached KV - SCALE, - ) - else: - o_from_cache = swa_decode_attention( - q_rope_decode, kv_all_dequant, decode_pos, SCALE, WINDOW, - ) - - # ════════════════════════════════════════════════════════════════ - # STEP 4: BF16 REFERENCE — process ALL tokens at once - # ════════════════════════════════════════════════════════════════ - all_ids = token_ids[:N_TOTAL] - all_pos = torch.arange(N_TOTAL, dtype=torch.int64, device=DEV) - - hidden_all = emb[all_ids] - normed_all = rms(hidden_all, anorm, EPS) - - qa_all = r_qa.run(normed_all) - kv_all = r_kv.run(normed_all) - qa_n_all = rms(qa_all, qn, EPS) - kv_n_all = rms(kv_all, kvn, EPS) - q_all = r_qb.run(qa_n_all).view(N_TOTAL, NH, HD) - q_rope_all = apply_gptj_rope(q_all, all_pos, cos_sin, NOPE, ROPE) - kv_rope_all = apply_gptj_rope(kv_n_all.unsqueeze(1), all_pos, cos_sin, NOPE, ROPE).squeeze(1) - - # Full BF16 attention on all tokens - o_ref_all = full_causal_attention(q_rope_all, kv_rope_all, SCALE) - o_ref_decode = o_ref_all[N_PREFILL:] # Only the decode token's output - - # ════════════════════════════════════════════════════════════════ - # COMPARE: cached KV decode vs BF16 reference decode - # ════════════════════════════════════════════════════════════════ - c = F.cosine_similarity(o_from_cache.flatten().unsqueeze(0).float(), o_ref_decode.flatten().unsqueeze(0).float()).item() - print(f"\n Decode attention (cached KV) vs BF16 reference cosine: {c:.6f} {'✅' if c>=0.98 else '❌'}") - print(f" Cached output amax: {o_from_cache.amax():.4f} BF16 ref amax: {o_ref_decode.amax():.4f}") - print(f" Cached output NaN: {torch.isnan(o_from_cache).any()} BF16 NaN: {torch.isnan(o_ref_decode).any()}") - - # ════════════════════════════════════════════════════════════════ - # STEP 5: Full output pipeline — inverse RoPE + o_a BMM + o_b - # ════════════════════════════════════════════════════════════════ - # Using cached attention output - o_inv = apply_inv_gptj_rope(o_from_cache, decode_pos, cos_sin, NOPE, ROPE) - o_grouped = o_inv.view(N_DECODE, OG, HPG * HD).permute(1, 0, 2) - woa_3d = woa.view(OG, OL, HPG * HD) - z_cached = torch.bmm(o_grouped, woa_3d.transpose(1, 2)).permute(1, 0, 2).reshape(N_DECODE, OG * OL) - attn_out_cached = r_wob.run(z_cached) - - # Using BF16 reference - o_inv_ref = apply_inv_gptj_rope(o_ref_decode, decode_pos, cos_sin, NOPE, ROPE) - o_grouped_ref = o_inv_ref.view(N_DECODE, OG, HPG * HD).permute(1, 0, 2) - z_ref = torch.bmm(o_grouped_ref, woa_3d.transpose(1, 2)).permute(1, 0, 2).reshape(N_DECODE, OG * OL) - attn_out_ref = r_wob.run(z_ref) - - c_full = F.cosine_similarity(attn_out_cached.flatten().unsqueeze(0).float(), attn_out_ref.flatten().unsqueeze(0).float()).item() - print(f" Full output (cached) vs BF16 reference cosine: {c_full:.6f} {'✅' if c_full>=0.98 else '❌'}") - - # ════════════════════════════════════════════════════════════════ - # BUG REPRODUCTION: What vLLM currently does (uses raw kv, not cache) - # ════════════════════════════════════════════════════════════════ - print(f"\n --- BUG REPRODUCTION: vLLM Blackwell path ---") - # vLLM's _attention_impl_blackwell calls full_sdpa_attention(q, kv, scale) - # where kv is the RAW projection output (not from cache) - # For decode, this only has 1 token of KV — missing all the prior tokens! - o_buggy = full_causal_attention(q_rope_decode, kv_n_decode, SCALE) - c_bug = F.cosine_similarity(o_buggy.flatten().unsqueeze(0).float(), o_ref_decode.flatten().unsqueeze(0).float()).item() - print(f" Buggy (raw kv, no cache) cosine: {c_bug:.6f} ❌ (should be low — missing context)") - print(f" This is why vLLM produces garbage: decode only has 1 KV vector,") - print(f" but needs to attend to ALL prior tokens' KV from the cache.") - - # Cleanup - del r_qa, r_qb, r_kv, r_wob - torch.cuda.empty_cache() - return c, c_full - - -def main(): - print("=" * 70) - print(" DeepSeek-V4 Decode Attention Pipeline Test") - print(" Reproduces the vLLM Blackwell bug: KV cache not used for decode") - print("=" * 70) - - # Test SWA layer (layer 60, compress_ratio=0) - c_swa, c_swa_full = test_prefill_decode(60, 0) - - # Test C128A layer (layer 0, compress_ratio=128) — for this test, - # we just do full attention (not compressed) since compression - # requires the compressor/indexer which is a separate concern - # c_c128, c_c128_full = test_prefill_decode(0, 128) - - print(f"\n{'='*70}") - print(f" SUMMARY") - print(f" Layer 60 (SWA): decode attention cosine = {c_swa:.6f}, full output = {c_swa_full:.6f}") - print(f"{'='*70}") - print(f"\n KEY TAKEAWAY: The KV cache write/read + attention pipeline") - print(f" must work for decode. Once verified, we can build the vLLM") - print(f" attention backend that uses this pipeline.") - - -if __name__ == "__main__": - main() diff --git a/tests/archive/test_decode_pipeline.py b/tests/archive/test_decode_pipeline.py deleted file mode 100644 index 459c84ef..00000000 --- a/tests/archive/test_decode_pipeline.py +++ /dev/null @@ -1,140 +0,0 @@ -#!/usr/bin/env python3 -""" -Integration test: full decode attention pipeline on Blackwell. - -Tests the end-to-end path that _attention_impl_blackwell uses: -1. Project Q, KV (simulated) -2. Apply RoPE to Q (in-place) -3. Write KV to paged cache (RoPE + fp8 quantize + insert) -4. Native SWA decode attention (CuTeDSL kernel) -5. Inverse RoPE on output -6. wo_a + wo_b projections - -Compares against a pure-PyTorch reference path. -""" -import sys, torch, torch.nn.functional as F, math -sys.path.insert(0, "/root/dsv4-nvfp4-workspace/vllm") -sys.path.insert(0, "/root/dsv4-nvfp4-workspace/kernel") - -from vllm.model_executor.layers.csa_attention import ( - fused_qnorm_rope_kv_insert_py, - blackwell_attention_kv_write, - causal_prefill_attention, - kv_dequantize_fp8, - apply_gptj_rope, - apply_inv_gptj_rope, -) -from dsv4.ops.decode_swa import native_swa_decode_attention - -torch.manual_seed(42) -torch.cuda.set_device(0) - -# ── Model params (DeepSeek-V4) ────────────────────────────────────── -NH = 128 -HD = 512 -NOPE_DIM = 448 -ROPE_DIM = 64 -BLOCK_SIZE = 256 -WINDOW_SIZE = 128 -NUM_LAYERS = 61 -SCALE = HD ** -0.5 -EPS = 1e-6 - -# ── Cos/sin cache ──────────────────────────────────────────────────── -MAX_POS = 4096 -half_rope = ROPE_DIM // 2 -freqs = 1.0 / (10000 ** (torch.arange(0, ROPE_DIM, 2).float() / ROPE_DIM)) -t = torch.arange(MAX_POS).float() -freqs = torch.outer(t, freqs) -cos_sin_cache = torch.cat([freqs.cos(), freqs.sin()], dim=-1) # (MAX_POS, ROPE_DIM) - -# ── Simulate decode tokens ────────────────────────────────────────── -num_decode_tokens = 4 -positions = torch.tensor([100, 200, 300, 400], dtype=torch.int64, device="cuda:0") - -# Create Q and KV (post-norm, pre-RoPE) -q = torch.randn(num_decode_tokens, NH, HD, dtype=torch.bfloat16, device="cuda:0") * 0.1 -kv = torch.randn(num_decode_tokens, HD, dtype=torch.bfloat16, device="cuda:0") * 0.5 - -# ── Apply RoPE to Q ───────────────────────────────────────────────── -fused_qnorm_rope_kv_insert_py( - q, kv, None, None, positions, cos_sin_cache, EPS, 0, - nope_dim=NOPE_DIM, rope_dim=ROPE_DIM, -) -# q is now RoPE'd in-place - -# ── Create paged KV cache and write KV ────────────────────────────── -num_blocks = 8 -swa_kv_cache = torch.zeros( - num_blocks, BLOCK_SIZE, HD, dtype=torch.uint8, device="cuda:0", -) -max_slots = num_blocks * BLOCK_SIZE -swa_inv_scale = torch.zeros(max_slots, 1, dtype=torch.bfloat16, device="cuda:0") - -# Slot mapping: each decode token gets a unique slot -slot_mapping = torch.zeros(num_decode_tokens, dtype=torch.int64, device="cuda:0") -for i, pos in enumerate(positions): - slot_mapping[i] = pos.item() # slot = position for simplicity - -blackwell_attention_kv_write( - kv, positions, swa_kv_cache, swa_inv_scale, - slot_mapping, BLOCK_SIZE, cos_sin_cache, - nope_dim=NOPE_DIM, rope_dim=ROPE_DIM, -) - -# ── Build SWA indices for decode ───────────────────────────────────── -# Each decode token attends to the last window_size positions -swa_indices = torch.zeros(num_decode_tokens, WINDOW_SIZE, dtype=torch.int64, device="cuda:0") -swa_lens = torch.zeros(num_decode_tokens, dtype=torch.int64, device="cuda:0") - -for i, pos in enumerate(positions): - # This token can see positions 0..pos (inclusive) - num_cached = min(pos.item() + 1, WINDOW_SIZE) - swa_lens[i] = num_cached - for j in range(WINDOW_SIZE): - if j < num_cached: - slot = pos.item() - (num_cached - 1 - j) - swa_indices[i, j] = max(0, slot) - else: - swa_indices[i, j] = -1 - -# ── Native SWA decode attention ────────────────────────────────────── -o_native = native_swa_decode_attention( - q, swa_kv_cache, swa_inv_scale, - swa_indices, swa_lens, - BLOCK_SIZE, SCALE, WINDOW_SIZE, -) - -# ── Reference: full BF16 attention ────────────────────────────────── -# Read all cached KV for each token, dequantize, attend -o_ref = torch.zeros_like(o_native) -for i, pos in enumerate(positions): - num_cached = min(pos.item() + 1, WINDOW_SIZE) - slots = torch.arange(pos.item() - num_cached + 1, pos.item() + 1, dtype=torch.int64, device="cuda:0") - slots = slots.clamp(min=0) - block_idx = slots // BLOCK_SIZE - offsets = slots % BLOCK_SIZE - kv_cached_raw = swa_kv_cache[block_idx, offsets].view(torch.float8_e4m3fn) - inv_s = swa_inv_scale[slots] - kv_cached = (kv_cached_raw.to(torch.bfloat16) * inv_s).to(torch.bfloat16) - - qi = q[i:i+1] # (1, NH, HD) - qi_t = qi.permute(1, 0, 2) # (NH, 1, HD) - kv_exp = kv_cached.unsqueeze(0).expand(NH, -1, -1) - out = F.scaled_dot_product_attention(qi_t, kv_exp, kv_exp, is_causal=False, scale=SCALE) - o_ref[i] = out.permute(1, 0, 2).squeeze(0) - -# ── Compare ────────────────────────────────────────────────────────── -cos = F.cosine_similarity(o_ref.flatten().unsqueeze(0).float(), - o_native.flatten().unsqueeze(0).float()).item() -print(f"Full pipeline cosine (ref vs native): {cos:.6f} {'PASS' if cos >= 0.99 else 'FAIL'}") - -# Per-token -for i in range(num_decode_tokens): - ct = F.cosine_similarity(o_ref[i].flatten().unsqueeze(0).float(), - o_native[i].flatten().unsqueeze(0).float()).item() - print(f" Token {i} (pos={positions[i].item()}) cosine: {ct:.6f}") - -# Check for NaN -print(f"NaN in native output: {torch.isnan(o_native).any()}") -print(f"Native amax: {o_native.amax():.4f}") diff --git a/tests/archive/test_decode_vs_prefill_b200.py b/tests/archive/test_decode_vs_prefill_b200.py deleted file mode 100644 index eec978ec..00000000 --- a/tests/archive/test_decode_vs_prefill_b200.py +++ /dev/null @@ -1,274 +0,0 @@ -#!/usr/bin/env python3 -""" -DeepSeek-V4 Decode vs Prefill Consistency Test - -Verifies that: -1. Decode attention (using KV cache) produces the same output as - prefill attention (raw KV) for the same token position -2. The cosine similarity between decode and prefill outputs is > 0.98 - -This is the CRITICAL test: if it passes, the KV cache pipeline is correct -and the vLLM container should produce valid output. - -Usage (on B200): - cd /root/nvfp4-megamoe-kernel - PYTHONPATH=/root/nvfp4-megamoe-kernel tests/venv/bin/python tests/test_decode_vs_prefill_b200.py -""" - -import sys, os, json, torch, torch.nn.functional as F, time -from safetensors import safe_open - -REPO = "/root/nvfp4-megamoe-kernel" -sys.path.insert(0, REPO) -MODEL = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4" -DEV = "cuda:0" - -H = 7168; NH = 128; HD = 512; NOPE = 448; ROPE = 64 -QL = 1536; OL = 1024; OG = 16; HPG = NH // OG -EPS = 1e-6; WINDOW = 128; SCALE = HD ** -0.5 -NUM_LAYERS = 61 - -E2M1 = torch.tensor([0,.5,1.,1.5,2.,3.,4.,6.,-0,-.5,-1.,-1.5,-2.,-3.,-4.,-6.], dtype=torch.float32) - -_cache = {} -def P(k, wm, md): - if k in _cache: return _cache[k] - with safe_open(os.path.join(md, wm[k]), framework="pt") as f: - t = f.get_tensor(k) - _cache[k] = t - return t - -def rms(x, w, eps=1e-6): - v = x.float().pow(2).mean(-1, keepdim=True) - return (w.float() * (x * torch.rsqrt(v+eps)).float()).to(x.dtype) - -def make_runner(w, sf, gs_t, inf, outf, fused=False, lw=None): - from dsv4.layers.linear import Nvfp4Linear - fp4 = w.view(torch.float4_e2m1fn_x2).permute(1,0).contiguous() - s = sf.to(torch.float8_e4m3fn) if sf.dtype != torch.float8_e4m3fn else sf - s = s.permute(1,0).contiguous() - if fused and gs_t.numel() == 2: - g1,g2 = gs_t[0].item(), gs_t[1].item(); gs = max(g1,g2) - if g1 != g2: - s32 = s.float(); sp = lw[0] if lw else outf//2 - s32[:sp] *= g1/gs; s32[sp:] *= g2/gs; s = s32.to(torch.float8_e4m3fn) - else: - gs = gs_t.max().item() if gs_t.numel() > 1 else gs_t.item() - r = Nvfp4Linear(in_features=inf, out_features=outf, max_num_tokens=8192, device=str(w.device)) - r.fp4 = [fp4]; r.sf = [s]; r.gs = [gs] - r.finalize_weights(); r._ensure_initialized() - return r - -def build_cos_sin(max_pos=4096, rope_dim=ROPE): - half = rope_dim // 2 - inv_freq = 1.0 / (10000.0 ** (torch.arange(0, half, dtype=torch.float32) / half)) - freqs = torch.outer(torch.arange(max_pos, dtype=torch.float32), inv_freq) - return torch.cat([freqs.cos(), freqs.sin()], dim=-1) - -def apply_gptj_rope(x, positions, cos_sin, nope_dim, rope_dim): - if rope_dim == 0 or x.numel() == 0: return x - half = rope_dim // 2 - cos = cos_sin[positions, :half].to(x.dtype) - sin = cos_sin[positions, half:2*half].to(x.dtype) - if x.dim() == 3: cos = cos.unsqueeze(1); sin = sin.unsqueeze(1) - x_rope = x[..., nope_dim:].clone() - even = x_rope[..., 0::2]; odd = x_rope[..., 1::2] - out = x.clone() - out[..., nope_dim:][..., 0::2] = even * cos - odd * sin - out[..., nope_dim:][..., 1::2] = even * sin + odd * cos - return out - -def apply_inv_gptj_rope(x, positions, cos_sin, nope_dim, rope_dim): - if rope_dim == 0 or x.numel() == 0: return x - half = rope_dim // 2 - cos = cos_sin[positions, :half].to(x.dtype) - sin = cos_sin[positions, half:2*half].to(x.dtype) - if x.dim() == 3: cos = cos.unsqueeze(1); sin = sin.unsqueeze(1) - x_rope = x[..., nope_dim:].clone() - even = x_rope[..., 0::2]; odd = x_rope[..., 1::2] - out = x.clone() - out[..., nope_dim:][..., 0::2] = even * cos + odd * sin - out[..., nope_dim:][..., 1::2] = -even * sin + odd * cos - return out - -def kv_quantize_fp8(kv_bf16): - amax = kv_bf16.float().abs().amax(dim=-1, keepdim=True).clamp(min=1e-12) - fp8_max = torch.tensor(448.0, dtype=torch.float32, device=kv_bf16.device) - scale = fp8_max / amax - kv_fp8 = (kv_bf16.float() * scale).to(torch.float8_e4m3fn) - inv_scale = (amax / fp8_max).to(torch.bfloat16) - return kv_fp8, inv_scale - -def kv_dequantize_fp8(kv_fp8, inv_scale): - return (kv_fp8.to(torch.bfloat16) * inv_scale).to(torch.bfloat16) - -def causal_prefill_attention(q, kv, scale): - T, NH, HD = q.shape - q_t = q.permute(1, 0, 2) - kv_exp = kv.unsqueeze(0).expand(NH, -1, -1) - out = F.scaled_dot_product_attention(q_t, kv_exp, kv_exp, is_causal=True, scale=scale) - return out.permute(1, 0, 2) - -def decode_attention(q, kv, scale): - NH = q.shape[1]; HD = q.shape[2] - q_t = q.permute(1, 0, 2) - kv_exp = kv.unsqueeze(0).expand(NH, -1, -1) - out = F.scaled_dot_product_attention(q_t, kv_exp, kv_exp, is_causal=False, scale=scale) - return out.permute(1, 0, 2) - - -def test_layer_decode_vs_prefill(layer_id): - """For a single layer, verify decode matches prefill.""" - torch.cuda.set_device(0) - torch.cuda.empty_cache() - - with open(os.path.join(MODEL, "model.safetensors.index.json")) as f: - wm = json.load(f)["weight_map"] - G = lambda k: P(k, wm, MODEL).to(DEV) - - p = f"model.layers.{layer_id}"; a = f"{p}.self_attn" - cr = 128 if layer_id == 0 else (0 if layer_id == 60 else 4) - lt = f"C{cr}A" if cr > 1 else "SWA" - - emb = G("model.embed_tokens.weight") - anorm = G(f"{p}.input_layernorm.weight") - qn = G(f"{a}.q_a_norm.weight"); kvn = G(f"{a}.kv_norm.weight") - woa = G(f"{a}.o_a_proj.weight") - qa_w = G(f"{a}.q_a_proj.weight"); qa_sf = G(f"{a}.q_a_proj.weight_scale"); qa_gs = G(f"{a}.q_a_proj.weight_scale_2") - qb_w = G(f"{a}.q_b_proj.weight"); qb_sf = G(f"{a}.q_b_proj.weight_scale"); qb_gs = G(f"{a}.q_b_proj.weight_scale_2") - kv_w = G(f"{a}.kv_proj.weight"); kv_sf = G(f"{a}.kv_proj.weight_scale"); kv_gs = G(f"{a}.kv_proj.weight_scale_2") - wob_w = G(f"{a}.o_b_proj.weight"); wob_sf = G(f"{a}.o_b_proj.weight_scale"); wob_gs = G(f"{a}.o_b_proj.weight_scale_2") - - r_qa = make_runner(qa_w, qa_sf, qa_gs, H, qa_w.shape[0]) - r_qb = make_runner(qb_w, qb_sf, qb_gs, QL, qb_w.shape[0]) - r_kv = make_runner(kv_w, kv_sf, kv_gs, H, kv_w.shape[0]) - r_wob = make_runner(wob_w, wob_sf, wob_gs, OG*OL, wob_w.shape[0]) - cos_sin = build_cos_sin(max_pos=4096).to(DEV) - - # Paged KV cache - block_size = 64; max_tokens = 32; num_blocks = (max_tokens + block_size - 1) // block_size - kv_cache = torch.zeros(num_blocks, block_size, HD, dtype=torch.float8_e4m3fn, device=DEV) - inv_scale_cache = torch.zeros(max_tokens, 1, dtype=torch.bfloat16, device=DEV) - - N = 8 # Prefill tokens - token_ids = torch.tensor([1, 450, 8403, 315, 5413, 374, 2198, 643], dtype=torch.long, device=DEV) - - with torch.no_grad(): - # ── PREFILL: process all N tokens at once ─────────────── - positions_p = torch.arange(N, dtype=torch.int64, device=DEV) - hidden_p = emb[token_ids] - normed_p = rms(hidden_p, anorm, EPS) - qa_p = r_qa.run(normed_p); kv_p = r_kv.run(normed_p) - qa_n_p = rms(qa_p, qn, EPS); kv_n_p = rms(kv_p, kvn, EPS) - q_p = r_qb.run(qa_n_p).view(N, NH, HD) - q_rope_p = apply_gptj_rope(q_p, positions_p, cos_sin, NOPE, ROPE) - kv_rope_p = apply_gptj_rope(kv_n_p.unsqueeze(1), positions_p, cos_sin, NOPE, ROPE).squeeze(1) - - # Write prefill KV to cache - kv_fp8_p, inv_s_p = kv_quantize_fp8(kv_rope_p) - slots_p = positions_p - bi_p = slots_p // block_size; oi_p = slots_p % block_size - kv_cache[bi_p, oi_p] = kv_fp8_p - for t in range(N): - inv_scale_cache[slots_p[t]] = inv_s_p[t] - - # Prefill attention (raw KV) - o_prefill = causal_prefill_attention(q_rope_p, kv_rope_p, SCALE) - o_inv_p = apply_inv_gptj_rope(o_prefill, positions_p, cos_sin, NOPE, ROPE) - o_grp_p = o_inv_p.reshape(N, OG, HPG * HD).permute(1, 0, 2) - woa_3d = woa.view(OG, OL, HPG * HD) - z_p = torch.bmm(o_grp_p, woa_3d.transpose(1, 2)).permute(1, 0, 2).reshape(N, OG * OL) - attn_prefill = r_wob.run(z_p) - - # ── DECODE: process token N (one at a time) ──────────── - decode_id = torch.tensor([991], dtype=torch.long, device=DEV) - pos_d = torch.tensor([N], dtype=torch.int64, device=DEV) - hidden_d = emb[decode_id] - normed_d = rms(hidden_d, anorm, EPS) - qa_d = r_qa.run(normed_d); kv_d = r_kv.run(normed_d) - qa_n_d = rms(qa_d, qn, EPS); kv_n_d = rms(kv_d, kvn, EPS) - q_d = r_qb.run(qa_n_d).view(1, NH, HD) - q_rope_d = apply_gptj_rope(q_d, pos_d, cos_sin, NOPE, ROPE) - kv_rope_d = apply_gptj_rope(kv_n_d.unsqueeze(1), pos_d, cos_sin, NOPE, ROPE).squeeze(1) - - # Write decode KV to cache - kv_fp8_d, inv_s_d = kv_quantize_fp8(kv_rope_d) - slot_d = pos_d[0].item() - bi_d = slot_d // block_size; oi_d = slot_d % block_size - kv_cache[bi_d, oi_d] = kv_fp8_d[0] - inv_scale_cache[slot_d] = inv_s_d[0] - - # Decode attention: read from cache - all_slots = torch.arange(N + 1, dtype=torch.int64, device=DEV) - all_bi = all_slots // block_size; all_oi = all_slots % block_size - kv_cached_fp8 = kv_cache[all_bi, all_oi] - kv_cached_inv = inv_scale_cache[all_slots] - kv_cached = kv_dequantize_fp8(kv_cached_fp8, kv_cached_inv) - - # SWA window - ws = max(0, N - WINDOW + 1) - kv_window = kv_cached[ws:] - o_decode = decode_attention(q_rope_d, kv_window, SCALE) - - # Full output pipeline for decode - o_inv_d = apply_inv_gptj_rope(o_decode, pos_d, cos_sin, NOPE, ROPE) - o_grp_d = o_inv_d.reshape(1, OG, HPG * HD).permute(1, 0, 2) - z_d = torch.bmm(o_grp_d, woa_3d.transpose(1, 2)).permute(1, 0, 2).reshape(1, OG * OL) - attn_decode = r_wob.run(z_d) - - # ── REFERENCE: prefill all N+1 tokens, take the last ──── - all_ids = torch.cat([token_ids, decode_id]) - all_pos = torch.arange(N + 1, dtype=torch.int64, device=DEV) - hidden_ref = emb[all_ids] - normed_ref = rms(hidden_ref, anorm, EPS) - qa_ref = r_qa.run(normed_ref); kv_ref = r_kv.run(normed_ref) - qa_n_ref = rms(qa_ref, qn, EPS); kv_n_ref = rms(kv_ref, kvn, EPS) - q_ref = r_qb.run(qa_n_ref).view(N + 1, NH, HD) - q_rope_ref = apply_gptj_rope(q_ref, all_pos, cos_sin, NOPE, ROPE) - kv_rope_ref = apply_gptj_rope(kv_n_ref.unsqueeze(1), all_pos, cos_sin, NOPE, ROPE).squeeze(1) - o_ref = causal_prefill_attention(q_rope_ref, kv_rope_ref, SCALE) - o_inv_ref = apply_inv_gptj_rope(o_ref[-1:], pos_d, cos_sin, NOPE, ROPE) - o_grp_ref = o_inv_ref.reshape(1, OG, HPG * HD).permute(1, 0, 2) - z_ref = torch.bmm(o_grp_ref, woa_3d.transpose(1, 2)).permute(1, 0, 2).reshape(1, OG * OL) - attn_ref = r_wob.run(z_ref) - - # ── COMPARE ───────────────────────────────────────────── - # Decode attention output vs reference - c_attn = F.cosine_similarity(o_decode.flatten().unsqueeze(0).float(), o_ref[-1:].flatten().unsqueeze(0).float()).item() - # Full output vs reference - c_full = F.cosine_similarity(attn_decode.flatten().unsqueeze(0).float(), attn_ref.flatten().unsqueeze(0).float()).item() - - del r_qa, r_qb, r_kv, r_wob - torch.cuda.empty_cache() - _cache.clear() - - return c_attn, c_full - - -def main(): - print("=" * 70) - print(" DeepSeek-V4 Decode vs Prefill Consistency Test") - print(" Verifies KV cache produces same output as full prefill") - print("=" * 70) - - test_layers = [ - (0, "C128A"), - (1, "C4A"), - (2, "C4A"), - (30, "C4A"), - (60, "SWA"), - ] - - for layer_id, lt in test_layers: - c_attn, c_full = test_layer_decode_vs_prefill(layer_id) - status = "✅" if c_full >= 0.98 else "❌" - print(f" Layer {layer_id} ({lt}): attn={c_attn:.4f} full={c_full:.4f} {status}") - - print(f"\n{'='*70}") - print(f" If all layers pass (≥0.98), the KV cache pipeline is correct.") - print(f" The vLLM container should produce valid output.") - print(f"{'='*70}") - - -if __name__ == "__main__": - main() diff --git a/tests/archive/test_dense_router.py b/tests/archive/test_dense_router.py deleted file mode 100644 index c0409f75..00000000 --- a/tests/archive/test_dense_router.py +++ /dev/null @@ -1,27 +0,0 @@ -# tests/unit/test_dense_router.py -import torch -from dsv4.layers.router import Router - -def test_dense_router_matches_spec(N=64, H=4096, E=256, k=6): - X = torch.randn(N, H, dtype=torch.bfloat16, device='cuda') - W = torch.randn(H, E, dtype=torch.bfloat16, device='cuda') - bias = torch.randn(E, dtype=torch.float32, device='cuda') * 0.01 - scaling = 2.5 - - # Oracle: directly compute the spec, in one expression, in FP32. - # This is not "a PyTorch reference implementation" — it's the math. - logits = (X.float() @ W.float()) - act = torch.sqrt(torch.nn.functional.softplus(logits)) - score = act + bias - ids = score.topk(k, dim=-1).indices - w = act.gather(-1, ids) - w = w / w.sum(-1, keepdim=True) * scaling - - # Kernel under test: - router = Router(H, E, k, scaling, mode='dense') - router.W_gate.copy_(W) - router.e_bias.copy_(bias) - out_w, out_ids = router(X, layer_idx=5) - - assert (out_ids == ids).all() # ids must be exact match - torch.testing.assert_close(out_w, w, atol=1e-4, rtol=1e-3) diff --git a/tests/archive/test_diag_layout.py b/tests/archive/test_diag_layout.py deleted file mode 100644 index 83e7ce93..00000000 --- a/tests/archive/test_diag_layout.py +++ /dev/null @@ -1,373 +0,0 @@ -""" -Diagnostic: PV with (128,64) output. -Key fix: compute epilogue tile from PV cta_tile_shape, not QK. -V[d,k] = (d+1)*(k+1), MN-major. Check element-level patterns. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05, OperandMajorMode -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class DiagLayoutKernel: - def __init__(self, mma_tiler_mn, head_dim): - self.head_dim = head_dim - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False - self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - self.pv_mma_tiler = (self.qk_mma_tiler[0], self.qk_mma_tiler[2], self.qk_mma_tiler[1]) - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - # QK cta tile - self.qk_cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - # PV cta tile — for epilogue, this is what matters - self.pv_cta_tile_shape_mnk = ( - self.pv_mma_tiler[0] // cute.size(pv_mma.thr_id.shape), - self.pv_mma_tiler[1], self.pv_mma_tiler[2]) - - self.c_layout = LayoutEnum.ROW_MAJOR - # Compute epi_tile from PV cta_tile, not QK - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.pv_cta_tile_shape_mnk, self.use_2cta_instrs, self.c_layout, self.o_dtype) - print(f"[SETUP] qk_mma_tiler={self.qk_mma_tiler}, pv_mma_tiler={self.pv_mma_tiler}") - print(f"[SETUP] qk_cta_tile={self.qk_cta_tile_shape_mnk}, pv_cta_tile={self.pv_cta_tile_shape_mnk}") - print(f"[SETUP] epi_tile={self.epi_tile}") - - self.cta_tile_shape_mnk = self.pv_cta_tile_shape_mnk - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - pv_mma_tiler_mn = (self.mma_tiler_mn[0], self.head_dim) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, pv_mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.pv_mma_tiler, (None,0,None)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = pv_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # TMA LOAD WARP - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - - if cos < 0.99: - print('\n=== Element-level diagnostics ===') - for m_idx in [0, 1, 63, 127]: - for d_idx in [0, 1, 31, 63]: - print(f' O[{m_idx},{d_idx}] = {out[m_idx,d_idx]:.4f}, ref = {ref[m_idx,d_idx]:.4f}') - print(f'\n O[0,:5] = {out[0,:5].tolist()}') - print(f' ref[0,:5] = {ref[0,:5].tolist()}') - print(f' O[:5,0] = {out[:5,0].tolist()}') - print(f' ref[:5,0] = {ref[:5,0].tolist()}') - -if __name__ == '__main__': - test_diag_v() diff --git a/tests/archive/test_diag_multitile.py b/tests/archive/test_diag_multitile.py deleted file mode 100644 index d0d974bc..00000000 --- a/tests/archive/test_diag_multitile.py +++ /dev/null @@ -1,30 +0,0 @@ -"""Test the identity diag for multi-tile n=256,384""" -import torch, math, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline, cutlass.torch as ct, cuda.bindings.driver as cuda -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean -from cutlass.utils import LayoutEnum -from test_fmha_v3_diag import FmhaV3Diag - -HEAD_DIM = 64 -for n in [128, 256, 384]: - torch.manual_seed(42) - q = torch.randn(128, HEAD_DIM, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(n, HEAD_DIM, 1, dtype=torch.bfloat16, device='cuda') - v = torch.ones(n, HEAD_DIM, dtype=torch.bfloat16, device='cuda') - v_kernel = v.unsqueeze(-1) - c = torch.zeros(128, HEAD_DIM, 1, dtype=torch.bfloat16, device='cuda') - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - kernel = FmhaV3Diag(s_k=n) - print(f'n={n}: Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - out = c[:,:,0].float() - qf = q[:,:,0].float(); kf = k[:,:,0].float() - ref = (qf @ kf.T * (1.0/math.sqrt(HEAD_DIM))) @ v.float() - cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - print(f'Identity diag n={n}: cos {cos:.6f} {"PASS" if cos >= 0.99 else "FAIL"}') diff --git a/tests/archive/test_diag_permute.py b/tests/archive/test_diag_permute.py deleted file mode 100644 index 9662fd45..00000000 --- a/tests/archive/test_diag_permute.py +++ /dev/null @@ -1,80 +0,0 @@ -""" -Quick diagnostic: truncated identity V with 128x64 PV. -Check if output columns match a permutation of reference columns. -If O[m,d] = P[m, perm(d)], then the PV MMA is reading P from wrong TMEM addresses. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05, OperandMajorMode -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - -# Reuse the DiagVTruncIdKernel from test_diag_v_truncid.py -# (just run it and do more analysis on the output) - -# Actually, let me just re-run the truncid test and do the permutation analysis in Python -# First run the kernel, then analyze - -# We already ran it and have the results. Let me just do the analysis with the numbers we have. -# O[0,:5] = [6.0625, 11.875, -9.5625, -4.6875, -14.9375] -# ref[0,:5] = [6.0625, 10.5625, 11.875, -11.75, -9.5625] -# P[0] (full Q@K^T row 0) needs to be computed - -torch.manual_seed(42) -m, n, head_dim = 128, 128, 64 -q = torch.randn(m, head_dim, 1, dtype=torch.bfloat16, device='cuda') -k = torch.randn(n, head_dim, 1, dtype=torch.bfloat16, device='cuda') - -qf = q[:,:,0].float() -kf = k[:,:,0].float() -P = (qf @ kf.T) # (128, 128) — the P matrix - -# Now check: does O[0, d] = P[0, perm(d)] for some permutation? -# O[0,0] = 6.0625 → matches P[0,0] = 6.0625 -# O[0,1] = 11.875 → matches P[0,2] = 11.875 -# O[0,2] = -9.5625 → matches P[0,4] = -9.5625 -# So O[0, d] = P[0, 2*d]? Let me check more. - -O_row0 = [6.0625, 11.875, -9.5625, -4.6875, -14.9375] -P_row0 = P[0, :10].tolist() -print(f"P[0, :10] = {P_row0}") -print(f"O[0, :5] = {O_row0}") - -# Check: O[0, d] = P[0, 2*d]? -for d in range(5): - print(f" O[0,{d}] = {O_row0[d]:.4f}, P[0,{2*d}] = {P_row0[2*d]:.4f}, match = {abs(O_row0[d] - P_row0[2*d]) < 0.01}") - -# Also check full P row 0 vs O -# We can't get O without running the kernel again, but the pattern is clear: -# O[m, d] = P[m, 2*d] for the truncated identity V case -# This means the PV MMA is reading P from every other TMEM column - -# Why 2*d? Because with (128,64) MMA, the A fragment reads TMEM with stride 2 in the K dimension. -# The (128,64,16) MMA atom has N=64, which means it reads 64 columns of P per K-tile -# But P has 128 columns. The MMA reads the first 64, but with the wrong stride. -# -# Actually, with (128,64,16) MMA: -# - A operand: (M=128, K=128) → MMA reads 128/16 = 8 K-tiles -# - Each K-tile reads P[:, k*16:(k+1)*16] = 16 columns of P -# - The A fragment for K-tile kb reads from TMEM column offset based on N_MMA -# -# The (128,64,16) MMA's TMEM A fragment layout might be: -# (128, N_MMA) where N_MMA relates to the N dimension of the MMA -# If N_MMA = 64 (half of 128), then P's 128 BF16 values in K are stored -# in 128 BF16 TMEM columns = 64 FP32 TMEM columns -# But the (128,64,16) A fragment might only address 32 FP32 TMEM columns -# because the MMA only uses 64 columns for the C output -# So P's 128 K values don't fit in 32 TMEM columns, and the layout is different - -# The root cause: the (128,64) MMA's A fragment in TMEM packs 128 BF16 K values -# into fewer TMEM columns than the (128,128) MMA. The softmax packing writes P -# using the (128,128) layout, but the PV MMA reads with the (128,64) layout. - -print("\n=== HYPOTHESIS ===") -print("The (128,64,16) MMA atom reads P from TMEM with a DIFFERENT layout") -print("than the softmax packing writes P with (QK C fragment layout).") -print("The (128,128,16) MMA atom's A fragment layout matches the QK C fragment layout,") -print("so the 128x128 case works. The (128,64,16) layout differs, causing the bug.") -print("Fix: softmax packing should write P using the PV MMA's A fragment layout.") diff --git a/tests/archive/test_diag_smem_layout.py b/tests/archive/test_diag_smem_layout.py deleted file mode 100644 index d1209e8f..00000000 --- a/tests/archive/test_diag_smem_layout.py +++ /dev/null @@ -1,72 +0,0 @@ -"""Print V SMEM layouts for (128,64) and (128,128) PV. Must run inside JIT.""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05, OperandMajorMode -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - -class SmemLayoutKernel: - def __init__(self): - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.qk_acc_dtype = Float32 - self.use_2cta_instrs = False; self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.num_c_stage = 2 - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type - a_major = LayoutEnum.from_tensor(q).mma_major_mode() - b_major = LayoutEnum.from_tensor(k).mma_major_mode() - v_major = LayoutEnum.from_tensor(v).mma_major_mode() - c_layout = LayoutEnum.from_tensor(c) - - # QK - qk_mma = utils.sm100.make_trivial_tiled_mma( - BFloat16, BFloat16, a_major, b_major, - Float32, tcgen05.CtaGroup.ONE, (128, 128), tcgen05.OperandSource.SMEM) - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - qk_mma_tiler = (128, 128, qk_inst_k * 4) - b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, qk_mma_tiler, BFloat16, 1) - print(f"QK B SMEM: outer={b_smem_s.outer}, inner={b_smem_s.inner}") - - # PV (128, 64) - pv_mma_64 = utils.sm100.make_trivial_tiled_mma( - BFloat16, BFloat16, OperandMajorMode.K, v_major, - Float32, tcgen05.CtaGroup.ONE, (128, 64), tcgen05.OperandSource.TMEM) - pv_mma_tiler_64 = (128, 64, 128) - v_smem_64 = utils.sm100.make_smem_layout_b(pv_mma_64, pv_mma_tiler_64, BFloat16, 1) - print(f"PV(128,64) V SMEM: outer={v_smem_64.outer}, inner={v_smem_64.inner}") - - # PV (128, 128) - pv_mma_128 = utils.sm100.make_trivial_tiled_mma( - BFloat16, BFloat16, OperandMajorMode.K, v_major, - Float32, tcgen05.CtaGroup.ONE, (128, 128), tcgen05.OperandSource.TMEM) - pv_mma_tiler_128 = (128, 128, 128) - v_smem_128 = utils.sm100.make_smem_layout_b(pv_mma_128, pv_mma_tiler_128, BFloat16, 1) - print(f"PV(128,128) V SMEM: outer={v_smem_128.outer}, inner={v_smem_128.inner}") - - # Also print the PV MMA atom shapes - print(f"PV(128,64) MMA shape_mnk={pv_mma_64.shape_mnk}") - print(f"PV(128,128) MMA shape_mnk={pv_mma_128.shape_mnk}") - -torch.manual_seed(42) -m, n, head_dim = 128, 128, 64 -q = torch.randn(m, head_dim, 1, dtype=torch.bfloat16, device='cuda') -k = torch.randn(n, head_dim, 1, dtype=torch.bfloat16, device='cuda') -v_data = torch.zeros(head_dim, n, dtype=torch.bfloat16, device='cuda') -v_data[0, 0] = 1.0 -v = v_data.as_strided((head_dim, n), (1, head_dim)).unsqueeze(-1) -c = torch.zeros(m, head_dim, 1, dtype=torch.bfloat16, device='cuda') - -mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) -mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) -mV = ct.from_dlpack(v).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v)) -mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) -stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) -kernel = SmemLayoutKernel() -compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) diff --git a/tests/archive/test_diag_v_mma128.py b/tests/archive/test_diag_v_mma128.py deleted file mode 100644 index cf0aa6d4..00000000 --- a/tests/archive/test_diag_v_mma128.py +++ /dev/null @@ -1,374 +0,0 @@ -""" -Diagnostic: PV with (128,64) output but using (128,128) MMA for PV. -This keeps the A fragment (P) TMEM layout the same as the softmax packing path. -V is truncated identity (64,128) MN-major. If cosine ~0.999, confirms TMEM alias mismatch. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05, OperandMajorMode -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class DiagVMma128Kernel: - def __init__(self, head_dim): - self.head_dim = head_dim - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = (128, 128); self.mma_tiler = (128, 128, 1) - self.use_2cta_instrs = False - self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (128, 128, qk_inst_k * 4) - # PV uses (128, 128) MMA even though output is (128, 64) - self.pv_mma_tiler = (128, 128, 128) # Same as 128x128 case - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, self.use_2cta_instrs, self.c_layout, self.o_dtype) - print(f"[SETUP] qk_mma_tiler={self.qk_mma_tiler}, pv_mma_tiler={self.pv_mma_tiler}") - print(f"[SETUP] epi_tile={self.epi_tile}") - - self.num_ab_stage = 1; self.num_acc_stage = 1 - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_column_offset(tStS) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, (128, 128), tcgen05.OperandSource.SMEM) - # PV MMA (128, 128) - same as QK, so A fragment layout matches softmax packing - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, (128, 128), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - # PV (128,128) output partitioned with pv_thr — c tensor is (128,64) but we write full (128,128) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = pv_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # TMA LOAD WARP - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - # Check that last 64 cols are ~0 - out_last64 = out[:, head_dim:] - last64_max = out_last64.abs().max().item() - print('Last 64 cols max abs value: {:.6f} (should be ~0)'.format(last64_max)) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_diag_v_ones.py b/tests/archive/test_diag_v_ones.py deleted file mode 100644 index 583cb85b..00000000 --- a/tests/archive/test_diag_v_ones.py +++ /dev/null @@ -1,62 +0,0 @@ -""" -Diagnostic: PV with V = all ones (64, 128) MN-major. -O[m, d] should be sum_k P[m, k] for all d (all columns identical). -If columns differ, the PV MMA is reading V incorrectly. -Also test: PV with V = single element (d=0, k=0) = 1, rest 0. -O[m, 0] = P[m, 0], all other entries 0. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05, OperandMajorMode -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - -# Reuse the same kernel as test_diag_v_truncid.py but with different V - -def run_test(v_data, desc, head_dim=64): - m, n = 128, 128 - q = torch.randn(m, head_dim, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(n, head_dim, 1, dtype=torch.bfloat16, device='cuda') - v = v_data.as_strided((head_dim, n), (1, head_dim)).unsqueeze(-1) - c = torch.zeros(m, head_dim, 1, dtype=torch.bfloat16, device='cuda') - - qf = q[:,:,0].float(); kf = k[:,:,0].float(); vf = v_data.float() - ref = (qf @ kf.T).bfloat16().float() @ vf.T - - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - # Import the kernel from test_diag_v_truncid - from test_diag_v_truncid import DiagVTruncIdKernel - kernel = DiagVTruncIdKernel(mma_tiler_mn=(128, 128), head_dim=head_dim) - print(f'Compiling {desc}...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - print(f'Running {desc}...', flush=True) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - out = c[:,:,0].float() - cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - print(f'{desc}: cosine {cos:.6f} {"PASS" if cos >= 0.99 else "FAIL"}') - - if cos < 0.99: - print(f' O[0,:5] = {out[0,:5].tolist()}') - print(f' ref[0,:5] = {ref[0,:5].tolist()}') - # Check if columns are all the same (for all-ones V) - if v_data.abs().sum() > 0: - col0 = out[:, 0] - all_same = all(torch.allclose(out[:, d], col0, atol=0.1) for d in range(min(8, head_dim))) - print(f' All columns same? {all_same}') - -# Test 1: All ones V -v_ones = torch.ones(64, 128, dtype=torch.bfloat16, device='cuda') -run_test(v_ones, "All-ones V") - -# Test 2: Single element V -v_single = torch.zeros(64, 128, dtype=torch.bfloat16, device='cuda') -v_single[0, 0] = 1.0 -run_test(v_single, "Single-element V") diff --git a/tests/archive/test_diag_v_truncid.py b/tests/archive/test_diag_v_truncid.py deleted file mode 100644 index 78f05939..00000000 --- a/tests/archive/test_diag_v_truncid.py +++ /dev/null @@ -1,369 +0,0 @@ -""" -Diagnostic: PV with truncated identity V (64,128). -V[d,k] = 1 if d==k, 0 otherwise. MN-major. -With identity softmax: O = P[:, :64] = (Q@K^T).bfloat16()[:, :64]. -If PV MMA is correct, cosine should be ~0.999. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05, OperandMajorMode -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class DiagVTruncIdKernel: - def __init__(self, mma_tiler_mn, head_dim): - self.head_dim = head_dim - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False - self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - self.pv_mma_tiler = (self.qk_mma_tiler[0], self.qk_mma_tiler[2], self.qk_mma_tiler[1]) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.qk_cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - # PV cta tile for epilogue - self.pv_cta_tile_shape_mnk = ( - self.pv_mma_tiler[0] // cute.size(pv_mma.thr_id.shape), - self.pv_mma_tiler[1], self.pv_mma_tiler[2]) - self.cta_tile_shape_mnk = self.pv_cta_tile_shape_mnk - - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.pv_cta_tile_shape_mnk, self.use_2cta_instrs, self.c_layout, self.o_dtype) - print(f"[SETUP] qk_mma_tiler={self.qk_mma_tiler}, pv_mma_tiler={self.pv_mma_tiler}") - print(f"[SETUP] epi_tile={self.epi_tile}") - - self.num_ab_stage = 1; self.num_acc_stage = 1 - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - pv_mma_tiler_mn = (self.mma_tiler_mn[0], self.head_dim) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, pv_mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage *2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.pv_mma_tiler, (None,0,None)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = pv_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # TMA LOAD WARP - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - - if cos < 0.99: - print('\n=== Element-level ===') - for m_idx in [0, 1, 63, 127]: - for d_idx in [0, 1, 31, 63]: - print(f' O[{m_idx},{d_idx}] = {out[m_idx,d_idx]:.4f}, ref = {ref[m_idx,d_idx]:.4f}') - print(f' O[0,:5] = {out[0,:5].tolist()}') - print(f' ref[0,:5] = {ref[0,:5].tolist()}') - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_e2e_decode_b200.py b/tests/archive/test_e2e_decode_b200.py deleted file mode 100644 index f5096c8d..00000000 --- a/tests/archive/test_e2e_decode_b200.py +++ /dev/null @@ -1,425 +0,0 @@ -#!/usr/bin/env python3 -""" -DeepSeek-V4 End-to-End Decode Test - -Generates actual tokens using our KV cache pipeline: -1. Prefill: process N tokens through all 61 layers, write KV to paged cache -2. Decode: generate tokens one at a time using cached KV -3. Verify: check that generated tokens form coherent text (not garbage) - -This is the test that MUST pass before we touch the vLLM container. - -Usage (on B200): - cd /root/nvfp4-megamoe-kernel - PYTHONPATH=/root/nvfp4-megamoe-kernel tests/venv/bin/python tests/test_e2e_decode_b200.py -""" - -import sys, os, json, torch, torch.nn.functional as F, time -from safetensors import safe_open - -REPO = "/root/nvfp4-megamoe-kernel" -sys.path.insert(0, REPO) -MODEL = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4" -DEV = "cuda:0" - -H = 7168; NH = 128; HD = 512; NOPE = 448; ROPE = 64 -QL = 1536; OL = 1024; OG = 16; HPG = NH // OG -EPS = 1e-6; WINDOW = 128; SCALE = HD ** -0.5 -NUM_LAYERS = 61 -NUM_TEST_LAYERS = 3 # Test with 3 layers first (0, 1, 60 = C128A, C4A, SWA) -NUM_EXPERTS = 384; TOPK = 6 -VOCAB = 129024 - -E2M1 = torch.tensor([0,.5,1.,1.5,2.,3.,4.,6.,-0,-.5,-1.,-1.5,-2.,-3.,-4.,-6.], dtype=torch.float32) - -_cache = {} -def P(k, wm, md): - if k in _cache: return _cache[k] - with safe_open(os.path.join(md, wm[k]), framework="pt") as f: - t = f.get_tensor(k) - _cache[k] = t - return t - -def rms(x, w, eps=1e-6): - v = x.float().pow(2).mean(-1, keepdim=True) - return (w.float() * (x * torch.rsqrt(v+eps)).float()).to(x.dtype) - -def make_runner(w, sf, gs_t, inf, outf, fused=False, lw=None): - from dsv4.layers.linear import Nvfp4Linear - fp4 = w.view(torch.float4_e2m1fn_x2).permute(1,0).contiguous() - s = sf.to(torch.float8_e4m3fn) if sf.dtype != torch.float8_e4m3fn else sf - s = s.permute(1,0).contiguous() - if fused and gs_t.numel() == 2: - g1,g2 = gs_t[0].item(), gs_t[1].item(); gs = max(g1,g2) - if g1 != g2: - s32 = s.float(); sp = lw[0] if lw else outf//2 - s32[:sp] *= g1/gs; s32[sp:] *= g2/gs; s = s32.to(torch.float8_e4m3fn) - else: - gs = gs_t.max().item() if gs_t.numel() > 1 else gs_t.item() - r = Nvfp4Linear(in_features=inf, out_features=outf, max_num_tokens=8192, device=str(w.device)) - r.fp4 = [fp4]; r.sf = [s]; r.gs = [gs] - r.finalize_weights(); r._ensure_initialized() - return r - -def build_cos_sin(max_pos=4096, rope_dim=ROPE): - half = rope_dim // 2 - inv_freq = 1.0 / (10000.0 ** (torch.arange(0, half, dtype=torch.float32) / half)) - freqs = torch.outer(torch.arange(max_pos, dtype=torch.float32), inv_freq) - return torch.cat([freqs.cos(), freqs.sin()], dim=-1) - -def apply_gptj_rope(x, positions, cos_sin, nope_dim, rope_dim): - if rope_dim == 0 or x.numel() == 0: return x - half = rope_dim // 2 - cos = cos_sin[positions, :half].to(x.dtype) - sin = cos_sin[positions, half:2*half].to(x.dtype) - if x.dim() == 3: cos = cos.unsqueeze(1); sin = sin.unsqueeze(1) - x_rope = x[..., nope_dim:].clone() - even = x_rope[..., 0::2]; odd = x_rope[..., 1::2] - out = x.clone() - out[..., nope_dim:][..., 0::2] = even * cos - odd * sin - out[..., nope_dim:][..., 1::2] = even * sin + odd * cos - return out - -def apply_inv_gptj_rope(x, positions, cos_sin, nope_dim, rope_dim): - if rope_dim == 0 or x.numel() == 0: return x - half = rope_dim // 2 - cos = cos_sin[positions, :half].to(x.dtype) - sin = cos_sin[positions, half:2*half].to(x.dtype) - if x.dim() == 3: cos = cos.unsqueeze(1); sin = sin.unsqueeze(1) - x_rope = x[..., nope_dim:].clone() - even = x_rope[..., 0::2]; odd = x_rope[..., 1::2] - out = x.clone() - out[..., nope_dim:][..., 0::2] = even * cos + odd * sin - out[..., nope_dim:][..., 1::2] = -even * sin + odd * cos - return out - -# ── KV Cache ───────────────────────────────────────────────────────── - -def kv_quantize_fp8(kv_bf16): - amax = kv_bf16.float().abs().amax(dim=-1, keepdim=True).clamp(min=1e-12) - fp8_max = torch.tensor(448.0, dtype=torch.float32, device=kv_bf16.device) - scale = fp8_max / amax - kv_fp8 = (kv_bf16.float() * scale).to(torch.float8_e4m3fn) - inv_scale = (amax / fp8_max).to(torch.bfloat16) - return kv_fp8, inv_scale - -def kv_dequantize_fp8(kv_fp8, inv_scale): - return (kv_fp8.to(torch.bfloat16) * inv_scale).to(torch.bfloat16) - -def paged_kv_write(kv_data, slot_mapping, cache, inv_scale_cache, block_size): - if cache.dtype == torch.uint8 and kv_data.dtype == torch.float8_e4m3fn: - kv_to_write = kv_data.view(torch.uint8) - else: - kv_to_write = kv_data - block_indices = slot_mapping // block_size - offsets = slot_mapping % block_size - cache[block_indices, offsets] = kv_to_write - # Write inv_scale - for t in range(kv_data.shape[0]): - inv_scale_cache[slot_mapping[t].item()] = kv_data # placeholder - # Actually write inv_scale per-token - if hasattr(inv_scale_cache, '__setitem__'): - for t in range(kv_data.shape[0]): - inv_scale_cache[slot_mapping[t].item()] = ... # need the inv_scale tensor - -def paged_kv_read(slot_mapping, cache, inv_scale_cache, block_size, num_tokens, head_dim): - block_indices = slot_mapping // block_size - offsets = slot_mapping % block_size - kv = cache[block_indices, offsets] - if cache.dtype == torch.uint8: - kv = kv.view(torch.float8_e4m3fn) - # Read inv_scale - inv_scales = inv_scale_cache[slot_mapping] # (T, 1) - return kv, inv_scales - - -# ── Attention ───────────────────────────────────────────────────────── - -def causal_prefill_attention(q, kv, scale): - T, NH, HD = q.shape - q_t = q.permute(1, 0, 2) - kv_exp = kv.unsqueeze(0).expand(NH, -1, -1) - out = F.scaled_dot_product_attention(q_t, kv_exp, kv_exp, is_causal=True, scale=scale) - return out.permute(1, 0, 2) - -def decode_attention(q, kv, scale): - NH = q.shape[1]; HD = q.shape[2] - q_t = q.permute(1, 0, 2) - kv_exp = kv.unsqueeze(0).expand(NH, -1, -1) - out = F.scaled_dot_product_attention(q_t, kv_exp, kv_exp, is_causal=False, scale=scale) - return out.permute(1, 0, 2) - - -# ── Layer type mapping ──────────────────────────────────────────────── - -def get_layer_type(layer_id): - """Return (compress_ratio, has_compressor) for each layer.""" - if layer_id == 60: - return 0, False # SWA (last layer) - if layer_id == 0: - return 128, True # HCA (C128A) - return 4, True # CSA (C4A) — most layers - - -def run_layer(hidden, layer_id, runners, weights, cos_sin, positions, - kv_caches, inv_scale_caches, block_size, is_prefill=True): - """Run one transformer layer. Returns updated hidden states. - - Writes KV to the paged cache. Uses cache for decode, raw KV for prefill. - """ - p = f"model.layers.{layer_id}" - a = f"{p}.self_attn" - - r_qa = runners[layer_id]['qa'] - r_qb = runners[layer_id]['qb'] - r_kv = runners[layer_id]['kv'] - r_wob = runners[layer_id]['wob'] - woa = weights[layer_id]['woa'] - qn_w = weights[layer_id]['qn'] - kvn_w = weights[layer_id]['kvn'] - anorm_w = weights[layer_id]['anorm'] - fnorm_w = weights[layer_id]['fnorm'] - - NT = hidden.shape[0] - - # ── Attention ────────────────────────────────────────────── - normed = rms(hidden, anorm_w, EPS) - qa = r_qa.run(normed) - kv = r_kv.run(normed) - qa_n = rms(qa, qn_w, EPS) - kv_n = rms(kv, kvn_w, EPS) - q = r_qb.run(qa_n).view(NT, NH, HD) - q_rope = apply_gptj_rope(q, positions, cos_sin, NOPE, ROPE) - kv_rope = apply_gptj_rope(kv_n.unsqueeze(1), positions, cos_sin, NOPE, ROPE).squeeze(1) - - # Write KV to paged cache - kv_fp8, kv_inv_s = kv_quantize_fp8(kv_rope) - slots = positions # slot = position (simplified) - block_indices = slots // block_size - offsets = slots % block_size - cache = kv_caches[layer_id] - inv_sc = inv_scale_caches[layer_id] - if cache.dtype == torch.uint8: - cache[block_indices, offsets] = kv_fp8.view(torch.uint8) - else: - cache[block_indices, offsets] = kv_fp8 - for t in range(NT): - inv_sc[slots[t].item()] = kv_inv_s[t] - - # Attention - if is_prefill: - o_attn = causal_prefill_attention(q_rope, kv_rope, SCALE) - else: - # Decode: read ALL cached KV from position 0 to current - pos = positions[0].item() - all_slots = torch.arange(pos + 1, dtype=torch.int64, device=DEV) - all_bi = all_slots // block_size - all_oi = all_slots % block_size - kv_cached_fp8 = cache[all_bi, all_oi] - if cache.dtype == torch.uint8: - kv_cached_fp8 = kv_cached_fp8.view(torch.float8_e4m3fn) - kv_cached_inv = inv_sc[all_slots] - kv_cached = kv_dequantize_fp8(kv_cached_fp8, kv_cached_inv) - # SWA window - window_start = max(0, pos - WINDOW + 1) - kv_window = kv_cached[window_start:] - o_attn = decode_attention(q_rope, kv_window, SCALE) - - # Output projection: inverse RoPE + o_a BMM + o_b - o_inv = apply_inv_gptj_rope(o_attn, positions, cos_sin, NOPE, ROPE) - o_grouped = o_inv.reshape(NT, OG, HPG * HD).permute(1, 0, 2) - woa_3d = woa.view(OG, OL, HPG * HD) - z = torch.bmm(o_grouped, woa_3d.transpose(1, 2)).permute(1, 0, 2).reshape(NT, OG * OL) - attn_out = r_wob.run(z) - - hidden = hidden + attn_out - - # ── MoE (shared expert only for speed) ───────────────────── - fnormed = rms(hidden, fnorm_w, EPS) - - r_se_gate = runners[layer_id]['se_gate'] - r_se_up = runners[layer_id]['se_up'] - r_se_down = runners[layer_id]['se_down'] - - gate_out = r_se_gate.run(fnormed) - up_out = r_se_up.run(fnormed) - se_activated = F.silu(gate_out) * up_out - se_final = r_se_down.run(se_activated) - hidden = hidden + se_final - - return hidden - - -def main(): - print("=" * 70) - print(" DeepSeek-V4 End-to-End Decode Test") - print(" Prefill → KV Cache → Decode → Generate Tokens") - print("=" * 70) - - torch.cuda.set_device(0) - torch.manual_seed(42) - - with open(os.path.join(MODEL, "model.safetensors.index.json")) as f: - wm = json.load(f)["weight_map"] - G = lambda k: P(k, wm, MODEL).to(DEV) - - # Load shared weights - emb = G("model.embed_tokens.weight") - lm_head = G("lm_head.weight") - fnorm_w = G("model.norm.weight") - cos_sin = build_cos_sin(max_pos=4096).to(DEV) - - # ── Load per-layer weights and create runners ────────────── - print("\n Loading weights and creating runners...") - runners = {} - weights = {} - - # Test with all 61 layers (shared experts only) - test_layers = list(range(NUM_LAYERS)) - - for layer_id in test_layers: - p = f"model.layers.{layer_id}" - a = f"{p}.self_attn" - m = f"{p}.mlp" - - # Attention weights - qa_w = G(f"{a}.q_a_proj.weight"); qa_sf = G(f"{a}.q_a_proj.weight_scale"); qa_gs = G(f"{a}.q_a_proj.weight_scale_2") - qb_w = G(f"{a}.q_b_proj.weight"); qb_sf = G(f"{a}.q_b_proj.weight_scale"); qb_gs = G(f"{a}.q_b_proj.weight_scale_2") - kv_w = G(f"{a}.kv_proj.weight"); kv_sf = G(f"{a}.kv_proj.weight_scale"); kv_gs = G(f"{a}.kv_proj.weight_scale_2") - wob_w = G(f"{a}.o_b_proj.weight"); wob_sf = G(f"{a}.o_b_proj.weight_scale"); wob_gs = G(f"{a}.o_b_proj.weight_scale_2") - woa = G(f"{a}.o_a_proj.weight") - qn = G(f"{a}.q_a_norm.weight") - kvn = G(f"{a}.kv_norm.weight") - anorm = G(f"{p}.input_layernorm.weight") - fnorm = G(f"{p}.post_attention_layernorm.weight") - - # Shared expert weights (separate gate_proj + up_proj + down_proj) - se_gate_w = G(f"{m}.shared_experts.gate_proj.weight"); se_gate_sf = G(f"{m}.shared_experts.gate_proj.weight_scale"); se_gate_gs = G(f"{m}.shared_experts.gate_proj.weight_scale_2") - se_up_w = G(f"{m}.shared_experts.up_proj.weight"); se_up_sf = G(f"{m}.shared_experts.up_proj.weight_scale"); se_up_gs = G(f"{m}.shared_experts.up_proj.weight_scale_2") - se_down_w = G(f"{m}.shared_experts.down_proj.weight"); se_down_sf = G(f"{m}.shared_experts.down_proj.weight_scale"); se_down_gs = G(f"{m}.shared_experts.down_proj.weight_scale_2") - - r_qa = make_runner(qa_w, qa_sf, qa_gs, H, qa_w.shape[0]) - r_qb = make_runner(qb_w, qb_sf, qb_gs, QL, qb_w.shape[0]) - r_kv = make_runner(kv_w, kv_sf, kv_gs, H, kv_w.shape[0]) - r_wob = make_runner(wob_w, wob_sf, wob_gs, OG*OL, wob_w.shape[0]) - r_se_gate = make_runner(se_gate_w, se_gate_sf, se_gate_gs, H, se_gate_w.shape[0]) - r_se_up = make_runner(se_up_w, se_up_sf, se_up_gs, H, se_up_w.shape[0]) - r_se_down = make_runner(se_down_w, se_down_sf, se_down_gs, 3072, se_down_w.shape[0]) - - runners[layer_id] = { - 'qa': r_qa, 'qb': r_qb, 'kv': r_kv, 'wob': r_wob, - 'se_gate': r_se_gate, 'se_up': r_se_up, 'se_down': r_se_down, - } - weights[layer_id] = { - 'woa': woa, 'qn': qn, 'kvn': kvn, - 'anorm': anorm, 'fnorm': fnorm, - } - - if layer_id % 10 == 0: - print(f" Layer {layer_id} loaded") - - # ── Allocate KV caches ───────────────────────────────────── - block_size = 64 # Match vLLM's SWA cache block size - max_tokens = 256 - num_blocks = (max_tokens + block_size - 1) // block_size - - kv_caches = {} - inv_scale_caches = {} - for layer_id in test_layers: - kv_caches[layer_id] = torch.zeros(num_blocks, block_size, HD, dtype=torch.uint8, device=DEV) - inv_scale_caches[layer_id] = torch.zeros(max_tokens, 1, dtype=torch.bfloat16, device=DEV) - - print(f"\n KV caches allocated: {NUM_LAYERS} layers × {num_blocks} blocks × {block_size} tokens × {HD} dims") - - # ── PREFILL ──────────────────────────────────────────────── - print(f"\n === PREFILL ===") - prompt = "The capital of France is" - # Tokenize manually (use simple wordpiece-style IDs for testing) - # For a real test, we'd use the tokenizer, but this works for verifying the pipeline - token_ids = torch.tensor([1, 450, 8403, 315, 5413, 374], dtype=torch.long, device=DEV) - positions = torch.arange(len(token_ids), dtype=torch.int64, device=DEV) - - hidden = emb[token_ids] - print(f" Input: {len(token_ids)} tokens") - - t0 = time.time() - with torch.no_grad(): - for layer_id in test_layers: - hidden = run_layer(hidden, layer_id, runners, weights, cos_sin, positions, - kv_caches, inv_scale_caches, block_size, is_prefill=True) - if layer_id % 10 == 0: - print(f" Layer {layer_id}: amax={hidden.amax():.4f} NaN={torch.isnan(hidden).any()}") - - # Final norm + LM head - hidden = rms(hidden, fnorm_w, EPS) - logits = hidden @ lm_head.T - t1 = time.time() - - print(f" Prefill time: {t1-t0:.2f}s") - print(f" Logits: amax={logits.amax():.4f} std={logits[-1].float().std():.4f}") - top5 = torch.topk(logits[-1], 5) - print(f" Top 5 tokens: {top5.indices.tolist()}") - - # ── DECODE ───────────────────────────────────────────────── - print(f"\n === DECODE (5 tokens) ===") - - # Sample first decode token - next_token = top5.indices[0].unsqueeze(0) # Greedy - generated = [next_token.item()] - current_pos = len(token_ids) - - for step in range(5): - pos = torch.tensor([current_pos], dtype=torch.int64, device=DEV) - hidden = emb[next_token] - - with torch.no_grad(): - for layer_id in test_layers: - hidden = run_layer(hidden, layer_id, runners, weights, cos_sin, pos, - kv_caches, inv_scale_caches, block_size, is_prefill=False) - - hidden = rms(hidden, fnorm_w, EPS) - logits = hidden @ lm_head.T - - next_token = logits[-1].argmax().unsqueeze(0) - generated.append(next_token.item()) - current_pos += 1 - - log_std = logits[-1].float().std().item() - print(f" Step {step}: token={next_token.item()} logit_std={log_std:.4f} {'✅' if 0.5 < log_std < 50 else '❌'}") - - print(f"\n Generated tokens: {generated}") - print(f" Logit check: {'✅ All reasonable' if all(0.5 < 1 < 50 for _ in generated) else '❌ Check for NaN/garbage'}") - - # ── Verification: decode with cache should match full prefill ── - print(f"\n === VERIFICATION: decode vs full prefill ===") - # Run all tokens at once (prefill) and compare the last token's logits - all_tokens = torch.cat([token_ids, torch.tensor(generated[:-1], dtype=torch.long, device=DEV)]) - all_positions = torch.arange(len(all_tokens), dtype=torch.int64, device=DEV) - - # Reset caches - for layer_id in test_layers: - kv_caches[layer_id].zero_() - inv_scale_caches[layer_id].zero_() - - hidden_ref = emb[all_tokens] - with torch.no_grad(): - for layer_id in test_layers: - hidden_ref = run_layer(hidden_ref, layer_id, runners, weights, cos_sin, all_positions, - kv_caches, inv_scale_caches, block_size, is_prefill=True) - - hidden_ref = rms(hidden_ref, fnorm_w, EPS) - logits_ref = hidden_ref @ lm_head.T - - # Compare the decode token's logits - # (This isn't a perfect comparison because decode uses fp8 cached KV vs prefill uses raw KV, - # but cosine should be > 0.95) - # We'd need to re-run decode to get the exact comparison, but the logit std check above - # is sufficient to verify the pipeline works. - - print(f"\n{'='*70}") - print(f" DONE") - print(f"{'='*70}") - - -if __name__ == "__main__": - main() diff --git a/tests/archive/test_error_pattern.py b/tests/archive/test_error_pattern.py deleted file mode 100644 index 3c7632df..00000000 --- a/tests/archive/test_error_pattern.py +++ /dev/null @@ -1,85 +0,0 @@ -""" -BF16 Packing Diagnostic: Run identity softmax with K=V=randn, -compare output vs reference to identify the error pattern. -""" -import torch, sys -sys.path.insert(0, '/root/dsv4-nvfp4-workspace/kernel/tests') - -import cutlass.cute as cute -import cutlass.torch as ct -import cuda.bindings.driver as cuda -from test_stage_b_v7 import StageBIdentitySoftmax - -torch.manual_seed(42) -m, n, k = 128, 128, 128 -q = torch.randn(m, k, 1, dtype=torch.bfloat16, device='cuda') -kv = torch.randn(n, k, 1, dtype=torch.bfloat16, device='cuda') -c = torch.zeros(m, n, 1, dtype=torch.bfloat16, device='cuda') - -qf = q[:,:,0].float() -kvf = kv[:,:,0].float() -ref = qf @ kvf.T @ kvf # identity softmax: (Q @ K^T) @ V - -mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) -mK = ct.from_dlpack(kv).mark_layout_dynamic(leading_dim=ct.get_leading_dim(kv)) -mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) -stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - -kernel = StageBIdentitySoftmax(mma_tiler_mn=(128, 128), use_2cta_instrs=False, use_tma_store=True) -print('Compiling...', flush=True) -compiled = cute.compile(kernel, mQ, mK, mC, stream) -print('Running...', flush=True) -compiled(mQ, mK, mC, stream) -torch.cuda.synchronize() - -out = c[:,:,0].float() -cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - -print(f'\nCosine: {cos:.6f}') -print(f'Output row 0[:8]: {out[0,:8].tolist()}') -print(f'Ref row 0[:8]: {ref[0,:8].tolist()}') - -# Key diagnostic: compare Q@K^T stage (which we know is correct) vs PV stage -# If Q@K^T is correct but PV is wrong, the output will show the PV error pattern -# With identity softmax, P = Q@K^T. So output = P @ V = (Q@K^T) @ V -# If V is being read wrong, the output will be P @ V_permuted - -# Check: is the output a permutation of the reference? -# Sort both and compare -out_sorted = out.flatten().sort()[0] -ref_sorted = ref.flatten().sort()[0] -cos_sorted = torch.nn.functional.cosine_similarity(out_sorted.unsqueeze(0), ref_sorted.unsqueeze(0)).item() -print(f'\nCosine after sorting: {cos_sorted:.6f}') -print(f'(If ~1.0, output is a permutation of reference)') - -# Check: does output match P @ V^T? (V transposed) -ref_vt = qf @ kvf.T @ kvf.T # (Q@K^T) @ V^T -cos_vt = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref_vt.flatten().unsqueeze(0)).item() -print(f'Cosine with P @ V^T: {cos_vt:.6f}') - -# Check: does output match P^T @ V? (P transposed) -# P = Q@K^T, so P^T = K@Q^T -ref_pt = kvf @ qf.T @ kvf -cos_pt = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref_pt.flatten().unsqueeze(0)).item() -print(f'Cosine with P^T @ V: {cos_pt:.6f}') - -# Check: is output simply the Q@K^T scores (MMA2 produced identity)? -# If MMA2 didn't run or produced P unchanged, output = P = Q@K^T -qkt = qf @ kvf.T -cos_qkt = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), qkt.flatten().unsqueeze(0)).item() -print(f'Cosine with just Q@K^T: {cos_qkt:.6f}') - -# Check: all output rows identical? (means P has identical rows, like all-ones) -all_same = torch.allclose(out[0], out[1], atol=1e-3) -print(f'All output rows identical: {all_same}') -if not all_same: - cos_r01 = torch.nn.functional.cosine_similarity(out[0].unsqueeze(0), out[1].unsqueeze(0)).item() - print(f'Cosine between row 0 and row 1: {cos_r01:.6f}') - -# Check: is output just V (P=I case)? -cos_v = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), kvf.flatten().unsqueeze(0)).item() -print(f'Cosine with V alone: {cos_v:.6f}') - -# Print some output statistics -print(f'\nOutput stats: min={out.min().item():.4f}, max={out.max().item():.4f}, mean={out.mean().item():.4f}') -print(f'Ref stats: min={ref.min().item():.4f}, max={ref.max().item():.4f}, mean={ref.mean().item():.4f}') diff --git a/tests/archive/test_fmha_pipeline.py b/tests/archive/test_fmha_pipeline.py deleted file mode 100644 index b08e39ab..00000000 --- a/tests/archive/test_fmha_pipeline.py +++ /dev/null @@ -1,354 +0,0 @@ -""" -Stage B — FMHA-style KV-tile interleaved attention kernel. - -Following CUTLASS FMHA reference architecture: -- Q: (seq_q, head_dim) — loaded once -- K, V: tiled over sequence dimension, V overwrites K in SMEM (FMHA trick) -- For each KV-tile: - 1. TMA load K[tile] into sK SMEM - 2. QK MMA: sQ @ sK^T → S in TMEM - 3. Softmax: S → P in TMEM (with online softmax rescaling of O in TMEM) - 4. V overwrites sK SMEM (after QK, K no longer needed) - 5. PV MMA: P @ sV → O in TMEM (accumulate) -- Epilogue: divide O by row_sum, store to GMEM - -This properly handles non-(128,128) PV because V SMEM always has the correct -data for the current KV-tile — it's loaded right before PV, not stale from -the beginning. - -Warp layout: - Warp 0-3: Softmax (4 warps) - Warp 4: MMA - Warp 5: TMA load -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class FmhaPipelineKernel: - def __init__(self, qk_mma_tiler, pv_mma_tiler): - self.acc_dtype = Float32 - self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16 - self.o_dtype = BFloat16 - self.c_dtype = BFloat16 - self.qk_mma_tiler = qk_mma_tiler - self.pv_mma_tiler = pv_mma_tiler - self.use_2cta_instrs = False - self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.softmax_warp_ids = (0, 1, 2, 3) - self.mma_warp_id = 4 - self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.kv_stage = 2 # double-buffered KV - self.q_stage = 1 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.qk_mma_tiler[:2], qk_inst_k * 4) - pv_inst_k = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (*self.pv_mma_tiler[:2], pv_inst_k * 4) - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide( - cute.make_layout((1, 1, 1)), (qk_mma.thr_id.shape,)) - self.epi_tile = self.pv_mma_tiler[:2] - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.pv_mma_tiler[1], - self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.num_ab_stage = 1 - self.num_acc_stage = 1 - self.num_c_stage = 2 - - self.q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, self.q_dtype, self.q_stage) - self.k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, self.q_dtype, self.kv_stage) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, self.kv_stage) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, self.num_c_stage) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = o_cols - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - a_smem = cute.slice_(self.q_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.k_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + - cute.size_in_bytes(self.q_dtype, b_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.qk_mma_tiler[:2], - tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, self.pv_mma_tiler[:2], - tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.q_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.k_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.qk_mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.qk_mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom( - cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel( - qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.q_smem_s, self.k_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1, 1, 1), block=[self.threads_per_cta, 1, 1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, q_smem_s, k_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q) - cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v) - cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - q_bar: cute.struct.MemRange[cutlass.Int64, self.q_stage * 2] - kv_bar: cute.struct.MemRange[cutlass.Int64, self.kv_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator() - st = smem.allocate(SS) - - q_prod, q_cons = pipeline.PipelineTmaUmma.create( - barrier_storage=st.q_bar.data_ptr(), num_stages=self.q_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - kv_prod, kv_cons = pipeline.PipelineTmaUmma.create( - barrier_storage=st.kv_bar.data_ptr(), num_stages=self.kv_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, 32 * len(self.softmax_warp_ids)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.softmax_warp_ids) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier( - barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.softmax_warp_ids))) - tmem = utils.TmemAllocator( - st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.softmax_warp_ids[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=q_smem_s.outer, byte_alignment=128, swizzle=q_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=k_smem_s.outer, byte_alignment=128, swizzle=k_smem_s.inner) - # V overwrites K SMEM (FMHA trick) - sV_ptr = cute.recast_ptr(sK.iterator, v_smem_s.inner) - sV = cute.make_tensor(sV_ptr, v_smem_s.outer) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None, 0, None)), (None, None, None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0, None, None)), (None, None, None)) - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0, None, None)), (None, None, None)) - gC = cute.local_tile(mC, cute.slice_(self.pv_mma_tiler, (None, None, 0)), (None, None, None)) - n_kv_tiles = cute.size(gK, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - - tCgQ = qk_thr.partition_A(gQ) - tCgK = qk_thr.partition_B(gK) - tCgV = pv_thr.partition_B(gV) - tCgC = pv_thr.partition_C(gC) - - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0, 0, None, 0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ, 0, 3), cute.group_modes(tCgQ, 0, 3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0, None, 0, 0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK, 0, 3), cute.group_modes(tCgK, 0, 3)) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV, 0, 3), cute.group_modes(tCgV, 0, 3)) - tAgQ = tAgQ[(None, 0, None, 0)] - tBgK = tBgK[(None, 0, None, 0)] - tVgV = tVgV[(None, 0, None, 0)] - - tCrQ = qk_mma.make_fragment_A(sQ) - tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - # Load Q once - q_prod.reset() - qh = q_prod.acquire_and_advance() - cute.copy(tma_q, tAgQ[(None, qh.count)], tAsQ[(None, qh.index)], tma_bar_ptr=qh.barrier) - q_prod.tail() - - # Load KV tiles: for each tile, load K then V - # K and V share SMEM, so V overwrites K after QK consumes it - kv_prod.reset() - peek = kv_prod.try_acquire() - for kt in cutlass.range(n_kv_tiles, unroll=1): - # Load K[tile] - kvh = kv_prod.acquire_and_advance(peek) - cute.copy(tma_k, tBgK[(None, kvh.count)], tBsK[(None, kvh.index)], tma_bar_ptr=kvh.barrier) - # Load V[tile] into the SAME SMEM (overwrites K after QK) - # Wait — we need QK to finish before V overwrites K. - # FMHA uses a SEPARATE pipeline entry for V. The MMA warp - # consumes K first (QK), then V (PV). The pipeline ordering - # ensures V doesn't overwrite K before QK is done. - cute.copy(tma_v, tVgV[(None, kvh.count)], tVsV[(None, kvh.index)], tma_bar_ptr=kvh.barrier) - peek = cutlass.Boolean(1) - if kvh.count + 1 < 2 * n_kv_tiles: - peek = kv_prod.try_acquire() - kv_prod.tail() - - # ═══ MMA WARP ═══ - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - - q_cons.reset() - qh = q_cons.wait_and_advance() - qh.release() - - kv_cons.reset() - peek = kv_cons.try_wait() - - acc_prod_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_stage) - acc_pipe.producer_acquire(acc_prod_st) - - for kt in range(n_kv_tiles): - # Wait for K[tile] - kvh = kv_cons.wait_and_advance(peek) - peek = cutlass.Boolean(1) - - # ─── QK: Q @ K[tile]^T → S ─── - s0_handle = mma_si_prod.acquire_and_advance() - qk_mma.set(tcgen05.Field.ACCUMULATE, kt != 0) - nblk = cute.size(tCrQ, mode=[2]) - for kb in cutlass.range(nblk, unroll_full=True): - cute.gemm(qk_mma, tStS0, - tCrQ[(None, None, kb, 0)], - tCrK[(None, None, kb, kvh.index)], - tStS0) - qk_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - s0_handle.commit() - - # ─── Wait for softmax: S → P done ─── - s0_handle = mma_si_prod.acquire_and_advance() - - # ─── Wait for V[tile] ─── - vvh = kv_cons.wait_and_advance(peek) - peek = cutlass.Boolean(1) - - # ─── PV: P @ V[tile] → O ─── - pv_mma.set(tcgen05.Field.ACCUMULATE, kt != 0) - nblk_pv = cute.size(tOrP0, mode=[2]) - for kb in cutlass.range(nblk_pv, unroll_full=True): - cute.gemm(pv_mma, tOtO0, - tOrP0[(None, None, kb)], - tCrV[(None, None, kb, vvh.index)], - tOtO0) - pv_mma.set(tcgen05.Field.ACCUMULATE, True) - - kvh.release() - vvh.release() - - acc_pipe.producer_commit(acc_prod_st) - acc_prod_st.advance() - acc_pipe.producer_tail(acc_prod_st) - - # ═══ SOFTMAX WARPS ═══ - if warp_idx < self.mma_warp_id: - tmem.allocate(self.num_tmem_alloc_cols) - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) - sfw_idx = tidx % (32 * len(self.softmax_warp_ids)) - - tmem_load_atom = cute.make_copy_atom( - tcgen05.copy.Ld32x32bOp(tcgen \ No newline at end of file diff --git a/tests/archive/test_fmha_v1.py b/tests/archive/test_fmha_v1.py deleted file mode 100644 index cf05f3ea..00000000 --- a/tests/archive/test_fmha_v1.py +++ /dev/null @@ -1,253 +0,0 @@ -""" -FMHA Pipeline v1: Modified v30 with pv_mma_tiler=(128,64) for head_dim=64. -Step 1: Test if the V SMEM layout works for head_dim=64 PV. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - -HEAD_DIM = 64 - -class PvHeadDimKernel: - def __init__(self): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0,1,2,3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (128, 128, qk_ik * 4) - pv_ik = cute.size(pv_mma.shape_mnk, mode=[2]) - # PV tiler: (M, N, K) = (QK_M, head_dim, QK_N) - # K of PV = sequence per tile = QK's N = 128 - pv_nphases = 128 // pv_ik # 128/16 = 8 k-phases - self.pv_mma_tiler = (128, HEAD_DIM, pv_ik * pv_nphases) # (128, 64, 128) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0]//cute.size(qk_mma.thr_id.shape), HEAD_DIM, self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - v_s = cute.slice_(self.v_smem_s, (None,None,None,0)) - v_sz = cute.size_in_bytes(self.q_dtype, v_s) - print(f"[DIAG] pv_mma_tiler={self.pv_mma_tiler} V SMEM per stage: {v_sz} bytes ({v_sz//2} BF16)") - pv_thr = pv_mma.get_slice(0) - tCrV = pv_mma.make_fragment_B(v_s) - print(f"[DIAG] tCrV shape={tCrV.shape} size={cute.size(tCrV)} k_phases={cute.size(tCrV,mode=[2])}") - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - qk_thr = qk_mma.get_slice(0); qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - pv_thr = pv_mma.get_slice(0); pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0; self.tmem_p0_offset = 32 - self.tmem_o0_offset = find_tmem_tensor_col_offset(tOtO) - tCS = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage)) - tCO = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCS, tCO], arch="sm_100") - a = cute.slice_(self.a_smem_s,(None,None,None,0)); b = cute.slice_(self.b_smem_s,(None,None,None,0)) - vs = cute.slice_(self.v_smem_s,(None,None,None,0)) - self.num_tma_load_bytes = (cute.size_in_bytes(self.q_dtype,a)+cute.size_in_bytes(self.q_dtype,b)+cute.size_in_bytes(self.q_dtype,vs))*cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, self.a_major, self.b_major, self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, self.qk_acc_dtype, self.cta_group, (128,HEAD_DIM), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - q_s = cute.slice_(self.a_smem_s,(None,None,None,0)); k_s = cute.slice_(self.b_smem_s,(None,None,None,0)) - v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id),q,q_s,self.mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id),k,k_s,self.mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape) - epi_s = cute.select(self.c_smem_s,mode=[0,1]) - tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile) - self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.a_smem_s,self.b_smem_s,self.v_smem_s,self.p_tmem_s,self.c_smem_s,self.epi_tile).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx,_,_ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage*2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage*2] - tmem_dealloc: cutlass.Int64; holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - ab_p,ab_c = pipeline.PipelineTmaUmma.create(barrier_storage=st.ab_bar.data_ptr(),num_stages=self.num_ab_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.num_tma_load_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - mma_si_prod,mma_si_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.mma_si_bar.data_ptr(),num_stages=1,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.epilogue_warp_id))).make_participants() - acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(),num_stages=self.num_acc_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,len(self.epilogue_warp_id)*(2 if cute.size(qk_mma.thr_id.shape)==2 else 1)),cta_layout_vmnk=cl_vmnk,defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2,num_threads=32*len((self.mma_warp_id,*self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr,barrier_for_retrieve=tmem_bar,allocator_warp_id=self.epilogue_warp_id[0],is_two_cta=cute.size(qk_mma.thr_id.shape)==2,two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk,is_relaxed=True) - sQ = smem.allocate_tensor(element_type=self.q_dtype,layout=a_smem_s.outer,byte_alignment=128,swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype,layout=b_smem_s.outer,byte_alignment=128,swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype,layout=v_smem_s.outer,byte_alignment=128,swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype,layout=c_smem_s.outer,byte_alignment=128,swizzle=c_smem_s.inner) - gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None)) - gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None)) - gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None)) - k_cnt = cute.size(gQ,mode=[3]) - qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = pv_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,0,None,0)).shape) - tAsQ,tAgQ = cpasync.tma_partition(tma_q,0,a_lay,cute.group_modes(sQ,0,3),cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape) - tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - tStS0 = cute.make_tensor(tStS.iterator+self.tmem_s0_offset,tStS.layout) - pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - tOtO0 = cute.make_tensor(tOtO.iterator+self.tmem_o0_offset,tOtO.layout) - tP = cute.make_tensor(tStS.iterator,p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None,None,None,0)] - tOrP0 = cute.make_tensor(tOrP.iterator+self.qk_acc_dtype.width//self.q_dtype.width*self.tmem_p0_offset,tOrP.layout) - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_as,self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_as,self.num_acc_stage)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt,unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q,tAgQ[(None,h.count)],tAsQ[(None,h.index)],tma_bar_ptr=h.barrier) - cute.copy(tma_k,tBgK[(None,h.count)],tBsK[(None,h.index)],tma_bar_ptr=h.barrier) - cute.copy(tma_v,tVgV[(None,h.count)],tVsV[(None,h.index)],tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else "FAIL"}') - if cos < 0.99: - print(f' out[0,:4] = {out[0,:4].tolist()}') - print(f' ref[0,:4] = {ref[0,:4].tolist()}') - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_fmha_v2.py b/tests/archive/test_fmha_v2.py deleted file mode 100644 index 873ecdf2..00000000 --- a/tests/archive/test_fmha_v2.py +++ /dev/null @@ -1,245 +0,0 @@ -""" -FMHA Pipeline v2: KV-tile interleaved with V overwriting K in SMEM. -K and V use separate pipeline entries from the same kv pipeline. -MMA: wait for K → QK → wait for softmax → wait for V → PV per KV-tile. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - -HEAD_DIM = 64 - -class FmhaKernel: - def __init__(self): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0,1,2,3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.num_c_stage = 2 - self.kv_stage = 2; self.q_stage = 1 - - def _setup(self, qk_mma, pv_mma): - qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (128, 128, qk_ik * 4) - pv_ik = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (128, HEAD_DIM, pv_ik * (128 // pv_ik)) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0]//cute.size(qk_mma.thr_id.shape), HEAD_DIM, self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - self.q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, self.q_dtype, self.q_stage) - self.k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, self.q_dtype, self.kv_stage) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, self.kv_stage) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - qk_thr = qk_mma.get_slice(0); qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - pv_thr = pv_mma.get_slice(0); pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0; self.tmem_p0_offset = 32 - self.tmem_o0_offset = find_tmem_tensor_col_offset(tOtO) - tCS = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage)) - tCO = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCS, tCO], arch="sm_100") - a = cute.slice_(self.q_smem_s,(None,None,None,0)); b = cute.slice_(self.k_smem_s,(None,None,None,0)) - self.num_tma_load_bytes = (cute.size_in_bytes(self.q_dtype,a)+cute.size_in_bytes(self.q_dtype,b))*cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, self.a_major, self.b_major, self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, self.qk_acc_dtype, self.cta_group, (128,HEAD_DIM), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)) - k_s = cute.slice_(self.k_smem_s,(None,None,None,0)) - v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id),q,q_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id),k,k_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape) - epi_s = cute.select(self.c_smem_s,mode=[0,1]) - tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile) - self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.v_smem_s,self.p_tmem_s,self.c_smem_s,self.epi_tile).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, q_smem_s, k_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx,_,_ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - @cute.struct - class SS: - q_bar: cute.struct.MemRange[cutlass.Int64, self.q_stage*2] - kv_bar: cute.struct.MemRange[cutlass.Int64, self.kv_stage*2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage*2] - tmem_dealloc: cutlass.Int64; holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - qp,qc = pipeline.PipelineTmaUmma.create(barrier_storage=st.q_bar.data_ptr(),num_stages=self.q_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.num_tma_load_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - kvp,kvc = pipeline.PipelineTmaUmma.create(barrier_storage=st.kv_bar.data_ptr(),num_stages=self.kv_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.num_tma_load_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - mma_si_prod,mma_si_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.mma_si_bar.data_ptr(),num_stages=1,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.epilogue_warp_id))).make_participants() - acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(),num_stages=self.num_acc_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,len(self.epilogue_warp_id)),cta_layout_vmnk=cl_vmnk,defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2,num_threads=32*len((self.mma_warp_id,*self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr,barrier_for_retrieve=tmem_bar,allocator_warp_id=self.epilogue_warp_id[0],is_two_cta=cute.size(qk_mma.thr_id.shape)==2,two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk,is_relaxed=True) - sQ = smem.allocate_tensor(element_type=self.q_dtype,layout=q_smem_s.outer,byte_alignment=128,swizzle=q_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype,layout=k_smem_s.outer,byte_alignment=128,swizzle=k_smem_s.inner) - # V overwrites K SMEM - sV = cute.make_tensor(cute.recast_ptr(sK.iterator, v_smem_s.inner), v_smem_s.outer) - sC = smem.allocate_tensor(element_type=self.o_dtype,layout=c_smem_s.outer,byte_alignment=128,swizzle=c_smem_s.inner) - gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None)) - gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None)) - gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None)) - gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None)) - n_kv_tiles = cute.size(gK, mode=[3]) - qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - tCgV = pv_thr.partition_B(gV); tCgC = pv_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,0,None,0)).shape) - tAsQ,tAgQ = cpasync.tma_partition(tma_q,0,a_lay,cute.group_modes(sQ,0,3),cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape) - tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3)) - tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)]; tVgV = tVgV[(None,0,None,0)] - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - tStS0 = cute.make_tensor(tStS.iterator+self.tmem_s0_offset,tStS.layout) - pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - tOtO0 = cute.make_tensor(tOtO.iterator+self.tmem_o0_offset,tOtO.layout) - tP = cute.make_tensor(tStS.iterator,p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None,None,None,0)] - tOrP0 = cute.make_tensor(tOrP.iterator+self.qk_acc_dtype.width//self.q_dtype.width*self.tmem_p0_offset,tOrP.layout) - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_as,self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_as,self.num_acc_stage)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD ═══ - if warp_idx == self.tma_warp_id: - qp.reset(); qh = qp.acquire_and_advance() - cute.copy(tma_q,tAgQ[(None,qh.count)],tAsQ[(None,qh.index)],tma_bar_ptr=qh.barrier) - qp.tail() - kvp.reset(); pk = kvp.try_acquire() - for kt in cutlass.range(n_kv_tiles,unroll=1): - kh = kvp.acquire_and_advance(pk) - cute.copy(tma_k,tBgK[(None,kh.count)],tBsK[(None,kh.index)],tma_bar_ptr=kh.barrier) - vh = kvp.acquire_and_advance(cutlass.Boolean(1)) - cute.copy(tma_v,tVgV[(None,vh.count)],tVsV[(None,vh.index)],tma_bar_ptr=vh.barrier) - pk = cutlass.Boolean(1) - kvp.tail() - - # ═══ MMA ═══ - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - qc.reset(); qh = qc.wait_and_advance(); qh.release() - kvc.reset(); pk = kvc.try_wait() - acc_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer,self.num_acc_stage) - acc_pipe.producer_acquire(acc_st) - for kt in range(n_kv_tiles): - kh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - sh = mma_si_prod.acquire_and_advance() - qk_mma.set(tcgen05.Field.ACCUMULATE, kt != 0) - for kb in cutlass.range(cute.size(tCrQ,mode=[2]),unroll_full=True): - cute.gemm(qk_mma,tStS0,tCrQ[(None,None,kb,0)],tCrK[(None,None,kb,kh.index)],tStS0) - qk_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store(); sh.commit() - sh = mma_si_prod.acquire_and_advance() # wait softmax - vh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) # wait V - pv_mma.set(tcgen05.Field.ACCUMULATE, kt != 0) - for kb in cutlass.range(cute.size(tOrP0,mode=[2]),unroll_full=True): - cute.gemm(pv_mma,tOtO0,tOrP0[(None,None,kb)],tCrV[(None,None,kb,vh.index)],tOtO0) - pv_mma.set(tcgen05.Field.ACCUMULATE, True) - kh.release(); vh.release() - acc_pipe.producer_commit(acc_st); acc_st.advance(); acc_pipe.producer_tail(acc_st) - - # ═══ EPILOGUE ═══ - if warp_idx < self.mma_warp_id: - tmem.allocate(self.num_tmem_alloc_cols) - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) - sfw_idx = tidx % (32 * len(self.epilogue_warp_id)) - tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS0) - thr_load = tiled_tmem_load.get_slice(sfw_idx) - tTMEM_LOADtS = thr_load.partition_S(tStS0) - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tTMEM_LOADcS = thr_load.partition_D(tScS) - tStS_P_layout = cute.composition(tStS.layout, cute.make_layout((128, self.tilePlikeFP32))) - tStS_P = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStS_P_layout) - tmem_store_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStS_P) - thr_store = tiled_tmem_store.get_slice(sfw_idx) - tTMEM_STOREtS_x4 = thr_store.partition_D(tStS_P) - tScS_P_layout = cute.composition(tScS.layout, cute.make_layout((128, self.tilePlikeFP32))) - tScS_P = cute.make_tensor(tScS.iterator, tScS_P_layout) - tTMEM_STOREcS = thr_store.partition_S(tScS_P) - for kt in range(n_kv_tiles): - si_handle = mma_si_cons.wait_and_advance() - tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS) - tTMEM_STORErS_x4 = cute.make_rmem_tensor(tTMEM_STOREcS.shape, self.qk_acc_dtype) - tTMEM_STORErS_x4_e = cute.make_tensor(cute.recast_ptr(tTMEM_STORErS_x4.iterator, dtype=self.q_dtype), tTMEM_LOADrS.layout) - frg_cnt = 4; frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt - tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile)) - tTMEM_STORErS_x4_e_frg = cute.logical_divide(tTMEM_STORErS_x4_e, cute.make_layout(frg_tile)) - for j in range(frg_cnt): - s_vec = tTMEM_LOADrS_frg[None, j].load() - tTMEM_STORErS_x4_e_frg[None, j].store(s_vec.to(self.q_dtype)) - cute.copy(tiled_tmem_store, tTMEM_STORErS_x4, tTMEM_STOREtS_x4) - cute.arch.fence_view_async_tmem_store() - si_handle.release() - tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout) - acc_cons_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage) - c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) - c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp) - acc_cons_st = utils.gemm.sm100.epilogue_tma_store(self, tidx, warp_idx, tma_c, tCtO_base, sC, tCgC, epi_tile, 0, const_expr(lambda x: x), (0,0,0), acc_cons_st, acc_pipe, c_pipe) - c_pipe.producer_tail() - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def test(): - torch.manual_seed(42) - for n in [128]: - m, hd = 128, HEAD_DIM - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(n, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.ones(n, hd, 1, dtype=torch.bfloat16, device='cuda') - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - qf = q[:,:,0].float(); kf = k[:,:,0].float() - ref = (qf @ kf.T).bfloat16().float() @ v[:,:,0].float() - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - kernel = FmhaKernel() - print(f'n={n}: Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - print(f'n={n}: Running...', flush=True) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - out = c[:,:,0].float() - cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - print(f'FMHA v2 n={n} V=ones: cosine {cos:.6f} {"PASS" if cos >= 0.99 else "FAIL"}') - if cos < 0.99: - print(f' out[0,:4]={out[0,:4].tolist()} ref[0,:4]={ref[0,:4].tolist()}') - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_fmha_v2_fixed.py b/tests/archive/test_fmha_v2_fixed.py deleted file mode 100644 index 0c077a48..00000000 --- a/tests/archive/test_fmha_v2_fixed.py +++ /dev/null @@ -1,277 +0,0 @@ -""" -FMHA Pipeline v2 FIXED: - 1. tx_count per pipeline (Q gets Q bytes, KV gets K bytes) - 2. Separate SMEM for K and V (V no longer recast_ptr into K) - 3. kv_stage=2: K in stage 0, V in stage 1 (different pipeline stages, no SMEM overlap) - 4. s_pipe (PipelineUmmaAsync): MMA→epilogue "scores ready" - softmax_done_bar (NamedBarrier): epilogue→MMA "softmax done" -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - -HEAD_DIM = 64 - -class FmhaKernel: - def __init__(self): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0,1,2,3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.num_c_stage = 2 - self.kv_stage = 2; self.q_stage = 1 - - def _setup(self, qk_mma, pv_mma): - qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (128, 128, qk_ik * 4) - pv_ik = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (128, HEAD_DIM, pv_ik * (128 // pv_ik)) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0]//cute.size(qk_mma.thr_id.shape), HEAD_DIM, self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - self.q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, self.q_dtype, self.q_stage) - self.k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, self.q_dtype, self.kv_stage) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, self.kv_stage) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - qk_thr = qk_mma.get_slice(0); qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - pv_thr = pv_mma.get_slice(0); pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0; self.tmem_p0_offset = 32 - self.tmem_o0_offset = find_tmem_tensor_col_offset(tOtO) - tCS = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage)) - tCO = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCS, tCO], arch="sm_100") - # Per-pipeline tx_count (FIX #1) - cta = cute.size(qk_mma.thr_id.shape) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)) - k_s = cute.slice_(self.k_smem_s,(None,None,None,0)) - v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - self.q_tx_bytes = cute.size_in_bytes(self.q_dtype, q_s) * cta - # KV pipeline: each slot holds K OR V, not both. K and V are the same size per slot - # when head_dim=64 (both 128*64*2 = 16384 bytes). - self.kv_tx_bytes = cute.size_in_bytes(self.q_dtype, k_s) * cta - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, self.a_major, self.b_major, self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, self.qk_acc_dtype, self.cta_group, (128,HEAD_DIM), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)) - k_s = cute.slice_(self.k_smem_s,(None,None,None,0)) - v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id),q,q_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id),k,k_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape) - epi_s = cute.select(self.c_smem_s,mode=[0,1]) - tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile) - self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.v_smem_s,self.p_tmem_s,self.c_smem_s,self.epi_tile).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, q_smem_s, k_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx,_,_ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - @cute.struct - class SS: - q_bar: cute.struct.MemRange[cutlass.Int64, self.q_stage*2] - kv_bar: cute.struct.MemRange[cutlass.Int64, self.kv_stage*2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage*2] - tmem_dealloc: cutlass.Int64; holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - # FIX #1: per-pipeline tx_count - qp,qc = pipeline.PipelineTmaUmma.create(barrier_storage=st.q_bar.data_ptr(),num_stages=self.q_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.q_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - kvp,kvc = pipeline.PipelineTmaUmma.create(barrier_storage=st.kv_bar.data_ptr(),num_stages=self.kv_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.kv_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - - # FIX #2: s_pipe for MMA→epilogue, NamedBarrier for epilogue→MMA - mma_si_prod,mma_si_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.mma_si_bar.data_ptr(),num_stages=1,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.epilogue_warp_id))).make_participants() - softmax_done_bar = pipeline.NamedBarrier(barrier_id=3, num_threads=32 + 32*len(self.epilogue_warp_id)) - - acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(),num_stages=self.num_acc_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,len(self.epilogue_warp_id)),cta_layout_vmnk=cl_vmnk,defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2,num_threads=32*len((self.mma_warp_id,*self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr,barrier_for_retrieve=tmem_bar,allocator_warp_id=self.epilogue_warp_id[0],is_two_cta=cute.size(qk_mma.thr_id.shape)==2,two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk,is_relaxed=True) - - # FIX #3: Separate SMEM for K and V (no recast_ptr overlap) - sQ = smem.allocate_tensor(element_type=self.q_dtype,layout=q_smem_s.outer,byte_alignment=128,swizzle=q_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype,layout=k_smem_s.outer,byte_alignment=128,swizzle=k_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype,layout=v_smem_s.outer,byte_alignment=128,swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype,layout=c_smem_s.outer,byte_alignment=128,swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None)) - gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None)) - gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None)) - gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None)) - n_kv_tiles = cute.size(gK, mode=[3]) - qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - tCgV = pv_thr.partition_B(gV); tCgC = pv_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,0,None,0)).shape) - tAsQ,tAgQ = cpasync.tma_partition(tma_q,0,a_lay,cute.group_modes(sQ,0,3),cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape) - tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3)) - tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)]; tVgV = tVgV[(None,0,None,0)] - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - tStS0 = cute.make_tensor(tStS.iterator+self.tmem_s0_offset,tStS.layout) - pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - tOtO0 = cute.make_tensor(tOtO.iterator+self.tmem_o0_offset,tOtO.layout) - tP = cute.make_tensor(tStS.iterator,p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None,None,None,0)] - tOrP0 = cute.make_tensor(tOrP.iterator+self.qk_acc_dtype.width//self.q_dtype.width*self.tmem_p0_offset,tOrP.layout) - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_as,self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_as,self.num_acc_stage)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD ═══ - # K and V on separate SMEM, loaded back-to-back into kv pipeline - # K goes to pipeline stage k_handle.index, V to v_handle.index - if warp_idx == self.tma_warp_id: - qp.reset(); qh = qp.acquire_and_advance() - cute.copy(tma_q,tAgQ[(None,qh.count)],tAsQ[(None,qh.index)],tma_bar_ptr=qh.barrier) - qp.tail() - kvp.reset(); pk = kvp.try_acquire() - for kt in cutlass.range(n_kv_tiles,unroll=1): - kh = kvp.acquire_and_advance(pk) - cute.copy(tma_k,tBgK[(None,kh.count)],tBsK[(None,kh.index)],tma_bar_ptr=kh.barrier) - vh = kvp.acquire_and_advance(cutlass.Boolean(1)) - cute.copy(tma_v,tVgV[(None,vh.count)],tVsV[(None,vh.index)],tma_bar_ptr=vh.barrier) - pk = cutlass.Boolean(1) - kvp.tail() - - # ═══ MMA ═══ - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - qc.reset(); qh = qc.wait_and_advance(); qh.release() - kvc.reset(); pk = kvc.try_wait() - acc_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer,self.num_acc_stage) - acc_pipe.producer_acquire(acc_st) - for kt in range(n_kv_tiles): - # Wait for K - kh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - # QK MMA - sh = mma_si_prod.acquire_and_advance() - qk_mma.set(tcgen05.Field.ACCUMULATE, kt != 0) - for kb in cutlass.range(cute.size(tCrQ,mode=[2]),unroll_full=True): - cute.gemm(qk_mma,tStS0,tCrQ[(None,None,kb,0)],tCrK[(None,None,kb,kh.index)],tStS0) - qk_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - sh.commit() - kh.release() - - # Wait for softmax - softmax_done_bar.wait() - - # Wait for V - vh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - # PV MMA - pv_mma.set(tcgen05.Field.ACCUMULATE, kt != 0) - for kb in cutlass.range(cute.size(tOrP0,mode=[2]),unroll_full=True): - cute.gemm(pv_mma,tOtO0,tOrP0[(None,None,kb)],tCrV[(None,None,kb,vh.index)],tOtO0) - pv_mma.set(tcgen05.Field.ACCUMULATE, True) - vh.release() - acc_pipe.producer_commit(acc_st); acc_st.advance(); acc_pipe.producer_tail(acc_st) - - # ═══ EPILOGUE ═══ - if warp_idx < self.mma_warp_id: - tmem.allocate(self.num_tmem_alloc_cols) - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) - sfw_idx = tidx % (32 * len(self.epilogue_warp_id)) - tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS0) - thr_load = tiled_tmem_load.get_slice(sfw_idx) - tTMEM_LOADtS = thr_load.partition_S(tStS0) - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tTMEM_LOADcS = thr_load.partition_D(tScS) - tStS_P_layout = cute.composition(tStS.layout, cute.make_layout((128, self.tilePlikeFP32))) - tStS_P = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStS_P_layout) - tmem_store_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStS_P) - thr_store = tiled_tmem_store.get_slice(sfw_idx) - tTMEM_STOREtS_x4 = thr_store.partition_D(tStS_P) - tScS_P_layout = cute.composition(tScS.layout, cute.make_layout((128, self.tilePlikeFP32))) - tScS_P = cute.make_tensor(tScS.iterator, tScS_P_layout) - tTMEM_STOREcS = thr_store.partition_S(tScS_P) - for kt in range(n_kv_tiles): - si_handle = mma_si_cons.wait_and_advance() - # Identity softmax: FP32→BF16, write P to TMEM - tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS) - tTMEM_STORErS_x4 = cute.make_rmem_tensor(tTMEM_STOREcS.shape, self.qk_acc_dtype) - tTMEM_STORErS_x4_e = cute.make_tensor(cute.recast_ptr(tTMEM_STORErS_x4.iterator, dtype=self.q_dtype), tTMEM_LOADrS.layout) - frg_cnt = 4; frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt - tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile)) - tTMEM_STORErS_x4_e_frg = cute.logical_divide(tTMEM_STORErS_x4_e, cute.make_layout(frg_tile)) - for j in range(frg_cnt): - s_vec = tTMEM_LOADrS_frg[None, j].load() - tTMEM_STORErS_x4_e_frg[None, j].store(s_vec.to(self.q_dtype)) - cute.copy(tiled_tmem_store, tTMEM_STORErS_x4, tTMEM_STOREtS_x4) - cute.arch.fence_view_async_tmem_store() - si_handle.release() - softmax_done_bar.wait() - - tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout) - acc_cons_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage) - c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) - c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp) - acc_cons_st = utils.gemm.sm100.epilogue_tma_store(self, tidx, warp_idx, tma_c, tCtO_base, sC, tCgC, epi_tile, 0, const_expr(lambda x: x), (0,0,0), acc_cons_st, acc_pipe, c_pipe) - c_pipe.producer_tail() - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def test(): - torch.manual_seed(42) - for n in [128]: - m, hd = 128, HEAD_DIM - q = torch.randn(m, hd,1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(n, hd,1, dtype=torch.bfloat16, device='cuda') - v = torch.ones(n, hd,1, dtype=torch.bfloat16, device='cuda') - c = torch.zeros(m, hd,1, dtype=torch.bfloat16, device='cuda') - qf = q[:,:,0].float(); kf = k[:,:,0].float() - ref = (qf @ kf.T).bfloat16().float() @ v[:,:,0].float() - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - kernel = FmhaKernel() - print(f'n={n}: Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - print(f'n={n}: Running...', flush=True) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - out = c[:,:,0].float() - cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - print(f'FMHA v2 FIXED n={n} V=ones: cosine {cos:.6f} {"PASS" if cos >= 0.99 else "FAIL"}') - if cos < 0.99: - print(f' out[0,:4]={out[0,:4].tolist()} ref[0,:4]={ref[0,:4].tolist()}') - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_fmha_v3_12w.py b/tests/archive/test_fmha_v3_12w.py deleted file mode 100644 index 69a7ab33..00000000 --- a/tests/archive/test_fmha_v3_12w.py +++ /dev/null @@ -1,327 +0,0 @@ -""" -FMHA v3: QK -> softmax -> PV with KV-tile interleaving. -Bug 4b fix (FMHA pattern): P store uses QK C-fragment layout composition, -NOT PV A-fragment layout. Register bridge: FP32 backing (store partition shape) -recast to BF16 view (QK-load layout). -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - -HEAD_DIM = 64 - -class FmhaV3: - def __init__(self): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0,1,2,3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.num_c_stage = 2 - self.kv_stage = 2; self.q_stage = 1; self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (128, 128, qk_ik * 4) - pv_ik = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (128, HEAD_DIM, pv_ik * (128 // pv_ik)) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0]//cute.size(qk_mma.thr_id.shape), HEAD_DIM, self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - self.q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, self.q_dtype, self.q_stage) - self.k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, self.q_dtype, self.kv_stage) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, self.kv_stage) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - qk_thr = qk_mma.get_slice(0); qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - pv_thr = pv_mma.get_slice(0); pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - self.tmem_s0_offset = 0; self.tmem_p0_offset = 32 - # P occupies [tmem_p0_offset, tmem_p0_offset + p_cols_fp32) - # S occupies [0, qk_mma_tiler[1]) = [0, 128) - # O must NOT overlap P. Place O after max(S end, P end), aligned to 32. - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - p_end = self.tmem_p0_offset + p_cols_fp32 # 32 + 64 = 96 - s_cols = self.qk_mma_tiler[1] # 128 - o_after = max(s_cols, p_end) # 128 - self.tmem_o0_offset = ((o_after + 31) // 32) * 32 # align to 32 = 128 - o_cols = find_tmem_tensor_col_offset(tOtO) # footprint of O - total = self.tmem_o0_offset + o_cols - # Must be multiple of 32 AND power of 2 - self.num_tmem_alloc_cols = 1 - while self.num_tmem_alloc_cols < total: - self.num_tmem_alloc_cols *= 2 - cta = cute.size(qk_mma.thr_id.shape) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)) - self.q_tx_bytes = cute.size_in_bytes(self.q_dtype, q_s) * cta - self.kv_tx_bytes = cute.size_in_bytes(self.q_dtype, k_s) * cta - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - # FMHA-style V: reconstruct as (HEAD_DIM, s_k, 1) MN-major - v_fmha = cute.make_tensor( - v.iterator, - cute.make_layout( - (HEAD_DIM, 128, 1), - stride=(1, HEAD_DIM, HEAD_DIM * 128), - ), - ) - self.v_major = LayoutEnum.from_tensor(v_fmha).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, self.a_major, self.b_major, self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, self.qk_acc_dtype, self.cta_group, (128,HEAD_DIM), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)); v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id),q,q_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id),k,k_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v_fmha,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape) - epi_s = cute.select(self.c_smem_s,mode=[0,1]) - tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile) - self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.v_smem_s,self.p_tmem_s,self.c_smem_s,self.epi_tile).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, q_smem_s, k_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx,_,_ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k); cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - q_bar: cute.struct.MemRange[cutlass.Int64, self.q_stage*2] - kv_bar: cute.struct.MemRange[cutlass.Int64, self.kv_stage*2] - s_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage*2] - tmem_dealloc: cutlass.Int64; holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - qp,qc = pipeline.PipelineTmaUmma.create(barrier_storage=st.q_bar.data_ptr(),num_stages=self.q_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.q_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - kvp,kvc = pipeline.PipelineTmaUmma.create(barrier_storage=st.kv_bar.data_ptr(),num_stages=self.kv_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.kv_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - s_prod,s_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.s_bar.data_ptr(),num_stages=1,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.epilogue_warp_id))).make_participants() - softmax_done_bar = pipeline.NamedBarrier(barrier_id=3, num_threads=32 + 32*len(self.epilogue_warp_id)) - acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(),num_stages=self.num_acc_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,len(self.epilogue_warp_id)),cta_layout_vmnk=cl_vmnk,defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2,num_threads=32*len((self.mma_warp_id,*self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr,barrier_for_retrieve=tmem_bar,allocator_warp_id=self.epilogue_warp_id[0],is_two_cta=cute.size(qk_mma.thr_id.shape)==2,two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk,is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype,layout=q_smem_s.outer,byte_alignment=128,swizzle=q_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype,layout=k_smem_s.outer,byte_alignment=128,swizzle=k_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype,layout=v_smem_s.outer,byte_alignment=128,swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype,layout=c_smem_s.outer,byte_alignment=128,swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None)) - gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None)) - gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None)) - gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None)) - n_kv_tiles = cute.size(gK, mode=[3]) - - qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - tCgV = pv_thr.partition_B(gV); tCgC = pv_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,0,None,0)).shape) - tAsQ,tAgQ = cpasync.tma_partition(tma_q,0,a_lay,cute.group_modes(sQ,0,3),cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape) - tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3)) - tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)]; tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - # --- PV read view (for MMA only, NOT for softmax store) --- - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None,None,None,0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # TMA LOAD - if warp_idx == self.tma_warp_id: - qp.reset(); qh = qp.acquire_and_advance() - cute.copy(tma_q,tAgQ[(None,qh.count)],tAsQ[(None,qh.index)],tma_bar_ptr=qh.barrier) - qp.tail() - kvp.reset(); pk = kvp.try_acquire() - for kt in cutlass.range(n_kv_tiles,unroll=1): - kh = kvp.acquire_and_advance(pk) - cute.copy(tma_k,tBgK[(None,kh.count)],tBsK[(None,kh.index)],tma_bar_ptr=kh.barrier) - pk = cutlass.Boolean(1) - vh = kvp.acquire_and_advance(pk) - cute.copy(tma_v,tVgV[(None,vh.count)],tVsV[(None,vh.index)],tma_bar_ptr=vh.barrier) - pk = cutlass.Boolean(1) - kvp.tail() - - # MMA - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - qc.reset(); qh = qc.wait_and_advance(); qh.release() - kvc.reset(); pk = kvc.try_wait() - acc_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_stage) - acc_pipe.producer_acquire(acc_st) - for kt in range(n_kv_tiles): - kh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - sh = s_prod.acquire_and_advance() - qk_mma.set(tcgen05.Field.ACCUMULATE, False) - for kb in cutlass.range(cute.size(tCrQ,mode=[2]), unroll_full=True): - cute.gemm(qk_mma, tStS0, tCrQ[(None,None,kb,0)], tCrK[(None,None,kb,kh.index)], tStS0) - qk_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - sh.commit(); kh.release() - softmax_done_bar.arrive_and_wait() - vh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - pv_mma.set(tcgen05.Field.ACCUMULATE, kt != 0) - for kb in cutlass.range(cute.size(tOrP0,mode=[2]), unroll_full=True): - cute.gemm(pv_mma, tOtO0, tOrP0[(None,None,kb)], tCrV[(None,None,kb,vh.index)], tOtO0) - pv_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - vh.release() - acc_pipe.producer_commit(acc_st); acc_st.advance() - acc_pipe.producer_tail(acc_st) - - # EPILOGUE - if warp_idx < self.mma_warp_id: - tmem.allocate(self.num_tmem_alloc_cols) - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) - sfw_idx = tidx % (32 * len(self.epilogue_warp_id)) - - # --- S load (QK C-fragment layout) --- - tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS0) - thr_load = tiled_tmem_load.get_slice(sfw_idx) - tTMEM_LOADtS = thr_load.partition_S(tStS0) - - # S coordinate tensor (QK C-fragment) - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tTMEM_LOADcS = thr_load.partition_D(tScS) - - # --- P store (QK C-fragment layout composition, FMHA pattern) --- - # P logical columns = PV K = QK N = pv_mma_tiler[2] - # Packed FP32 columns: BF16 pairs packed into FP32 words - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - # BF16: 128 * 16 / 32 = 64 - - # P TMEM destination: QK C-fragment layout composed with P sub-tile - tStP_layout = cute.composition( - tStS.layout, - cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32)), - ) - tStP0 = cute.make_tensor( - tStS.iterator + self.tmem_p0_offset, - tStP_layout, - ) - - # P TMEM store atom and tiled copy - tmem_store_atom = cute.make_copy_atom( - tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), - self.qk_acc_dtype, - ) - tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStP0) - thr_store = tiled_tmem_store.get_slice(sfw_idx) - tTMEM_STOREtP = thr_store.partition_D(tStP0) - - # P coordinate tensor: QK C-fragment coordinate composed with P sub-tile - tScP_layout = cute.composition( - tScS.layout, - cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32)), - ) - tScP = cute.make_tensor(tScS.iterator, tScP_layout) - tTMEM_STOREcP = thr_store.partition_S(tScP) - - for kt in range(n_kv_tiles): - si_handle = s_cons.wait_and_advance() - - # Load S from TMEM (FP32, QK C-fragment layout) - tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS) - - # Register bridge (FMHA pattern): - # rP_words: FP32 backing store with store-partition shape - # rP_bf16: BF16 view over same registers using QK-load layout - rP_words = cute.make_rmem_tensor(tTMEM_STOREcP.shape, self.qk_acc_dtype) - rP_bf16 = cute.make_tensor( - cute.recast_ptr(rP_words.iterator, dtype=self.q_dtype), - tTMEM_LOADrS.layout, - ) - - # Fragmented load→convert→store: - # Load S as FP32, convert to BF16, store through rP_bf16 view - frg_cnt = 4 - frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt - tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile)) - rP_bf16_frg = cute.logical_divide(rP_bf16, cute.make_layout(frg_tile)) - for j in range(frg_cnt): - s_vec = tTMEM_LOADrS_frg[None, j].load() - rP_bf16_frg[None, j].store(s_vec.to(self.q_dtype)) - - # Copy packed FP32 backing registers to TMEM - cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP) - cute.arch.fence_view_async_tmem_store() - si_handle.release() - softmax_done_bar.arrive() - - tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout) - acc_cons_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage) - c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) - c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp) - acc_cons_st = utils.gemm.sm100.epilogue_tma_store(self, tidx, warp_idx, tma_c, tCtO_base, sC, tCgC, epi_tile, 0, const_expr(lambda x: x), (0,0,0), acc_cons_st, acc_pipe, c_pipe) - c_pipe.producer_tail() - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def test(): - torch.manual_seed(42) - for n in [128]: - m, hd = 128, HEAD_DIM - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(n, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(n, hd, dtype=torch.bfloat16, device='cuda') - # V passed as (n, hd) row-major — FMHA-style reconstruction inside kernel - v_kernel = v.unsqueeze(-1) - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - qf = q[:,:,0].float(); kf = k[:,:,0].float() - ref = (qf @ kf.T).bfloat16().float() @ v.float() - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - kernel = FmhaV3() - print(f'n={n}: Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - print(f'n={n}: tmem_offsets: s0={kernel.tmem_s0_offset} p0={kernel.tmem_p0_offset} o0={kernel.tmem_o0_offset} alloc={kernel.num_tmem_alloc_cols}', flush=True) - print(f'n={n}: Running...', flush=True) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - out = c[:,:,0].float() - cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - print(f'FMHA v3 n={n} V=ones: cosine {cos:.6f} {"PASS" if cos >= 0.99 else "FAIL"}') - if cos < 0.99: - print(f' out[0,:4]={out[0,:4].tolist()} ref[0,:4]={ref[0,:4].tolist()}') - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_fmha_v3_debug.py b/tests/archive/test_fmha_v3_debug.py deleted file mode 100644 index 41a4c15b..00000000 --- a/tests/archive/test_fmha_v3_debug.py +++ /dev/null @@ -1,284 +0,0 @@ -""" -FMHA v3 debug: Test with HEAD_DIM=128 first (PV 128,128) to verify plumbing, -then HEAD_DIM=64 to debug Bug 4b. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - -class FmhaV3Debug: - def __init__(self, head_dim=64): - self.head_dim = head_dim - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0,1,2,3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.num_c_stage = 2 - self.kv_stage = 2; self.q_stage = 1; self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (128, 128, qk_ik * 4) - pv_ik = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (128, self.head_dim, pv_ik * (128 // pv_ik)) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0]//cute.size(qk_mma.thr_id.shape), self.head_dim, self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - self.q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, self.q_dtype, self.q_stage) - self.k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, self.q_dtype, self.kv_stage) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, self.kv_stage) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - qk_thr = qk_mma.get_slice(0); qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - pv_thr = pv_mma.get_slice(0); pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - self.tmem_s0_offset = 0; self.tmem_p0_offset = 32 - self.tmem_o0_offset = find_tmem_tensor_col_offset(tOtO) - tCS = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage)) - tCO = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCS, tCO], arch="sm_100") - cta = cute.size(qk_mma.thr_id.shape) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)) - self.q_tx_bytes = cute.size_in_bytes(self.q_dtype, q_s) * cta - self.kv_tx_bytes = cute.size_in_bytes(self.q_dtype, k_s) * cta - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - a_major = LayoutEnum.from_tensor(q).mma_major_mode() - b_major = LayoutEnum.from_tensor(k).mma_major_mode() - v_major = LayoutEnum.from_tensor(v).mma_major_mode() - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, a_major, b_major, self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, v_major, self.qk_acc_dtype, self.cta_group, (128,self.head_dim), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)); v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id),q,q_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id),k,k_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape) - epi_s = cute.select(self.c_smem_s,mode=[0,1]) - tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile) - self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.v_smem_s,self.p_tmem_s,self.c_smem_s,self.epi_tile).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, q_smem_s, k_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx,_,_ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k); cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - q_bar: cute.struct.MemRange[cutlass.Int64, self.q_stage*2] - kv_bar: cute.struct.MemRange[cutlass.Int64, self.kv_stage*2] - s_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage*2] - tmem_dealloc: cutlass.Int64; holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - qp,qc = pipeline.PipelineTmaUmma.create(barrier_storage=st.q_bar.data_ptr(),num_stages=self.q_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.q_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - kvp,kvc = pipeline.PipelineTmaUmma.create(barrier_storage=st.kv_bar.data_ptr(),num_stages=self.kv_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.kv_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - s_prod,s_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.s_bar.data_ptr(),num_stages=1,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.epilogue_warp_id))).make_participants() - softmax_done_bar = pipeline.NamedBarrier(barrier_id=3, num_threads=32 + 32*len(self.epilogue_warp_id)) - acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(),num_stages=self.num_acc_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,len(self.epilogue_warp_id)),cta_layout_vmnk=cl_vmnk,defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2,num_threads=32*len((self.mma_warp_id,*self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr,barrier_for_retrieve=tmem_bar,allocator_warp_id=self.epilogue_warp_id[0],is_two_cta=cute.size(qk_mma.thr_id.shape)==2,two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk,is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype,layout=q_smem_s.outer,byte_alignment=128,swizzle=q_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype,layout=k_smem_s.outer,byte_alignment=128,swizzle=k_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype,layout=v_smem_s.outer,byte_alignment=128,swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype,layout=c_smem_s.outer,byte_alignment=128,swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None)) - gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None)) - gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None)) - gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None)) - n_kv_tiles = cute.size(gK, mode=[3]) - - qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - tCgV = pv_thr.partition_B(gV); tCgC = pv_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,0,None,0)).shape) - tAsQ,tAgQ = cpasync.tma_partition(tma_q,0,a_lay,cute.group_modes(sQ,0,3),cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape) - tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3)) - tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)]; tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - # PV read view (MMA only) - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None,None,None,0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # TMA LOAD - if warp_idx == self.tma_warp_id: - qp.reset(); qh = qp.acquire_and_advance() - cute.copy(tma_q,tAgQ[(None,qh.count)],tAsQ[(None,qh.index)],tma_bar_ptr=qh.barrier) - qp.tail() - kvp.reset(); pk = kvp.try_acquire() - for kt in cutlass.range(n_kv_tiles,unroll=1): - kh = kvp.acquire_and_advance(pk) - cute.copy(tma_k,tBgK[(None,kh.count)],tBsK[(None,kh.index)],tma_bar_ptr=kh.barrier) - pk = cutlass.Boolean(1) - vh = kvp.acquire_and_advance(pk) - cute.copy(tma_v,tVgV[(None,vh.count)],tVsV[(None,vh.index)],tma_bar_ptr=vh.barrier) - pk = cutlass.Boolean(1) - kvp.tail() - - # MMA - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - qc.reset(); qh = qc.wait_and_advance(); qh.release() - kvc.reset(); pk = kvc.try_wait() - acc_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_stage) - acc_pipe.producer_acquire(acc_st) - for kt in range(n_kv_tiles): - kh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - sh = s_prod.acquire_and_advance() - qk_mma.set(tcgen05.Field.ACCUMULATE, False) - for kb in cutlass.range(cute.size(tCrQ,mode=[2]), unroll_full=True): - cute.gemm(qk_mma, tStS0, tCrQ[(None,None,kb,0)], tCrK[(None,None,kb,kh.index)], tStS0) - qk_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - sh.commit(); kh.release() - softmax_done_bar.arrive_and_wait() - vh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - pv_mma.set(tcgen05.Field.ACCUMULATE, kt != 0) - for kb in cutlass.range(cute.size(tOrP0,mode=[2]), unroll_full=True): - cute.gemm(pv_mma, tOtO0, tOrP0[(None,None,kb)], tCrV[(None,None,kb,vh.index)], tOtO0) - pv_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - vh.release() - acc_pipe.producer_commit(acc_st); acc_st.advance() - acc_pipe.producer_tail(acc_st) - - # EPILOGUE - if warp_idx < self.mma_warp_id: - tmem.allocate(self.num_tmem_alloc_cols) - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) - sfw_idx = tidx % (32 * len(self.epilogue_warp_id)) - - # S load - tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS0) - thr_load = tiled_tmem_load.get_slice(sfw_idx) - tTMEM_LOADtS = thr_load.partition_S(tStS0) - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tTMEM_LOADcS = thr_load.partition_D(tScS) - - # P store — FMHA pattern: QK C-fragment layout composition - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - tStP_layout = cute.composition( - tStS.layout, - cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32)), - ) - tStP0 = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStP_layout) - - tmem_store_atom = cute.make_copy_atom( - tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), - self.qk_acc_dtype, - ) - tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStP0) - thr_store = tiled_tmem_store.get_slice(sfw_idx) - tTMEM_STOREtP = thr_store.partition_D(tStP0) - - tScP_layout = cute.composition( - tScS.layout, - cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32)), - ) - tScP = cute.make_tensor(tScS.iterator, tScP_layout) - tTMEM_STOREcP = thr_store.partition_S(tScP) - - for kt in range(n_kv_tiles): - si_handle = s_cons.wait_and_advance() - tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS) - - # Register bridge - rP_words = cute.make_rmem_tensor(tTMEM_STOREcP.shape, self.qk_acc_dtype) - rP_bf16 = cute.make_tensor( - cute.recast_ptr(rP_words.iterator, dtype=self.q_dtype), - tTMEM_LOADrS.layout, - ) - - frg_cnt = 4 - frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt - tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile)) - rP_bf16_frg = cute.logical_divide(rP_bf16, cute.make_layout(frg_tile)) - for j in range(frg_cnt): - s_vec = tTMEM_LOADrS_frg[None, j].load() - rP_bf16_frg[None, j].store(s_vec.to(self.q_dtype)) - - cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP) - cute.arch.fence_view_async_tmem_store() - si_handle.release() - softmax_done_bar.arrive() - - tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout) - acc_cons_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage) - c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) - c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp) - acc_cons_st = utils.gemm.sm100.epilogue_tma_store(self, tidx, warp_idx, tma_c, tCtO_base, sC, tCgC, epi_tile, 0, const_expr(lambda x: x), (0,0,0), acc_cons_st, acc_pipe, c_pipe) - c_pipe.producer_tail() - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def test(): - torch.manual_seed(42) - for hd in [128, 64]: - m, n = 128, 128 - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(n, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.ones(n, hd, dtype=torch.bfloat16, device='cuda') - v = v.as_strided((n, hd), (1, n)).unsqueeze(-1) - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - qf = q[:,:,0].float(); kf = k[:,:,0].float() - ref = (qf @ kf.T).bfloat16().float() @ v[:,:,0].float() - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - kernel = FmhaV3Debug(head_dim=hd) - print(f'hd={hd}: Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - print(f'hd={hd}: pv_mma_tiler={kernel.pv_mma_tiler} p0={kernel.tmem_p0_offset} o0={kernel.tmem_o0_offset} alloc={kernel.num_tmem_alloc_cols}', flush=True) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - out = c[:,:,0].float() - cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - print(f'hd={hd} V=ones: cosine {cos:.6f} {"PASS" if cos >= 0.99 else "FAIL"}') - if cos < 0.99: - print(f' out[0,:4]={out[0,:4].tolist()} ref[0,:4]={ref[0,:4].tolist()}') - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_fmha_v3_diag.py b/tests/archive/test_fmha_v3_diag.py deleted file mode 100644 index 37bba2cb..00000000 --- a/tests/archive/test_fmha_v3_diag.py +++ /dev/null @@ -1,317 +0,0 @@ -""" -FMHA Stage-C multi-tile diagnostic. -Tests whether the combined K+V pipeline actually loads different KV tiles. -Runs identity softmax (no rescale, no online softmax) to isolate the -TMA + MMA pipeline from softmax logic. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct -import math - -HEAD_DIM = 64 - - -class FmhaV3Diag: - """Identity-softmax multi-tile with combined K+V barrier. - If this fails, the problem is in TMA/MMA, not softmax.""" - def __init__(self, s_k=256): - self.s_k = s_k - self.n_kv_tiles = s_k // 128 - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0,1,2,3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.num_c_stage = 2 - self.kv_stage = 2; self.q_stage = 1; self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (128, 128, qk_ik * 4) - pv_ik = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (128, HEAD_DIM, pv_ik * (128 // pv_ik)) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0]//cute.size(qk_mma.thr_id.shape), HEAD_DIM, self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - self.q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, self.q_dtype, self.q_stage) - self.k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, self.q_dtype, self.kv_stage) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, self.kv_stage) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - qk_thr = qk_mma.get_slice(0); qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - pv_thr = pv_mma.get_slice(0); pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - self.tmem_s0_offset = 0; self.tmem_p0_offset = 32 - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - p_end = self.tmem_p0_offset + p_cols_fp32 - s_cols = self.qk_mma_tiler[1] - o_after = max(s_cols, p_end) - self.tmem_o0_offset = ((o_after + 31) // 32) * 32 - o_cols = find_tmem_tensor_col_offset(tOtO) - total = self.tmem_o0_offset + o_cols - self.num_tmem_alloc_cols = 1 - while self.num_tmem_alloc_cols < total: - self.num_tmem_alloc_cols *= 2 - cta = cute.size(qk_mma.thr_id.shape) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)) - k_s = cute.slice_(self.k_smem_s,(None,None,None,0)) - v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - self.q_tx_bytes = cute.size_in_bytes(self.q_dtype, q_s) * cta - self.kv_tx_bytes = (cute.size_in_bytes(self.q_dtype, k_s) + - cute.size_in_bytes(self.q_dtype, v_s)) * cta - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - v_fmha = cute.make_tensor( - v.iterator, - cute.make_layout( - (HEAD_DIM, self.s_k, 1), - stride=(1, HEAD_DIM, HEAD_DIM * self.s_k), - ), - ) - self.v_major = LayoutEnum.from_tensor(v_fmha).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, self.a_major, self.b_major, self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, self.qk_acc_dtype, self.cta_group, (128,HEAD_DIM), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)); v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id),q,q_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id),k,k_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v_fmha,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape) - epi_s = cute.select(self.c_smem_s,mode=[0,1]) - tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile) - self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.v_smem_s,self.p_tmem_s,self.c_smem_s,self.epi_tile).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, q_smem_s, k_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx,_,_ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k); cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - q_bar: cute.struct.MemRange[cutlass.Int64, self.q_stage*2] - kv_bar: cute.struct.MemRange[cutlass.Int64, self.kv_stage*2] - s_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage*2] - tmem_dealloc: cutlass.Int64; holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - qp,qc = pipeline.PipelineTmaUmma.create(barrier_storage=st.q_bar.data_ptr(),num_stages=self.q_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.q_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - kvp,kvc = pipeline.PipelineTmaUmma.create(barrier_storage=st.kv_bar.data_ptr(),num_stages=self.kv_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.kv_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - s_prod,s_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.s_bar.data_ptr(),num_stages=1,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.epilogue_warp_id))).make_participants() - softmax_done_bar = pipeline.NamedBarrier(barrier_id=3, num_threads=32 + 32*len(self.epilogue_warp_id)) - acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(),num_stages=self.num_acc_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,len(self.epilogue_warp_id)),cta_layout_vmnk=cl_vmnk,defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2,num_threads=32*len((self.mma_warp_id,*self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr,barrier_for_retrieve=tmem_bar,allocator_warp_id=self.epilogue_warp_id[0],is_two_cta=cute.size(qk_mma.thr_id.shape)==2,two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk,is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype,layout=q_smem_s.outer,byte_alignment=128,swizzle=q_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype,layout=k_smem_s.outer,byte_alignment=128,swizzle=k_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype,layout=v_smem_s.outer,byte_alignment=128,swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype,layout=c_smem_s.outer,byte_alignment=128,swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None)) - gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None)) - gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None)) - gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None)) - n_kv_tiles = self.n_kv_tiles - - qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - tCgV = pv_thr.partition_B(gV); tCgC = pv_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,0,None,0)).shape) - tAsQ,tAgQ = cpasync.tma_partition(tma_q,0,a_lay,cute.group_modes(sQ,0,3),cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape) - tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3)) - tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3)) - # ===================================================================== - # ⚠️⚠️⚠️ CRITICAL: TMA PARTITION TENSOR MODE ORDERING ⚠️⚠️⚠️ - # After tma_partition, tBgK/tVgV have 4 modes: (V_grouped, ?, KV_tiles, ?) - # Mode 4 is the GMEM tile dimension. DO NOT pre-slice to fewer modes! - # See README.md for details. - # ===================================================================== - tAgQ = tAgQ[(None,0,None,0)] - tBgK = tBgK[(None,0,None,0)] - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None,None,None,0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # TMA LOAD (combined K+V barrier, Int32(kt) for GMEM coord) - if warp_idx == self.tma_warp_id: - qp.reset(); qh = qp.acquire_and_advance() - cute.copy(tma_q, tAgQ[(None, Int32(0))], tAsQ[(None, qh.index)], tma_bar_ptr=qh.barrier) - qp.tail() - kvp.reset(); pk = kvp.try_acquire() - # Force runtime Int32 (not literal) — option 3 from CUTLASS LLM - kv_coord = Int32(0 + 0) - for kt in cutlass.range(self.n_kv_tiles, unroll=1): - kvh = kvp.acquire_and_advance(pk) - # kv_coord indexes the surviving mode 1 (from original mode 2) of 2D tBgK/tVgV. - # Using (None, kv_coord) on a pre-sliced 4-mode tensor SILENTLY BREAKS multi-tile! - cute.copy(tma_k, tBgK[None, kv_coord], tBsK[(None, kvh.index)], tma_bar_ptr=kvh.barrier) - cute.copy(tma_v, tVgV[None, kv_coord], tVsV[(None, kvh.index)], tma_bar_ptr=kvh.barrier) - kv_coord += 1 - pk = cutlass.Boolean(1) - kvp.tail() - - # MMA (combined K+V slot) - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - qc.reset(); qh = qc.wait_and_advance(); qh.release() - kvc.reset(); pk = kvc.try_wait() - acc_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_stage) - acc_pipe.producer_acquire(acc_st) - for kt in range(self.n_kv_tiles): - kvh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - sh = s_prod.acquire_and_advance() - qk_mma.set(tcgen05.Field.ACCUMULATE, False) - for kb in cutlass.range(cute.size(tCrQ, mode=[2]), unroll_full=True): - cute.gemm(qk_mma, tStS0, tCrQ[(None,None,kb,0)], tCrK[(None,None,kb,kvh.index)], tStS0) - qk_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - sh.commit() - softmax_done_bar.arrive_and_wait() - pv_mma.set(tcgen05.Field.ACCUMULATE, kt != 0) - for kb in cutlass.range(cute.size(tOrP0, mode=[2]), unroll_full=True): - cute.gemm(pv_mma, tOtO0, tOrP0[(None,None,kb)], tCrV[(None,None,kb,kvh.index)], tOtO0) - pv_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - kvh.release() - acc_pipe.producer_commit(acc_st); acc_st.advance() - acc_pipe.producer_tail(acc_st) - - # SOFTMAX (identity: just load S, store as P, no scaling) - if warp_idx < self.mma_warp_id: - tmem.allocate(self.num_tmem_alloc_cols) - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) - sfw_idx = tidx % (32 * len(self.epilogue_warp_id)) - - # S load - tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS0) - thr_load = tiled_tmem_load.get_slice(sfw_idx) - tTMEM_LOADtS = thr_load.partition_S(tStS0) - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tTMEM_LOADcS = thr_load.partition_D(tScS) - - # P store - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - tStP_layout = cute.composition(tStS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tStP0 = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStP_layout) - tmem_store_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStP0) - thr_store = tiled_tmem_store.get_slice(sfw_idx) - tTMEM_STOREtP = thr_store.partition_D(tStP0) - tScP_layout = cute.composition(tScS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tScP = cute.make_tensor(tScS.iterator, tScP_layout) - tTMEM_STOREcP = thr_store.partition_S(tScP) - - for kt in range(self.n_kv_tiles): - si_handle = s_cons.wait_and_advance() - - # Load S from TMEM - tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS) - cute.arch.fence_view_async_tmem_load() - - # Identity: convert S to BF16 and store as P - rP_words = cute.make_rmem_tensor(tTMEM_STOREcP.shape, self.qk_acc_dtype) - rP_bf16 = cute.make_tensor(cute.recast_ptr(rP_words.iterator, dtype=self.q_dtype), tTMEM_LOADrS.layout) - frg_cnt = 4 - frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt - tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile)) - rP_bf16_frg = cute.logical_divide(rP_bf16, cute.make_layout(frg_tile)) - for j in range(frg_cnt): - s_vec = tTMEM_LOADrS_frg[None, j].load() - rP_bf16_frg[None, j].store(s_vec.to(self.q_dtype)) - - cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP) - cute.arch.fence_view_async_tmem_store() - si_handle.release() - softmax_done_bar.arrive() - - tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout) - acc_cons_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage) - c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) - c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp) - acc_cons_st = utils.gemm.sm100.epilogue_tma_store(self, tidx, warp_idx, tma_c, tCtO_base, sC, tCgC, epi_tile, 0, const_expr(lambda x: x), (0,0,0), acc_cons_st, acc_pipe, c_pipe) - c_pipe.producer_tail() - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def test(): - for n in [128, 256, 384]: - torch.manual_seed(42) - m, hd = 128, HEAD_DIM - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(n, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(n, hd, dtype=torch.bfloat16, device='cuda') - v_kernel = v.unsqueeze(-1) - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - - qf = q[:, :, 0].float() - kf = k[:, :, 0].float() - # Identity softmax reference: QK (no scale, no softmax) @ V - ref = (qf @ kf.T).bfloat16().float() @ v.float() - - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - kernel = FmhaV3Diag(s_k=n) - print(f'n={n}: Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - - out = c[:, :, 0].float() - cos = torch.nn.functional.cosine_similarity( - out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0) - ).item() - n_tiles = n // 128 - print(f'DIAG n={n} ({n_tiles} tiles): cos {cos:.6f} {"PASS" if cos >= 0.99 else "FAIL"}') - if cos < 0.99: - print(f' out[0,:4]={out[0,:4].tolist()}') - print(f' ref[0,:4]={ref[0,:4].tolist()}') - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_fmha_v3_diag_fixed.py b/tests/archive/test_fmha_v3_diag_fixed.py deleted file mode 100644 index 5251492c..00000000 --- a/tests/archive/test_fmha_v3_diag_fixed.py +++ /dev/null @@ -1,307 +0,0 @@ -""" -FMHA Stage-C multi-tile diagnostic. -Tests whether the combined K+V pipeline actually loads different KV tiles. -Runs identity softmax (no rescale, no online softmax) to isolate the -TMA + MMA pipeline from softmax logic. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct -import math - -HEAD_DIM = 64 - - -class FmhaV3Diag: - """Identity-softmax multi-tile with combined K+V barrier. - If this fails, the problem is in TMA/MMA, not softmax.""" - def __init__(self, s_k=256): - self.s_k = s_k - self.n_kv_tiles = s_k // 128 - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0,1,2,3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.num_c_stage = 2 - self.kv_stage = 2; self.q_stage = 1; self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (128, 128, qk_ik * 4) - pv_ik = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (128, HEAD_DIM, pv_ik * (128 // pv_ik)) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0]//cute.size(qk_mma.thr_id.shape), HEAD_DIM, self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - self.q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, self.q_dtype, self.q_stage) - self.k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, self.q_dtype, self.kv_stage) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, self.kv_stage) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - qk_thr = qk_mma.get_slice(0); qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - pv_thr = pv_mma.get_slice(0); pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - self.tmem_s0_offset = 0; self.tmem_p0_offset = 32 - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - p_end = self.tmem_p0_offset + p_cols_fp32 - s_cols = self.qk_mma_tiler[1] - o_after = max(s_cols, p_end) - self.tmem_o0_offset = ((o_after + 31) // 32) * 32 - o_cols = find_tmem_tensor_col_offset(tOtO) - total = self.tmem_o0_offset + o_cols - self.num_tmem_alloc_cols = 1 - while self.num_tmem_alloc_cols < total: - self.num_tmem_alloc_cols *= 2 - cta = cute.size(qk_mma.thr_id.shape) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)) - k_s = cute.slice_(self.k_smem_s,(None,None,None,0)) - v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - self.q_tx_bytes = cute.size_in_bytes(self.q_dtype, q_s) * cta - self.kv_tx_bytes = (cute.size_in_bytes(self.q_dtype, k_s) + - cute.size_in_bytes(self.q_dtype, v_s)) * cta - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - v_fmha = cute.make_tensor( - v.iterator, - cute.make_layout( - (HEAD_DIM, self.s_k, 1), - stride=(1, HEAD_DIM, HEAD_DIM * self.s_k), - ), - ) - self.v_major = LayoutEnum.from_tensor(v_fmha).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, self.a_major, self.b_major, self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, self.qk_acc_dtype, self.cta_group, (128,HEAD_DIM), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)); v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id),q,q_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id),k,k_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v_fmha,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape) - epi_s = cute.select(self.c_smem_s,mode=[0,1]) - tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile) - self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.v_smem_s,self.p_tmem_s,self.c_smem_s,self.epi_tile).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, q_smem_s, k_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx,_,_ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k); cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - q_bar: cute.struct.MemRange[cutlass.Int64, self.q_stage*2] - kv_bar: cute.struct.MemRange[cutlass.Int64, self.kv_stage*2] - s_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage*2] - tmem_dealloc: cutlass.Int64; holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - qp,qc = pipeline.PipelineTmaUmma.create(barrier_storage=st.q_bar.data_ptr(),num_stages=self.q_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.q_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - kvp,kvc = pipeline.PipelineTmaUmma.create(barrier_storage=st.kv_bar.data_ptr(),num_stages=self.kv_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.kv_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - s_prod,s_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.s_bar.data_ptr(),num_stages=1,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.epilogue_warp_id))).make_participants() - softmax_done_bar = pipeline.NamedBarrier(barrier_id=3, num_threads=32 + 32*len(self.epilogue_warp_id)) - acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(),num_stages=self.num_acc_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,len(self.epilogue_warp_id)),cta_layout_vmnk=cl_vmnk,defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2,num_threads=32*len((self.mma_warp_id,*self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr,barrier_for_retrieve=tmem_bar,allocator_warp_id=self.epilogue_warp_id[0],is_two_cta=cute.size(qk_mma.thr_id.shape)==2,two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk,is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype,layout=q_smem_s.outer,byte_alignment=128,swizzle=q_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype,layout=k_smem_s.outer,byte_alignment=128,swizzle=k_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype,layout=v_smem_s.outer,byte_alignment=128,swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype,layout=c_smem_s.outer,byte_alignment=128,swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None)) - gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None)) - gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None)) - gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None)) - n_kv_tiles = self.n_kv_tiles - - qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - tCgV = pv_thr.partition_B(gV); tCgC = pv_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,0,None,0)).shape) - tAsQ,tAgQ = cpasync.tma_partition(tma_q,0,a_lay,cute.group_modes(sQ,0,3),cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape) - tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3)) - tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,None,None,None,None,None,None,None)]; tVgV = tVgV[(None,None,None,None,None,None,None,None)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None,None,None,0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # TMA LOAD (combined K+V barrier, Int32(kt) for GMEM coord) - if warp_idx == self.tma_warp_id: - qp.reset(); qh = qp.acquire_and_advance() - cute.copy(tma_q, tAgQ[(None, Int32(0))], tAsQ[(None, qh.index)], tma_bar_ptr=qh.barrier) - qp.tail() - kvp.reset(); pk = kvp.try_acquire() - # Force runtime Int32 (not literal) — option 3 from CUTLASS LLM - kv_coord = Int32(0 + 0) - for kt in cutlass.range(self.n_kv_tiles, unroll=1): - kvh = kvp.acquire_and_advance(pk) - cute.copy(tma_k, tBgK[None, None, None, None, kv_coord, None, None, None], tBsK[(None, kvh.index)], tma_bar_ptr=kvh.barrier) - cute.copy(tma_v, tVgV[None, None, None, None, kv_coord, None, None, None], tVsV[(None, kvh.index)], tma_bar_ptr=kvh.barrier) - kv_coord += 1 - pk = cutlass.Boolean(1) - kvp.tail() - - # MMA (combined K+V slot) - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - qc.reset(); qh = qc.wait_and_advance(); qh.release() - kvc.reset(); pk = kvc.try_wait() - acc_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_stage) - acc_pipe.producer_acquire(acc_st) - for kt in range(self.n_kv_tiles): - kvh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - sh = s_prod.acquire_and_advance() - qk_mma.set(tcgen05.Field.ACCUMULATE, False) - for kb in cutlass.range(cute.size(tCrQ, mode=[2]), unroll_full=True): - cute.gemm(qk_mma, tStS0, tCrQ[(None,None,kb,0)], tCrK[(None,None,kb,kvh.index)], tStS0) - qk_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - sh.commit() - softmax_done_bar.arrive_and_wait() - pv_mma.set(tcgen05.Field.ACCUMULATE, kt != 0) - for kb in cutlass.range(cute.size(tOrP0, mode=[2]), unroll_full=True): - cute.gemm(pv_mma, tOtO0, tOrP0[(None,None,kb)], tCrV[(None,None,kb,kvh.index)], tOtO0) - pv_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - kvh.release() - acc_pipe.producer_commit(acc_st); acc_st.advance() - acc_pipe.producer_tail(acc_st) - - # SOFTMAX (identity: just load S, store as P, no scaling) - if warp_idx < self.mma_warp_id: - tmem.allocate(self.num_tmem_alloc_cols) - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) - sfw_idx = tidx % (32 * len(self.epilogue_warp_id)) - - # S load - tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS0) - thr_load = tiled_tmem_load.get_slice(sfw_idx) - tTMEM_LOADtS = thr_load.partition_S(tStS0) - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tTMEM_LOADcS = thr_load.partition_D(tScS) - - # P store - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - tStP_layout = cute.composition(tStS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tStP0 = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStP_layout) - tmem_store_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStP0) - thr_store = tiled_tmem_store.get_slice(sfw_idx) - tTMEM_STOREtP = thr_store.partition_D(tStP0) - tScP_layout = cute.composition(tScS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tScP = cute.make_tensor(tScS.iterator, tScP_layout) - tTMEM_STOREcP = thr_store.partition_S(tScP) - - for kt in range(self.n_kv_tiles): - si_handle = s_cons.wait_and_advance() - - # Load S from TMEM - tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS) - cute.arch.fence_view_async_tmem_load() - - # Identity: convert S to BF16 and store as P - rP_words = cute.make_rmem_tensor(tTMEM_STOREcP.shape, self.qk_acc_dtype) - rP_bf16 = cute.make_tensor(cute.recast_ptr(rP_words.iterator, dtype=self.q_dtype), tTMEM_LOADrS.layout) - frg_cnt = 4 - frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt - tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile)) - rP_bf16_frg = cute.logical_divide(rP_bf16, cute.make_layout(frg_tile)) - for j in range(frg_cnt): - s_vec = tTMEM_LOADrS_frg[None, j].load() - rP_bf16_frg[None, j].store(s_vec.to(self.q_dtype)) - - cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP) - cute.arch.fence_view_async_tmem_store() - si_handle.release() - softmax_done_bar.arrive() - - tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout) - acc_cons_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage) - c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) - c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp) - acc_cons_st = utils.gemm.sm100.epilogue_tma_store(self, tidx, warp_idx, tma_c, tCtO_base, sC, tCgC, epi_tile, 0, const_expr(lambda x: x), (0,0,0), acc_cons_st, acc_pipe, c_pipe) - c_pipe.producer_tail() - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def test(): - for n in [128, 256, 384]: - torch.manual_seed(42) - m, hd = 128, HEAD_DIM - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(n, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(n, hd, dtype=torch.bfloat16, device='cuda') - v_kernel = v.unsqueeze(-1) - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - - qf = q[:, :, 0].float() - kf = k[:, :, 0].float() - # Identity softmax reference: QK (no scale, no softmax) @ V - ref = (qf @ kf.T).bfloat16().float() @ v.float() - - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - kernel = FmhaV3Diag(s_k=n) - print(f'n={n}: Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - - out = c[:, :, 0].float() - cos = torch.nn.functional.cosine_similarity( - out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0) - ).item() - n_tiles = n // 128 - print(f'DIAG n={n} ({n_tiles} tiles): cos {cos:.6f} {"PASS" if cos >= 0.99 else "FAIL"}') - if cos < 0.99: - print(f' out[0,:4]={out[0,:4].tolist()}') - print(f' ref[0,:4]={ref[0,:4].tolist()}') - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_full_layer_b200.py b/tests/archive/test_full_layer_b200.py deleted file mode 100644 index c4db4920..00000000 --- a/tests/archive/test_full_layer_b200.py +++ /dev/null @@ -1,258 +0,0 @@ -#!/usr/bin/env python3 -""" -Full decoder layer 0 test: ALL components using CuTeDSL kernels, NO vLLM. - -Tests each attention + FFN projection individually (CuTeDSL vs BF16 ref). - -Usage (on B200): - source /root/nvfp4-megamoe-kernel/tests/.venv/bin/activate - python3 tests/test_full_layer_b200.py -""" - -import sys, os, json, math, torch, torch.nn.functional as F -from safetensors import safe_open - -REPO = "/root/nvfp4-megamoe-kernel" -sys.path.insert(0, REPO) -MODEL = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4" -DEV = "cuda:0" - -# Model config (layer 0, compress_ratio=128 → C4A) -H = 7168 -NH = 128 -HD = 512 -NOPE = 448 -ROPE = 64 -QL = 1536 -OL = 1024 -OG = 16 -HPG = NH // OG -HC = 4 -SL = 10.0 -EPS = 1e-6 -INTER = 3072 -NT = 4 - -E2M1 = torch.tensor([0,.5,1.,1.5,2.,3.,4.,6.,-0,-.5,-1.,-1.5,-2.,-3.,-4.,-6.], dtype=torch.float32) - -_cache = {} -def P(k, wm, md): - if k in _cache: return _cache[k] - with safe_open(os.path.join(md, wm[k]), framework="pt") as f: - t = f.get_tensor(k) - _cache[k] = t - return t - -def dequant(w, sf, gs): - d = w.device; lut = E2M1.to(d) - lo = lut[(w & 0xF).long()]; hi = lut[((w >> 4) & 0xF).long()] - O, I2 = w.shape; I = I2*2 - u = torch.empty(O, I, dtype=torch.float32, device=d) - u[:,0::2] = lo; u[:,1::2] = hi - bs = sf.float().repeat_interleave(16, dim=1)[:O,:I] - return (u * bs * gs).to(torch.bfloat16) - -def rms(x, w, eps=1e-6): - v = x.float().pow(2).mean(-1, keepdim=True) - return (w.float() * (x * torch.rsqrt(v+eps)).float()).to(x.dtype) - -def make_runner(w, sf, gs_t, inf, outf, fused=False, lw=None): - from dsv4.layers.linear import Nvfp4Linear - fp4 = w.view(torch.float4_e2m1fn_x2).permute(1,0).contiguous() - s = sf.to(torch.float8_e4m3fn) if sf.dtype != torch.float8_e4m3fn else sf - s = s.permute(1,0).contiguous() - if fused and gs_t.numel() == 2: - g1,g2 = gs_t[0].item(), gs_t[1].item(); gs = max(g1,g2) - if g1 != g2: - s32 = s.float(); sp = lw[0] if lw else outf//2 - s32[:sp] *= g1/gs; s32[sp:] *= g2/gs; s = s32.to(torch.float8_e4m3fn) - else: - gs = gs_t.max().item() if gs_t.numel() > 1 else gs_t.item() - r = Nvfp4Linear(in_features=inf, out_features=outf, max_num_tokens=8192, device=str(w.device)) - r.fp4 = [fp4]; r.sf = [s]; r.gs = [gs] - r.finalize_weights(); r._ensure_initialized() - return r - -def cosim(a, b): - return F.cosine_similarity(a.flatten().unsqueeze(0).float(), b.flatten().unsqueeze(0).float().to(a.device)).item() - -def main(): - torch.cuda.set_device(0); torch.manual_seed(42) - print("="*70) - print(" Layer 0: CuTeDSL NVFP4 vs BF16 Reference") - print("="*70) - - with open(os.path.join(MODEL, "model.safetensors.index.json")) as f: - wm = json.load(f)["weight_map"] - G = lambda k: P(k, wm, MODEL).to(DEV) - - p = "model.layers.0"; a = f"{p}.self_attn"; m = f"{p}.mlp" - - # ── Load attention weights (correct checkpoint key names) ───────── - print("\n--- Loading weights ---") - qa_w = G(f"{a}.q_a_proj.weight"); qa_sf = G(f"{a}.q_a_proj.weight_scale"); qa_gs = G(f"{a}.q_a_proj.weight_scale_2") - qb_w = G(f"{a}.q_b_proj.weight"); qb_sf = G(f"{a}.q_b_proj.weight_scale"); qb_gs = G(f"{a}.q_b_proj.weight_scale_2") - kv_w = G(f"{a}.kv_proj.weight"); kv_sf = G(f"{a}.kv_proj.weight_scale"); kv_gs = G(f"{a}.kv_proj.weight_scale_2") - woa = G(f"{a}.o_a_proj.weight") # BF16 - wob_w = G(f"{a}.o_b_proj.weight"); wob_sf = G(f"{a}.o_b_proj.weight_scale"); wob_gs = G(f"{a}.o_b_proj.weight_scale_2") - qn = G(f"{a}.q_a_norm.weight"); kvn = G(f"{a}.kv_norm.weight") - anorm = G(f"{p}.input_layernorm.weight"); fnorm = G(f"{p}.post_attention_layernorm.weight") - - # Compressor (C4A path) - ckv_w = G(f"{a}.compressor.kv_proj.weight"); ckv_sf = G(f"{a}.compressor.kv_proj.weight_scale"); ckv_gs = G(f"{a}.compressor.kv_proj.weight_scale_2") - cg_w = G(f"{a}.compressor.gate_proj.weight"); cg_sf = G(f"{a}.compressor.gate_proj.weight_scale"); cg_gs = G(f"{a}.compressor.gate_proj.weight_scale_2") - ckn = G(f"{a}.compressor.kv_norm.weight") - cpb = G(f"{a}.compressor.position_bias") - sinks = G(f"{a}.sinks") - - # MHC - hca_fn = G(f"{p}.attn_hc.fn"); hcf_fn = G(f"{p}.ffn_hc.fn") - hca_b = G(f"{p}.attn_hc.base"); hcf_b = G(f"{p}.ffn_hc.base") - hca_s = G(f"{p}.attn_hc.scale"); hcf_s = G(f"{p}.ffn_hc.scale") - - for nm, t in [("q_a_proj", qa_w), ("q_b_proj", qb_w), ("kv_proj", kv_w), - ("o_a_proj", woa), ("o_b_proj", wob_w), - ("comp.kv_proj", ckv_w), ("comp.gate_proj", cg_w), - ("sinks", sinks), ("comp.position_bias", cpb), - ("attn_hc.fn", hca_fn)]: - print(f" {nm}: shape={t.shape} dtype={t.dtype}") - - # ── Create CuTeDSL runners ──────────────────────────────────────── - print("\n--- Creating CuTeDSL runners ---") - r_qa = make_runner(qa_w, qa_sf, qa_gs, qa_w.shape[1]*2, qa_w.shape[0]) - r_qb = make_runner(qb_w, qb_sf, qb_gs, qb_w.shape[1]*2, qb_w.shape[0]) - r_kv = make_runner(kv_w, kv_sf, kv_gs, kv_w.shape[1]*2, kv_w.shape[0]) - r_wob = make_runner(wob_w, wob_sf, wob_gs, wob_w.shape[1]*2, wob_w.shape[0]) - - # Compressor runners - r_ckv = make_runner(ckv_w, ckv_sf, ckv_gs, ckv_w.shape[1]*2, ckv_w.shape[0]) - r_cg = make_runner(cg_w, cg_sf, cg_gs, cg_w.shape[1]*2, cg_w.shape[0]) - - print(f" q_a: in={qa_w.shape[1]*2} out={qa_w.shape[0]}") - print(f" q_b: in={qb_w.shape[1]*2} out={qb_w.shape[0]}") - print(f" kv: in={kv_w.shape[1]*2} out={kv_w.shape[0]}") - print(f" wo_b: in={wob_w.shape[1]*2} out={wob_w.shape[0]}") - print(f" comp.kv: in={ckv_w.shape[1]*2} out={ckv_w.shape[0]}") - print(f" comp.gate: in={cg_w.shape[1]*2} out={cg_w.shape[0]}") - - # Warmup - print(" Warming up...") - d1 = torch.randn(NT, H, dtype=torch.bfloat16, device=DEV)*2.0 - for r in [r_qa, r_kv, r_ckv, r_cg]: - r.compute_activation_global_scale(d1) - d2 = torch.randn(NT, QL, dtype=torch.bfloat16, device=DEV)*2.0 - r_qb.compute_activation_global_scale(d2) - d3 = torch.randn(NT, OG*OL, dtype=torch.bfloat16, device=DEV)*2.0 - r_wob.compute_activation_global_scale(d3) - print(" Done.") - - # ── Per-projection BF16 vs CuTeDSL comparison ──────────────────── - print("\n" + "="*70) - print(" PROJECTION-LEVEL: CuTeDSL vs BF16") - print("="*70) - torch.manual_seed(123) - tx = torch.randn(NT, H, dtype=torch.bfloat16, device=DEV)*2.0 - - results = {} - - # q_a_proj - with torch.no_grad(): co = r_qa.run(tx) - ref = tx @ dequant(qa_w, qa_sf, qa_gs.item()).T - c = cosim(co, ref); results['q_a_proj'] = c - print(f" q_a_proj: cosine={c:.6f} {'✅' if c>=0.98 else '❌'} amax={co.amax():.4f} ref={ref.amax():.4f}") - - # kv_proj - with torch.no_grad(): co = r_kv.run(tx) - ref = tx @ dequant(kv_w, kv_sf, kv_gs.item()).T - c = cosim(co, ref); results['kv_proj'] = c - print(f" kv_proj: cosine={c:.6f} {'✅' if c>=0.98 else '❌'} amax={co.amax():.4f} ref={ref.amax():.4f}") - - # q_b_proj - tq = torch.randn(NT, QL, dtype=torch.bfloat16, device=DEV)*2.0 - with torch.no_grad(): co = r_qb.run(tq) - ref = tq @ dequant(qb_w, qb_sf, qb_gs.item()).T - c = cosim(co, ref); results['q_b_proj'] = c - print(f" q_b_proj: cosine={c:.6f} {'✅' if c>=0.98 else '❌'} amax={co.amax():.4f} ref={ref.amax():.4f}") - - # wo_b_proj - tz = torch.randn(NT, OG*OL, dtype=torch.bfloat16, device=DEV)*2.0 - with torch.no_grad(): co = r_wob.run(tz) - ref = tz @ dequant(wob_w, wob_sf, wob_gs.item()).T - c = cosim(co, ref); results['wo_b_proj'] = c - print(f" wo_b_proj: cosine={c:.6f} {'✅' if c>=0.98 else '❌'} amax={co.amax():.4f} ref={ref.amax():.4f}") - - # compressor kv_proj - with torch.no_grad(): co = r_ckv.run(tx) - ref = tx @ dequant(ckv_w, ckv_sf, ckv_gs.item()).T - c = cosim(co, ref); results['comp.kv_proj'] = c - print(f" comp.kv_proj: cosine={c:.6f} {'✅' if c>=0.98 else '❌'} amax={co.amax():.4f} ref={ref.amax():.4f}") - - # compressor gate_proj - with torch.no_grad(): co = r_cg.run(tx) - ref = tx @ dequant(cg_w, cg_sf, cg_gs.item()).T - c = cosim(co, ref); results['comp.gate_proj'] = c - print(f" comp.gate: cosine={c:.6f} {'✅' if c>=0.98 else '❌'} amax={co.amax():.4f} ref={ref.amax():.4f}") - - # ── Shared expert ───────────────────────────────────────────────── - print("\n--- Shared Expert: CuTeDSL vs BF16 ---") - from dsv4.layers.shared_expert import Nvfp4SharedExpert - - sgw = G(f"{m}.shared_experts.gate_proj.weight"); sgsf = G(f"{m}.shared_experts.gate_proj.weight_scale") - sggs = G(f"{m}.shared_experts.gate_proj.weight_scale_2").item() - suw = G(f"{m}.shared_experts.up_proj.weight"); susf = G(f"{m}.shared_experts.up_proj.weight_scale") - sugs = G(f"{m}.shared_experts.up_proj.weight_scale_2").item() - sdw = G(f"{m}.shared_experts.down_proj.weight"); sdsf = G(f"{m}.shared_experts.down_proj.weight_scale") - sdgs = G(f"{m}.shared_experts.down_proj.weight_scale_2").item() - - si = INTER - sgu_w = torch.cat([sgw, suw], 0); sgu_sf = torch.cat([sgsf, susf], 0) - smgs = max(sggs, sugs) - if sggs != sugs: - s32 = sgu_sf.float(); s32[:si] *= sggs/smgs; s32[si:] *= sugs/smgs - sgu_sf = s32.to(torch.float8_e4m3fn) - - ser = Nvfp4SharedExpert(hidden_size=H, intermediate_size=si, max_num_tokens=8192, - device=DEV, swiglu_limit=SL) - ser.l1_fp4 = [sgu_w.view(torch.float4_e2m1fn_x2).permute(1,0).contiguous()] - ser.l1_sf = [sgu_sf.permute(1,0).contiguous()]; ser.l1_gs = [smgs] - ser.l2_fp4 = [sdw.view(torch.float4_e2m1fn_x2).permute(1,0).contiguous()] - ser.l2_sf = [sdsf.permute(1,0).contiguous()]; ser.l2_gs = [sdgs] - ser.finalize_weights(); ser._ensure_initialized() - - tse = torch.randn(NT, H, dtype=torch.bfloat16, device=DEV)*2.0 - ser.compute_activation_global_scales(tse) - with torch.no_grad(): so = ser.run(tse) - - gb = dequant(sgw, sgsf, sggs); ub = dequant(suw, susf, sugs); db = dequant(sdw, sdsf, sdgs) - with torch.no_grad(): - g_ = tse @ gb.T; u_ = tse @ ub.T - act = F.silu(g_.clamp(max=SL)) * u_.clamp(min=-SL, max=SL) - sref = act @ db.T - c = cosim(so, sref); results['shared_expert'] = c - print(f" shared_expert: cosine={c:.6f} {'✅' if c>=0.98 else '❌'} amax={so.amax():.4f} ref={sref.amax():.4f}") - - # ── MHC sanity ──────────────────────────────────────────────────── - print("\n--- MHC weight shapes ---") - print(f" attn_hc.fn: {hca_fn.shape} ffn_hc.fn: {hcf_fn.shape}") - print(f" attn_hc.base: {hca_b.shape} ffn_hc.base: {hcf_b.shape}") - print(f" attn_hc.scale: {hca_s.shape} ffn_hc.scale: {hcf_s.shape}") - print(f" attn_hc.scale values: {hca_s.tolist()}") - print(f" sinks: {sinks.shape}") - - # ── Summary ─────────────────────────────────────────────────────── - print("\n" + "="*70) - print(" SUMMARY") - print("="*70) - all_pass = True - for name, c in results.items(): - status = '✅' if c >= 0.98 else '❌' - if c < 0.98: all_pass = False - print(f" {name}: {c:.6f} {status}") - if all_pass: - print("\n All projections pass! CuTeDSL kernels match BF16 reference.") - print(" The bug is in vLLM's pipeline, not our kernels.") - else: - print("\n Some projections FAIL. Need to debug those specific kernels.") - -if __name__ == "__main__": - main() diff --git a/tests/archive/test_full_layer_nan_b200.py b/tests/archive/test_full_layer_nan_b200.py deleted file mode 100644 index 87aa5c67..00000000 --- a/tests/archive/test_full_layer_nan_b200.py +++ /dev/null @@ -1,348 +0,0 @@ -#!/usr/bin/env python3 -""" -DeepSeek-V4 Full Layer Forward Test - -Tests a complete transformer layer (attention + MoE) with real weights. -If this produces NaN, we can bisect which component causes it. - -Usage (on B200): - cd /root/nvfp4-megamoe-kernel - PYTHONPATH=/root/nvfp4-megamoe-kernel tests/venv/bin/python tests/test_full_layer_nan_b200.py -""" - -import sys, os, json, torch, torch.nn.functional as F -from safetensors import safe_open - -REPO = "/root/nvfp4-megamoe-kernel" -sys.path.insert(0, REPO) -MODEL = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4" -DEV = "cuda:0" - -H = 7168; NH = 128; HD = 512; NOPE = 448; ROPE = 64 -QL = 1536; OL = 1024; OG = 16; HPG = NH // OG -INTERMEDIATE = 3072 -TOPK = 6 -EPS = 1e-6; WINDOW = 128; SCALE = HD ** -0.5 - -E2M1 = torch.tensor([0,.5,1.,1.5,2.,3.,4.,6.,-0,-.5,-1.,-1.5,-2.,-3.,-4.,-6.], dtype=torch.float32) - -_cache = {} -def P(k, wm, md): - if k in _cache: return _cache[k] - with safe_open(os.path.join(md, wm[k]), framework="pt") as f: - t = f.get_tensor(k) - _cache[k] = t - return t - -def rms(x, w, eps=1e-6): - v = x.float().pow(2).mean(-1, keepdim=True) - return (w.float() * (x * torch.rsqrt(v+eps)).float()).to(x.dtype) - -def make_runner(w, sf, gs_t, inf, outf): - from dsv4.layers.linear import Nvfp4Linear - fp4 = w.view(torch.float4_e2m1fn_x2).permute(1,0).contiguous() - s = sf.to(torch.float8_e4m3fn) if sf.dtype != torch.float8_e4m3fn else sf - s = s.permute(1,0).contiguous() - gs = gs_t.max().item() if gs_t.numel() > 1 else gs_t.item() - r = Nvfp4Linear(in_features=inf, out_features=outf, max_num_tokens=8192, device=str(w.device)) - r.fp4 = [fp4]; r.sf = [s]; r.gs = [gs] - r.finalize_weights(); r._ensure_initialized() - return r - -def build_cos_sin(max_pos=4096, rope_dim=ROPE): - half = rope_dim // 2 - inv_freq = 1.0 / (10000.0 ** (torch.arange(0, half, dtype=torch.float32) / half)) - freqs = torch.outer(torch.arange(max_pos, dtype=torch.float32), inv_freq) - return torch.cat([freqs.cos(), freqs.sin()], dim=-1) - -def apply_gptj_rope(x, positions, cos_sin, nope_dim, rope_dim): - if rope_dim == 0 or x.numel() == 0: return x - half = rope_dim // 2 - cos = cos_sin[positions, :half].to(x.dtype) - sin = cos_sin[positions, half:2*half].to(x.dtype) - if x.dim() == 3: cos = cos.unsqueeze(1); sin = sin.unsqueeze(1) - x_rope = x[..., nope_dim:].clone() - even = x_rope[..., 0::2]; odd = x_rope[..., 1::2] - out = x.clone() - out[..., nope_dim:][..., 0::2] = even * cos - odd * sin - out[..., nope_dim:][..., 1::2] = even * sin + odd * cos - return out - -def apply_inv_gptj_rope(x, positions, cos_sin, nope_dim, rope_dim): - if rope_dim == 0 or x.numel() == 0: return x - half = rope_dim // 2 - cos = cos_sin[positions, :half].to(x.dtype) - sin = cos_sin[positions, half:2*half].to(x.dtype) - if x.dim() == 3: cos = cos.unsqueeze(1); sin = sin.unsqueeze(1) - x_rope = x[..., nope_dim:].clone() - even = x_rope[..., 0::2]; odd = x_rope[..., 1::2] - out = x.clone() - out[..., nope_dim:][..., 0::2] = even * cos + odd * sin - out[..., nope_dim:][..., 1::2] = -even * sin + odd * cos - return out - -def kv_quantize_fp8(kv_bf16): - amax = kv_bf16.float().abs().amax(dim=-1, keepdim=True).clamp(min=1e-12) - fp8_max = torch.tensor(448.0, dtype=torch.float32, device=kv_bf16.device) - scale = fp8_max / amax - kv_fp8 = (kv_bf16.float() * scale).to(torch.float8_e4m3fn) - inv_scale = (amax / fp8_max).to(torch.bfloat16) - return kv_fp8, inv_scale - -def kv_dequantize_fp8(kv_fp8, inv_scale): - return (kv_fp8.to(torch.bfloat16) * inv_scale).to(torch.bfloat16) - -def causal_prefill_attention(q, kv, scale): - T, NH, HD = q.shape - q_t = q.permute(1, 0, 2) - kv_exp = kv.unsqueeze(0).expand(NH, -1, -1) - out = F.scaled_dot_product_attention(q_t, kv_exp, kv_exp, is_causal=True, scale=scale) - return out.permute(1, 0, 2) - - -def test_full_layer(layer_id, num_tokens=8, num_moe_experts=16): - """Test a complete transformer layer with attention + MoE.""" - from dsv4.layers.moe import Nvfp4MoE - - torch.cuda.set_device(0) - torch.manual_seed(42) - torch.cuda.empty_cache() - _cache.clear() - - with open(os.path.join(MODEL, "model.safetensors.index.json")) as f: - wm = json.load(f)["weight_map"] - G = lambda k: P(k, wm, MODEL).to(DEV) - - p = f"model.layers.{layer_id}"; a = f"{p}.self_attn" - m = f"{p}.mlp" - cr = 128 if layer_id == 0 else (0 if layer_id == 60 else 4) - lt = f"C{cr}A" if cr > 1 else "SWA" - - emb = G("model.embed_tokens.weight") - anorm = G(f"{p}.input_layernorm.weight") - qn = G(f"{a}.q_a_norm.weight"); kvn = G(f"{a}.kv_norm.weight") - woa = G(f"{a}.o_a_proj.weight") - fnorm = G(f"{p}.post_attention_layernorm.weight") - - qa_w = G(f"{a}.q_a_proj.weight"); qa_sf = G(f"{a}.q_a_proj.weight_scale"); qa_gs = G(f"{a}.q_a_proj.weight_scale_2") - qb_w = G(f"{a}.q_b_proj.weight"); qb_sf = G(f"{a}.q_b_proj.weight_scale"); qb_gs = G(f"{a}.q_b_proj.weight_scale_2") - kv_w = G(f"{a}.kv_proj.weight"); kv_sf = G(f"{a}.kv_proj.weight_scale"); kv_gs = G(f"{a}.kv_proj.weight_scale_2") - wob_w = G(f"{a}.o_b_proj.weight"); wob_sf = G(f"{a}.o_b_proj.weight_scale"); wob_gs = G(f"{a}.o_b_proj.weight_scale_2") - - r_qa = make_runner(qa_w, qa_sf, qa_gs, H, qa_w.shape[0]) - r_qb = make_runner(qb_w, qb_sf, qb_gs, QL, qb_w.shape[0]) - r_kv = make_runner(kv_w, kv_sf, kv_gs, H, kv_w.shape[0]) - r_wob = make_runner(wob_w, wob_sf, wob_gs, OG*OL, wob_w.shape[0]) - cos_sin = build_cos_sin(max_pos=4096).to(DEV) - woa_3d = woa.view(OG, OL, HPG * HD) - - # MoE weights (only first num_moe_experts to fit in memory) - gate_ws, gate_sfs, gate_gss = [], [], [] - up_ws, up_sfs, up_gss = [], [], [] - down_ws, down_sfs, down_gss = [], [], [] - for i in range(num_moe_experts): - e = f"{m}.experts.{i}" - gate_ws.append(G(f"{e}.gate_proj.weight")) - gate_sfs.append(G(f"{e}.gate_proj.weight_scale")) - gate_gss.append(G(f"{e}.gate_proj.weight_scale_2")) - up_ws.append(G(f"{e}.up_proj.weight")) - up_sfs.append(G(f"{e}.up_proj.weight_scale")) - up_gss.append(G(f"{e}.up_proj.weight_scale_2")) - down_ws.append(G(f"{e}.down_proj.weight")) - down_sfs.append(G(f"{e}.down_proj.weight_scale")) - down_gss.append(G(f"{e}.down_proj.weight_scale_2")) - - w13_w = torch.cat([torch.stack(gate_ws), torch.stack(up_ws)], dim=1) - w13_sf = torch.cat([torch.stack(gate_sfs), torch.stack(up_sfs)], dim=1) - w13_gs = torch.cat([torch.stack(gate_gss), torch.stack(up_gss)], dim=0) - w2_w = torch.stack(down_ws) - w2_sf = torch.stack(down_sfs) - w2_gs = torch.stack(down_gss) - - # Free per-expert lists - del gate_ws, gate_sfs, gate_gss, up_ws, up_sfs, up_gss, down_ws, down_sfs, down_gss - - moe_runner = Nvfp4MoE( - num_experts=num_moe_experts, - hidden_size=H, - intermediate_size=INTERMEDIATE, - max_num_tokens=8192, - top_k=TOPK, - device=str(DEV), - ) - - l1_fp4 = w13_w.view(torch.float4_e2m1fn_x2) - l2_fp4 = w2_w.view(torch.float4_e2m1fn_x2) - l1_sf = w13_sf.to(torch.float8_e4m3fn) - l2_sf = w2_sf.to(torch.float8_e4m3fn) - - moe_runner.prepare_weights_from_stacked( - l1_fp4, l1_sf, w13_gs.flatten().tolist(), - l2_fp4, l2_sf, w2_gs.flatten().tolist(), - ) - - del w13_w, w13_sf, w13_gs, w2_w, w2_sf, w2_gs, l1_fp4, l2_fp4, l1_sf, l2_sf - torch.cuda.empty_cache() - - # Shared expert - se_gate_w = G(f"{m}.shared_experts.gate_proj.weight"); se_gate_sf = G(f"{m}.shared_experts.gate_proj.weight_scale"); se_gate_gs = G(f"{m}.shared_experts.gate_proj.weight_scale_2") - se_up_w = G(f"{m}.shared_experts.up_proj.weight"); se_up_sf = G(f"{m}.shared_experts.up_proj.weight_scale"); se_up_gs = G(f"{m}.shared_experts.up_proj.weight_scale_2") - se_down_w = G(f"{m}.shared_experts.down_proj.weight"); se_down_sf = G(f"{m}.shared_experts.down_proj.weight_scale"); se_down_gs = G(f"{m}.shared_experts.down_proj.weight_scale_2") - - r_se_gate = make_runner(se_gate_w, se_gate_sf, se_gate_gs, H, se_gate_w.shape[0]) - r_se_up = make_runner(se_up_w, se_up_sf, se_up_gs, H, se_up_w.shape[0]) - r_se_down = make_runner(se_down_w, se_down_sf, se_down_gs, INTERMEDIATE, se_down_w.shape[0]) - - # Run the layer - token_ids = torch.randint(1, 1000, (num_tokens,), dtype=torch.long, device=DEV) - positions = torch.arange(num_tokens, dtype=torch.int64, device=DEV) - - with torch.no_grad(): - hidden = emb[token_ids] - - # ── Attention ────────────────────────────────────────── - normed = rms(hidden, anorm, EPS) - qa = r_qa.run(normed); kv = r_kv.run(normed) - qa_n = rms(qa, qn, EPS); kv_n = rms(kv, kvn, EPS) - q = r_qb.run(qa_n).view(num_tokens, NH, HD) - q_rope = apply_gptj_rope(q, positions, cos_sin, NOPE, ROPE) - kv_rope = apply_gptj_rope(kv_n.unsqueeze(1), positions, cos_sin, NOPE, ROPE).squeeze(1) - - o_attn = causal_prefill_attention(q_rope, kv_rope, SCALE) - o_inv = apply_inv_gptj_rope(o_attn, positions, cos_sin, NOPE, ROPE) - o_grouped = o_inv.reshape(num_tokens, OG, HPG * HD).permute(1, 0, 2) - z = torch.bmm(o_grouped, woa_3d.transpose(1, 2)).permute(1, 0, 2).reshape(num_tokens, OG * OL) - attn_out = r_wob.run(z) - - hidden = hidden + attn_out - print(f" Layer {layer_id} ({lt}): after attention: amax={hidden.amax():.4f} NaN={torch.isnan(hidden).any()}") - - # ── MoE ──────────────────────────────────────────────── - fnormed = rms(hidden, fnorm, EPS) - - # Shared expert - gate_out = r_se_gate.run(fnormed) - up_out = r_se_up.run(fnormed) - activated = F.silu(gate_out) * up_out - se_out = r_se_down.run(activated) - - # Routed experts (using MoE runner with subset of experts) - topk_ids = torch.randint(0, num_moe_experts, (num_tokens, TOPK), device=DEV) - topk_weights = torch.softmax(torch.randn(num_tokens, TOPK, device=DEV), dim=-1) - moe_out = moe_runner.run(fnormed, topk_weights, topk_ids) - - hidden = hidden + se_out + moe_out - print(f" Layer {layer_id} ({lt}): after MoE: amax={hidden.amax():.4f} NaN={torch.isnan(hidden).any()}") - - del r_qa, r_qb, r_kv, r_wob, r_se_gate, r_se_up, r_se_down, moe_runner - torch.cuda.empty_cache() - _cache.clear() - - return not torch.isnan(hidden).any() - - -def test_multi_layer(): - """Test multiple layers chained together to see if NaN propagates.""" - emb = None - - # Load embedding once - with open(os.path.join(MODEL, "model.safetensors.index.json")) as f: - wm = json.load(f)["weight_map"] - G = lambda k: P(k, wm, MODEL).to(DEV) - emb = G("model.embed_tokens.weight") - - num_tokens = 8 - token_ids = torch.randint(1, 1000, (num_tokens,), dtype=torch.long, device=DEV) - hidden = emb[token_ids] - - # Test just layers 0, 2, 60 (one of each type) - # For each layer, do attention only (skip MoE to save memory) - for layer_id in [0, 2, 60]: - p = f"model.layers.{layer_id}"; a = f"{p}.self_attn" - cr = 128 if layer_id == 0 else (0 if layer_id == 60 else 4) - lt = f"C{cr}A" if cr > 1 else "SWA" - - anorm = G(f"{p}.input_layernorm.weight") - qn = G(f"{a}.q_a_norm.weight"); kvn = G(f"{a}.kv_norm.weight") - woa = G(f"{a}.o_a_proj.weight") - fnorm = G(f"{p}.post_attention_layernorm.weight") - - qa_w = G(f"{a}.q_a_proj.weight"); qa_sf = G(f"{a}.q_a_proj.weight_scale"); qa_gs = G(f"{a}.q_a_proj.weight_scale_2") - qb_w = G(f"{a}.q_b_proj.weight"); qb_sf = G(f"{a}.q_b_proj.weight_scale"); qb_gs = G(f"{a}.q_b_proj.weight_scale_2") - kv_w = G(f"{a}.kv_proj.weight"); kv_sf = G(f"{a}.kv_proj.weight_scale"); kv_gs = G(f"{a}.kv_proj.weight_scale_2") - wob_w = G(f"{a}.o_b_proj.weight"); wob_sf = G(f"{a}.o_b_proj.weight_scale"); wob_gs = G(f"{a}.o_b_proj.weight_scale_2") - - r_qa = make_runner(qa_w, qa_sf, qa_gs, H, qa_w.shape[0]) - r_qb = make_runner(qb_w, qb_sf, qb_gs, QL, qb_w.shape[0]) - r_kv = make_runner(kv_w, kv_sf, kv_gs, H, kv_w.shape[0]) - r_wob = make_runner(wob_w, wob_sf, wob_gs, OG*OL, wob_w.shape[0]) - cos_sin = build_cos_sin(max_pos=4096).to(DEV) - woa_3d = woa.view(OG, OL, HPG * HD) - - # Shared expert - m = f"{p}.mlp" - se_gate_w = G(f"{m}.shared_experts.gate_proj.weight"); se_gate_sf = G(f"{m}.shared_experts.gate_proj.weight_scale"); se_gate_gs = G(f"{m}.shared_experts.gate_proj.weight_scale_2") - se_up_w = G(f"{m}.shared_experts.up_proj.weight"); se_up_sf = G(f"{m}.shared_experts.up_proj.weight_scale"); se_up_gs = G(f"{m}.shared_experts.up_proj.weight_scale_2") - se_down_w = G(f"{m}.shared_experts.down_proj.weight"); se_down_sf = G(f"{m}.shared_experts.down_proj.weight_scale"); se_down_gs = G(f"{m}.shared_experts.down_proj.weight_scale_2") - - r_se_gate = make_runner(se_gate_w, se_gate_sf, se_gate_gs, H, se_gate_w.shape[0]) - r_se_up = make_runner(se_up_w, se_up_sf, se_up_gs, H, se_up_w.shape[0]) - r_se_down = make_runner(se_down_w, se_down_sf, se_down_gs, INTERMEDIATE, se_down_w.shape[0]) - - positions = torch.arange(num_tokens, dtype=torch.int64, device=DEV) - - with torch.no_grad(): - # Attention - normed = rms(hidden, anorm, EPS) - qa = r_qa.run(normed); kv = r_kv.run(normed) - qa_n = rms(qa, qn, EPS); kv_n = rms(kv, kvn, EPS) - q = r_qb.run(qa_n).view(num_tokens, NH, HD) - q_rope = apply_gptj_rope(q, positions, cos_sin, NOPE, ROPE) - kv_rope = apply_gptj_rope(kv_n.unsqueeze(1), positions, cos_sin, NOPE, ROPE).squeeze(1) - - o_attn = causal_prefill_attention(q_rope, kv_rope, SCALE) - o_inv = apply_inv_gptj_rope(o_attn, positions, cos_sin, NOPE, ROPE) - o_grouped = o_inv.reshape(num_tokens, OG, HPG * HD).permute(1, 0, 2) - z = torch.bmm(o_grouped, woa_3d.transpose(1, 2)).permute(1, 0, 2).reshape(num_tokens, OG * OL) - attn_out = r_wob.run(z) - hidden = hidden + attn_out - - # Shared expert MoE - fnormed = rms(hidden, fnorm, EPS) - gate_out = r_se_gate.run(fnormed) - up_out = r_se_up.run(fnormed) - activated = F.silu(gate_out) * up_out - se_out = r_se_down.run(activated) - hidden = hidden + se_out - - attn_nan = torch.isnan(attn_out).any().item() - moe_nan = torch.isnan(se_out).any().item() - hs_nan = torch.isnan(hidden).any().item() - print(f" Layer {layer_id} ({lt}): attn_nan={attn_nan} moe_nan={moe_nan} hidden_nan={hs_nan} amax={hidden.amax():.4f}") - - if hs_nan: - print(f" NaN detected at layer {layer_id}! Stopping.") - break - - del r_qa, r_qb, r_kv, r_wob, r_se_gate, r_se_up, r_se_down - torch.cuda.empty_cache() - _cache.clear() - - -def main(): - print("=" * 70) - print(" DeepSeek-V4 Full Layer NaN Test") - print(" Tests attention + MoE to find where NaN originates") - print("=" * 70) - - print("\n=== Test 1: Single full layer (attention + MoE) ===") - test_full_layer(layer_id=2, num_tokens=8, num_moe_experts=16) - - print("\n=== Test 2: Multi-layer chain (attention + shared expert only) ===") - test_multi_layer() - - print(f"\n{'='*70}") - - -if __name__ == "__main__": - main() diff --git a/tests/archive/test_full_model_b200.py b/tests/archive/test_full_model_b200.py deleted file mode 100644 index be4141f5..00000000 --- a/tests/archive/test_full_model_b200.py +++ /dev/null @@ -1,314 +0,0 @@ -#!/usr/bin/env python3 -""" -Full DeepSeek-V4 Model Forward Test - -Runs the ENTIRE model through our kernel pipeline: -- 61 layers: C128A, C4A, SWA attention + MoE -- All projections: CuTeDSL NVFP4 -- Attention: BF16 (SDPA for SWA, sparse for CSA/HCA) -- KV cache: FP8 quantize/dequant -- MoE: CuTeDSL NVFP4 -- LM head: BF16 - -Outputs logits and checks they're reasonable (not garbage). - -Usage (on B200): - cd /root/nvfp4-megamoe-kernel - PYTHONPATH=/root/nvfp4-megamoe-kernel tests/venv/bin/python tests/test_full_model_b200.py -""" - -import sys, os, json, torch, torch.nn.functional as F, math, time -from safetensors import safe_open - -REPO = "/root/nvfp4-megamoe-kernel" -sys.path.insert(0, REPO) -MODEL = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4" -DEV = "cuda:0" - -# Config -H = 7168; NH = 128; HD = 512; NOPE = 448; ROPE = 64 -QL = 1536; OL = 1024; OG = 16; HPG = NH // OG -EPS = 1e-6; WINDOW = 128; SCALE = HD ** -0.5 -NUM_LAYERS = 61 -NUM_EXPERTS = 384; TOPK = 6 - -E2M1 = torch.tensor([0,.5,1.,1.5,2.,3.,4.,6.,-0,-.5,-1.,-1.5,-2.,-3.,-4.,-6.], dtype=torch.float32) - -_cache = {} -def P(k, wm, md): - if k in _cache: return _cache[k] - with safe_open(os.path.join(md, wm[k]), framework="pt") as f: - t = f.get_tensor(k) - _cache[k] = t - return t - -def dequant(w, sf, gs): - d = w.device; lut = E2M1.to(d) - lo = lut[(w & 0xF).long()]; hi = lut[((w >> 4) & 0xF).long()] - O, I2 = w.shape; I = I2*2 - u = torch.empty(O, I, dtype=torch.float32, device=d) - u[:,0::2] = lo; u[:,1::2] = hi - bs = sf.float().repeat_interleave(16, dim=1)[:O,:I] - return (u * bs * gs).to(torch.bfloat16) - -def rms(x, w, eps=1e-6): - v = x.float().pow(2).mean(-1, keepdim=True) - return (w.float() * (x * torch.rsqrt(v+eps)).float()).to(x.dtype) - -def make_runner(w, sf, gs_t, inf, outf, fused=False, lw=None): - from dsv4.layers.linear import Nvfp4Linear - fp4 = w.view(torch.float4_e2m1fn_x2).permute(1,0).contiguous() - s = sf.to(torch.float8_e4m3fn) if sf.dtype != torch.float8_e4m3fn else sf - s = s.permute(1,0).contiguous() - if fused and gs_t.numel() == 2: - g1,g2 = gs_t[0].item(), gs_t[1].item(); gs = max(g1,g2) - if g1 != g2: - s32 = s.float(); sp = lw[0] if lw else outf//2 - s32[:sp] *= g1/gs; s32[sp:] *= g2/gs; s = s32.to(torch.float8_e4m3fn) - else: - gs = gs_t.max().item() if gs_t.numel() > 1 else gs_t.item() - r = Nvfp4Linear(in_features=inf, out_features=outf, max_num_tokens=8192, device=str(w.device)) - r.fp4 = [fp4]; r.sf = [s]; r.gs = [gs] - r.finalize_weights(); r._ensure_initialized() - return r - -def build_cos_sin(max_pos=8192, rope_dim=ROPE): - half = rope_dim // 2 - inv_freq = 1.0 / (10000.0 ** (torch.arange(0, half, dtype=torch.float32) / half)) - freqs = torch.outer(torch.arange(max_pos, dtype=torch.float32), inv_freq) - return torch.cat([freqs.cos(), freqs.sin()], dim=-1) - -def apply_gptj_rope(x, positions, cos_sin, nope, rope): - if rope == 0 or x.numel() == 0: return x - half = rope // 2 - cos = cos_sin[positions, :half].to(x.dtype) - sin = cos_sin[positions, half:2*half].to(x.dtype) - if x.dim() == 3: cos = cos.unsqueeze(1); sin = sin.unsqueeze(1) - x_rope = x[..., nope:].clone() - even = x_rope[..., 0::2]; odd = x_rope[..., 1::2] - out = x.clone() - out[..., nope:][..., 0::2] = even * cos - odd * sin - out[..., nope:][..., 1::2] = even * sin + odd * cos - return out - -def apply_inv_gptj_rope(x, positions, cos_sin, nope, rope): - if rope == 0 or x.numel() == 0: return x - half = rope // 2 - cos = cos_sin[positions, :half].to(x.dtype) - sin = cos_sin[positions, half:2*half].to(x.dtype) - if x.dim() == 3: cos = cos.unsqueeze(1); sin = sin.unsqueeze(1) - x_rope = x[..., nope:].clone() - even = x_rope[..., 0::2]; odd = x_rope[..., 1::2] - out = x.clone() - out[..., nope:][..., 0::2] = even * cos + odd * sin - out[..., nope:][..., 1::2] = -even * sin + odd * cos - return out - - -def bf16_causal_attention(q, kv, scale): - """Full causal self-attention.""" - T, NH, HD = q.shape - q_2d = q.reshape(T * NH, HD) - kv_exp = kv.unsqueeze(1).expand(-1, NH, -1).contiguous() - k_2d = kv_exp.permute(1, 0, 2).unsqueeze(1).expand(NH, T, T, -1).contiguous().reshape(T * NH, T, HD) - v_2d = k_2d.clone() - scores = torch.matmul(q_2d.unsqueeze(1), k_2d.transpose(-1, -2)) * scale - qpos = torch.arange(T, device=q.device).unsqueeze(1).repeat(1, NH).reshape(T * NH) - kpos = torch.arange(T, device=q.device).unsqueeze(0) - causal = kpos <= qpos.unsqueeze(1) - scores = scores.squeeze(1).masked_fill(~causal, float('-inf')) - weights = F.softmax(scores.float(), dim=-1).to(q.dtype) - out = torch.matmul(weights.unsqueeze(1), v_2d).squeeze(1) - return out.reshape(T, NH, HD) - - -def make_moe_runner(layer_id, wm, model_path): - """Create CuTeDSL MoE runner for a layer.""" - from dsv4.layers.moe import Nvfp4MoE - - p = f"model.layers.{layer_id}.mlp" - G = lambda k: P(k, wm, model_path).to(DEV) - - # Gate (router) weight - gate_w = G(f"{p}.gate.weight") # (384, 7168) BF16 - - # Expert weights (NVFP4) - w13_w = G(f"{p}.experts.w13_weight") # (384, 6144, 3584) uint8 - w13_sf = G(f"{p}.experts.w13_weight_scale") # (384, 6144, 448) fp8 - w13_gs = G(f"{p}.experts.w13_weight_scale_2") # (384, 2) - w2_w = G(f"{p}.experts.w2_weight") - w2_sf = G(f"{p}.experts.w2_weight_scale") - w2_gs = G(f"{p}.experts.w2_weight_scale_2") - - # Convert to runner format - l1_fp4 = w13_w.view(torch.float4_e2m1fn_x2).permute(1,0).contiguous() - l2_fp4 = w2_w.view(torch.float4_e2m1fn_x2).permute(1,0).contiguous() - l1_sf = w13_sf.to(torch.float8_e4m3fn).permute(1,0).contiguous() if w13_sf.dtype != torch.float8_e4m3fn else w13_sf.permute(1,0).contiguous() - l2_sf = w2_sf.to(torch.float8_e4m3fn).permute(1,0).contiguous() if w2_sf.dtype != torch.float8_e4m3fn else w2_sf.permute(1,0).contiguous() - - intermediate_size = 3072 # per expert - runner = Nvfp4MoE( - num_experts=NUM_EXPERTS, - hidden_size=H, - intermediate_size=intermediate_size, - max_num_tokens=8192, - top_k=TOPK, - device=DEV, - ) - - l1_gs_list = w13_gs.tolist() - l2_gs_list = w2_gs.tolist() - - runner.prepare_weights_from_stacked(l1_fp4, l1_sf, l1_gs_list, l2_fp4, l2_sf, l2_gs_list) - - # Shared expert - se_w13_w = G(f"{p}.shared_experts.gate_up_proj.weight") - se_w13_sf = G(f"{p}.shared_experts.gate_up_proj.weight_scale") - se_w13_gs = G(f"{p}.shared_experts.gate_up_proj.weight_scale_2") - se_w2_w = G(f"{p}.shared_experts.down_proj.weight") - se_w2_sf = G(f"{p}.shared_experts.down_proj.weight_scale") - se_w2_gs = G(f"{p}.shared_experts.down_proj.weight_scale_2") - - se_r_gate_up = make_runner(se_w13_w, se_w13_sf, se_w13_gs, H, se_w13_w.shape[0], fused=True, lw=[intermediate_size]) - se_r_down = make_runner(se_w2_w, se_w2_sf, se_w2_gs, intermediate_size, se_w2_w.shape[0]) - - return runner, gate_w, se_r_gate_up, se_r_down - - -def main(): - torch.cuda.set_device(0) - torch.manual_seed(42) - - print("=" * 70) - print(" Full DeepSeek-V4 Model Forward Test") - print(" 61 layers, all CuTeDSL NVFP4 kernels") - print("=" * 70) - - with open(os.path.join(MODEL, "model.safetensors.index.json")) as f: - wm = json.load(f)["weight_map"] - G = lambda k: P(k, wm, MODEL).to(DEV) - - # Load compress_ratios - with open(os.path.join(MODEL, "config.json")) as f: - config = json.load(f) - compress_ratios = config["compress_ratios"] - - # Global weights - emb = G("model.embed_tokens.weight") - fnorm_w = G("model.norm.weight") - lm_head = G("lm_head.weight") - cos_sin = build_cos_sin(max_pos=8192).to(DEV) - - # Input - NT = 6 - token_ids = torch.tensor([1, 450, 8403, 315, 5413, 374], dtype=torch.long, device=DEV) - positions = torch.arange(NT, dtype=torch.int64, device=DEV) - - print(f" Input: {NT} tokens: {token_ids.tolist()}") - print(f" Model: {NUM_LAYERS} layers, {NUM_EXPERTS} experts, top-{TOPK}") - - with torch.no_grad(): - hidden = emb[token_ids] - - for layer_id in range(NUM_LAYERS): - cr = max(1, compress_ratios[layer_id]) - layer_type = "SWA" if cr <= 1 else f"C{cr}A" - p = f"model.layers.{layer_id}" - a = f"{p}.self_attn" - m = f"{p}.mlp" - - # Layer norms - anorm = G(f"{p}.input_layernorm.weight") - fnorm = G(f"{p}.post_attention_layernorm.weight") - - # ── Attention ──────────────────────────────────────────── - qn = G(f"{a}.q_a_norm.weight") - kvn = G(f"{a}.kv_norm.weight") - woa = G(f"{a}.o_a_proj.weight") - - r_qa = make_runner(G(f"{a}.q_a_proj.weight"), G(f"{a}.q_a_proj.weight_scale"), G(f"{a}.q_a_proj.weight_scale_2"), H, G(f"{a}.q_a_proj.weight").shape[0]) - r_qb = make_runner(G(f"{a}.q_b_proj.weight"), G(f"{a}.q_b_proj.weight_scale"), G(f"{a}.q_b_proj.weight_scale_2"), QL, G(f"{a}.q_b_proj.weight").shape[0]) - r_kv = make_runner(G(f"{a}.kv_proj.weight"), G(f"{a}.kv_proj.weight_scale"), G(f"{a}.kv_proj.weight_scale_2"), H, G(f"{a}.kv_proj.weight").shape[0]) - r_wob = make_runner(G(f"{a}.o_b_proj.weight"), G(f"{a}.o_b_proj.weight_scale"), G(f"{a}.o_b_proj.weight_scale_2"), OG*OL, G(f"{a}.o_b_proj.weight").shape[0]) - - normed = rms(hidden, anorm, EPS) - - qa = r_qa.run(normed) - kv = r_kv.run(normed) - qa_n = rms(qa, qn, EPS) - kv_n = rms(kv, kvn, EPS) - q = r_qb.run(qa_n).view(NT, NH, HD) - q_rope = apply_gptj_rope(q, positions, cos_sin, NOPE, ROPE) - - # Attention (BF16 causal — simplified, no sparse index yet) - o_attn = bf16_causal_attention(q_rope, kv_n, SCALE) - - # o_a: inverse RoPE + BMM - o_inv = apply_inv_gptj_rope(o_attn, positions, cos_sin, NOPE, ROPE) - o_grouped = o_inv.view(NT, OG, HPG * HD).permute(1, 0, 2) - woa_3d = woa.view(OG, OL, HPG * HD) - z = torch.bmm(o_grouped, woa_3d.transpose(1, 2)).permute(1, 0, 2).reshape(NT, OG * OL) - attn_out = r_wob.run(z) - - # Residual - hidden = hidden + attn_out - - # ── MoE ────────────────────────────────────────────────── - # For speed, only load MoE for first 3 layers + last layer - if layer_id < 3 or layer_id == NUM_LAYERS - 1: - fnormed = rms(hidden, fnorm, EPS) - - # Simplified: just use BF16 shared expert for now - # (full MoE test is in layertest.py) - se_w13_w = G(f"{m}.shared_experts.gate_up_proj.weight") - se_w13_sf = G(f"{m}.shared_experts.gate_up_proj.weight_scale") - se_w13_gs = G(f"{m}.shared_experts.gate_up_proj.weight_scale_2") - se_w2_w = G(f"{m}.shared_experts.down_proj.weight") - se_w2_sf = G(f"{m}.shared_experts.down_proj.weight_scale") - se_w2_gs = G(f"{m}.shared_experts.down_proj.weight_scale_2") - - se_gate_up = make_runner(se_w13_w, se_w13_sf, se_w13_gs, H, se_w13_w.shape[0], fused=True, lw=[3072]) - se_down = make_runner(se_w2_w, se_w2_sf, se_w2_gs, 3072, se_w2_w.shape[0]) - - # Shared expert only (skip routed experts for speed) - se_out = se_gate_up.run(fnormed) - gate, up = se_out[:, :3072], se_out[:, 3072:] - se_activated = F.silu(gate) * up - se_final = se_down.run(se_activated) - - hidden = hidden + se_final - else: - # Skip MoE for middle layers (just use residual) - # This is WRONG for correctness but saves time - fnormed = rms(hidden, fnorm, EPS) - hidden = hidden + fnormed # placeholder - - if layer_id % 10 == 0 or layer_id == NUM_LAYERS - 1: - print(f" Layer {layer_id} ({layer_type}): hidden amax={hidden.amax():.4f} NaN={torch.isnan(hidden).any()}") - - # Cleanup per-layer weights - torch.cuda.empty_cache() - _cache.clear() - - # Final norm + LM head - x_n = rms(hidden, fnorm_w, EPS) - logits = x_n @ lm_head.T - - print(f"\n Final logits: amax={logits.amax():.4f} std={logits[-1].float().std():.4f}") - top5 = torch.topk(logits[-1], 5) - print(f" Top 5 tokens: {top5.indices.tolist()}") - print(f" Top 5 probs: {F.softmax(top5.values.float(), dim=0).tolist()}") - - log_std = logits[-1].float().std().item() - if 0.5 < log_std < 50: - print(f" ✅ Logits look reasonable (std={log_std:.4f})") - else: - print(f" ❌ Logits are garbage (std={log_std:.4f})") - - print(f"\n{'='*70}") - print(f" DONE") - print(f"{'='*70}") - - -if __name__ == "__main__": - main() diff --git a/tests/archive/test_fused_step1.py b/tests/archive/test_fused_step1.py deleted file mode 100644 index 4b83fb26..00000000 --- a/tests/archive/test_fused_step1.py +++ /dev/null @@ -1,92 +0,0 @@ -"""Test: Validate SiLU in registers (Step 1 of fused SwiGLU). - -Compiles the fused kernel with fused_swiglu=True, runs it, and compares -the BF16 output with PyTorch SiLU applied to the standard L1 GEMM output. -""" -import torch -import sys -sys.path.insert(0, '/root/dsv4-nvfp4-workspace/kernel') - -from dsv4.ops.quantize import ( - quantize_weight_to_nvfp4, - quantize_activation_nvfp4, -) -from dsv4.ops.layouts import ( - make_b_k_major, - assemble_scales_2d_side, - assemble_scales_3d_side, -) -from dsv4.ops.gemm_runner import ( - run_nvfp4_grouped_gemm, - run_fused_swiglu_grouped_gemm, - warmup_compilation, -) - - -def test_silu_step1(): - device = "cuda" - num_experts = 4 - hidden = 512 - intermediate = 256 - num_tokens = 32 - - torch.manual_seed(42) - x = torch.randn(num_tokens, hidden, dtype=torch.bfloat16, device=device) - l1_w = torch.randn(num_experts, 2 * intermediate, hidden, dtype=torch.bfloat16, device=device) - - l1_fp4_list, l1_sf_list, l1_gs_list = [], [], [] - for e in range(num_experts): - w_fp4, w_sf, w_gs = quantize_weight_to_nvfp4(l1_w[e].T) - l1_fp4_list.append(w_fp4) - l1_sf_list.append(w_sf) - l1_gs_list.append(w_gs) - - l1_mat_b = make_b_k_major(torch.stack(l1_fp4_list)) - l1_scale_b = assemble_scales_3d_side(l1_sf_list) - l1_gs = torch.tensor(l1_gs_list, dtype=torch.float32, device=device) - - gs_val = x.abs().max().item() / (6.0 * 448.0) - x_fp4, x_sf = quantize_activation_nvfp4(x, gs_val) - tokens_per_expert = [num_tokens // num_experts] * num_experts - scale_a = assemble_scales_2d_side([x_sf[i*tpe:(i+1)*tpe] for i, tpe in enumerate(tokens_per_expert)]) - expert_offsets = torch.tensor( - [sum(tokens_per_expert[:e+1]) for e in range(num_experts)], - dtype=torch.int32, device=device, - ) - global_scale_a = torch.full((num_experts,), gs_val, dtype=torch.float32, device=device) - - warmup_compilation(num_experts, hidden // 2, (2 * intermediate) // 2, device) - - # 1. Standard L1 GEMM (no SiLU) - out_bf16 = run_nvfp4_grouped_gemm( - mat_a=x_fp4, mat_b=l1_mat_b, - scale_a=scale_a, scale_b=l1_scale_b, - expert_offsets=expert_offsets, - global_scale_a=global_scale_a, global_scale_b=l1_gs, - ) - silu_ref = torch.nn.functional.silu(out_bf16) - print(f"Standard L1 output: shape={out_bf16.shape}, amax={out_bf16.abs().amax().item():.4f}") - print(f"PyTorch SiLU ref: amax={silu_ref.abs().amax().item():.4f}") - - # 2. Fused kernel with SiLU in registers - print("\nCompiling fused kernel (first time, may take a while)...") - out_fused = run_fused_swiglu_grouped_gemm( - mat_a=x_fp4, mat_b=l1_mat_b, - scale_a=scale_a, scale_b=l1_scale_b, - expert_offsets=expert_offsets, - global_scale_a=global_scale_a, global_scale_b=l1_gs, - ) - print(f"Fused SiLU output: shape={out_fused.shape}, amax={out_fused.abs().amax().item():.4f}") - - # 3. Compare - diff = (out_fused - silu_ref).float() - rel_err = diff.norm() / silu_ref.float().norm() - max_err = diff.abs().max() - print(f"\n=== Results ===") - print(f"Relative error: {rel_err.item():.6f}") - print(f"Max abs error: {max_err.item():.6f}") - print(f"PASS" if rel_err.item() < 0.1 else "FAIL (tolerance: 0.1 for NVFP4 quant noise)") - - -if __name__ == "__main__": - test_silu_step1() diff --git a/tests/archive/test_inspect_types.py b/tests/archive/test_inspect_types.py deleted file mode 100644 index 7ef47b12..00000000 --- a/tests/archive/test_inspect_types.py +++ /dev/null @@ -1,73 +0,0 @@ -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils -from cutlass.cute.nvgpu import tcgen05 -from cutlass import Float32, BFloat16 -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - -m, n, head_dim = 128, 128, 64 -q = torch.randn(m, head_dim, 1, dtype=torch.bfloat16, device='cuda') -k = torch.randn(n, head_dim, 1, dtype=torch.bfloat16, device='cuda') -v_base = torch.randn(head_dim, n, dtype=torch.bfloat16, device='cuda') -v = v_base.as_strided((head_dim, n), (1, head_dim)).unsqueeze(-1) -c = torch.zeros(m, head_dim, 1, dtype=torch.bfloat16, device='cuda') - -mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) -mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) -mV = ct.from_dlpack(v).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v)) -mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) -stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - -class InspectKernel: - def __init__(self): - self.q_dtype = BFloat16; self.acc_dtype = Float32 - - @cute.jit - def __call__(self, q, k, v, c, stream): - a_major = LayoutEnum.from_tensor(q).mma_major_mode() - b_major = LayoutEnum.from_tensor(k).mma_major_mode() - v_major = LayoutEnum.from_tensor(v).mma_major_mode() - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, a_major, b_major, - self.acc_dtype, tcgen05.CtaGroup.ONE, (128, 128), tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, v_major, - self.acc_dtype, tcgen05.CtaGroup.ONE, (128, 64), tcgen05.OperandSource.TMEM) - - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - qk_mma_tiler = (128, 128, qk_inst_k * 4) - pv_mma_tiler = (qk_mma_tiler[0], qk_mma_tiler[2], qk_mma_tiler[1]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - - qk_acc_shape = qk_thr.partition_shape_C(qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - - p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, pv_mma_tiler, self.q_dtype, 1) - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - - tmem_p0_offset = 32 - p_offset = self.acc_dtype.width // self.q_dtype.width * tmem_p0_offset - - # Print pointer values - if tOrP inherits FP32 type, +64 adds 256 bytes - # If BF16 type, +64 adds 128 bytes (correct, matches tStS+32 FP32 = 128 bytes) - cute.printf("tStS ptr value: %d", tStS.iterator) - cute.printf("tStS_P ptr (tStS+32): %d", tStS.iterator + tmem_p0_offset) - cute.printf("tOrP ptr value: %d", tOrP.iterator) - cute.printf("tOrP0 ptr (tOrP+64): %d", tOrP.iterator + p_offset) - cute.printf("p_offset: %d", p_offset) - - pv_acc_shape = pv_thr.partition_shape_C(pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - cute.printf("o_cols: %d", o_cols) - -kernel = InspectKernel() -compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) -compiled(mQ, mK, mV, mC, stream) -torch.cuda.synchronize() diff --git a/tests/archive/test_interleave.py b/tests/archive/test_interleave.py deleted file mode 100644 index a217f75f..00000000 --- a/tests/archive/test_interleave.py +++ /dev/null @@ -1,144 +0,0 @@ -"""Test: Verify weight interleave produces correct gate/up pairs in GEMM output. - -Stage 1 validation: If interleaved weights produce the same GEMM result -as non-interleaved weights (after de-interleaving the output), the -interleave is correct and the fused epilogue can safely assume gate/up -pairs are adjacent in registers. -""" -import torch -import sys -sys.path.insert(0 = '/root/dsv4-nvfp4-workspace/kernel') # FIXME - -from dsv4.ops.quantize import ( - quantize_to_nvfp4, - quantize_activation_nvfp4, - quantize_weight_to_nvfp4, -) -from dsv4.ops.layouts import ( - interleave_l1_weights, - deinterleave_l1_weights, - make_b_k_major, - assemble_scales_2d_side, - assemble_scales_3d_side, -) -from dsv4.ops.gemm_runner import ( - run_nvfp4_grouped_gemm, -) - - -def test_interleave_correctness(): - """Verify that interleaving weights and de-interleaving the GEMM output - gives the same result as non-interleaved weights. - """ - device = "cuda" - num_experts = 4 - hidden = 512 - intermediate = 256 - num_tokens = 32 - - # Create random BF16 input - x = torch.randn(num_tokens, hidden, dtype=torch.bfloat16, device=device) - - # Create random BF16 weights for gate and up - gate_w = torch.randn(num_experts, intermediate, hidden, dtype=torch.bfloat16, device=device) - up_w = torch.randn(num_experts, intermediate, hidden, dtype=torch.bfloat16, device=device) - - # === Path A: Non-interleaved (current production path) === - # Fuse gate+up: (E, 2*intermediate, hidden) - l1_bf16 = torch.cat([gate_w, up_w], dim=1) # (E, 6144, 7168) → (E, 2*inter, hidden) - - # Quantize weights - l1_fp4_list, l1_sf_list, l1_gs_list = [], [], [] - for e in range(num_experts): - w_fp4, w_sf, w_gs = quantize_weight_to_nvfp4(l1_bf16[e].T) # (K, N) - l1_fp4_list.append(w_fp4) - l1_sf_list.append(w_sf) - l1_gs_list.append(w_gs) - - # Stack and convert - l1_mat_b = make_b_k_major(torch.stack(l1_fp4_list)) - l1_scale_b = assemble_scales_3d_side(l1_sf_list) - l1_gs = torch.tensor(l1_gs_list, dtype=torch.float32, device=device) - - # Quantize activation - gs_val = x.abs().max().item() / (6.0 * 448.0) - x_fp4, x_sf = quantize_activation_nvfp4(x, gs_val) - - # Assemble scales - tokens_per_expert = [num_tokens // num_experts] * num_experts - scale_a = assemble_scales_2d_side([x_sf[i*tpe:(i+1)*tpe] for i, tpe in enumerate(tokens_per_expert)]) - - expert_offsets = torch.tensor( - [sum(tokens_per_expert[:e+1]) for e in range(num_experts)], - dtype=torch.int32, device=device, - ) - global_scale_a = torch.full((num_experts,), gs_val, dtype=torch.float32, device=device) - - # Run GEMM - out_a = run_nvfp4_grouped_gemm( - mat_a=x_fp4, mat_b=l1_mat_b, - scale_a=scale_a, scale_b=l1_scale_b, - expert_offsets=expert_offsets, - global_scale_a=global_scale_a, global_scale_b=l1_gs, - ) - # out_a: (num_tokens, 2*intermediate) BF16 - # gate = out_a[:, :intermediate], up = out_a[:, intermediate:] - gate_a = out_a[:, :intermediate] - up_a = out_a[:, intermediate:] - result_a = torch.nn.functional.silu(gate_a) * up_a # SwiGLU result - - # === Path B: Interleaved weights === - # Quantize gate and up separately, then interleave - gate_fp4, gate_sf, gate_gs = [], [], [] - up_fp4, up_sf, up_gs = [], [], [] - for e in range(num_experts): - g4, gs4, gg4 = quantize_weight_to_nvfp4(gate_w[e].T) - u4, us4, ug4 = quantize_weight_to_nvfp4(up_w[e].T) - gate_fp4.append(g4) - gate_sf.append(gs4) - gate_gs.append(gg4) - up_fp4.append(u4) - up_sf.append(us4) - up_gs.append(ug4) - - # Fuse and interleave - gate_stacked = torch.stack(gate_fp4) # (E, K_packed, N/2) - up_stacked = torch.stack(up_fp4) # (E, K_packed, N/2) - l1_bf16_fp4 = torch.cat([gate_stacked, up_stacked], dim=2) # (E, K, N) non-interleaved - l1_interleaved = interleave_l1_weights(l1_bf16_fp4) # interleaved - - # Make K-major - l1_mat_b_int = make_b_k_major(l1_interleaved) - - # Scale assembly: gate and up scales combined - l1_scale_b_int = assemble_scales_3d_side(gate_sf + up_sf) # interleave scales too? - # Actually, the scale interleaving needs to match the weight interleaving. - # This is more complex. For Stage 1, let's use a simpler approach. - - # Actually, for the interleaved path to produce the same GEMM output, - # we need the SFB to also be interleaved to match. - # The GEMM is: A (M, K) x B (E, K, N) = C (M, N) - # If we permute the N dimension of B, we permute the N dimension of C. - # So the output columns are also interleaved. - - # For this test, we just verify that the interleaved GEMM output, - # when de-interleaved, matches the non-interleaved output. - - # But the SFB (scale_b) must match the interleaved B. - # The B tensor has its N columns interleaved, so the SFB must be - # interleaved in the same way. - - # SFB for interleaved B: we need to interleave the scales too. - # Since scales are per-(K_sf, N) and we're interleaving N at granularity 4 FP4 cols, - # the scales need to be interleaved at the same granularity. - - # This is getting complex. Let me simplify: just test the interleave - # function itself, not the full GEMM. - - print("Interleave/deinterleave round-trip: PASSED (tested in bridge.py)") - print("Full GEMM interleave test: SKIPPED (requires SFB interleaving)") - print("Stage 1 kernel test will validate the full pipeline") - - -if __name__ == "__main__": - test_interleave_correctness() diff --git a/tests/archive/test_interleave_gemm.py b/tests/archive/test_interleave_gemm.py deleted file mode 100644 index d131f73a..00000000 --- a/tests/archive/test_interleave_gemm.py +++ /dev/null @@ -1,137 +0,0 @@ -"""Test: Verify that interleaved L1 weights produce the same GEMM result. - -The key insight: we quantize gate+up TOGETHER (same as non-interleaved), -then interleave the ALREADY-QUANTIZED FP4 bytes and scales in the N dimension. -This preserves quantization fidelity. -""" -import torch -import sys -sys.path.insert(0, '/root/dsv4-nvfp4-workspace/kernel') - -from dsv4.ops.quantize import ( - quantize_weight_to_nvfp4, - quantize_activation_nvfp4, -) -from dsv4.ops.layouts import ( - interleave_l1_weights, - make_b_k_major, - assemble_scales_2d_side, - assemble_scales_3d_side, -) -from dsv4.ops.gemm_runner import ( - run_nvfp4_grouped_gemm, - warmup_compilation, -) - - -def interleave_sfb(raw_scales, granularity_bf16=8): - """Interleave gate/up scales at the same granularity as the FP4 weights. - - raw_scales: list of (K_sf, N) float8_e4m3fn tensors where N = 2*intermediate_sf - Returns: list of (K_sf, N) float8_e4m3fn with interleaved gate/up - """ - g = granularity_bf16 // 2 # 4 FP4 scale columns per group - result = [] - for sf in raw_scales: - K_sf, N = sf.shape - N_half = N // 2 - gate = sf[:, :N_half].reshape(K_sf, N_half // g, g) - up = sf[:, N_half:].reshape(K_sf, N_half // g, g) - interleaved = torch.stack([gate, up], dim=2).reshape(K_sf, N) - result.append(interleaved) - return result - - -def test_interleave_gemm(): - device = "cuda" - num_experts = 4 - hidden = 512 - intermediate = 256 - num_tokens = 32 - - torch.manual_seed(42) - x = torch.randn(num_tokens, hidden, dtype=torch.bfloat16, device=device) - gate_w = torch.randn(num_experts, intermediate, hidden, dtype=torch.bfloat16, device=device) - up_w = torch.randn(num_experts, intermediate, hidden, dtype=torch.bfloat16, device=device) - - # === Path A: Non-interleaved === - l1_bf16 = torch.cat([gate_w, up_w], dim=1) # (E, 2*inter, hidden) - l1_fp4_list, l1_sf_list, l1_gs_list = [], [], [] - for e in range(num_experts): - w_fp4, w_sf, w_gs = quantize_weight_to_nvfp4(l1_bf16[e].T) - l1_fp4_list.append(w_fp4) - l1_sf_list.append(w_sf) - l1_gs_list.append(w_gs) - - l1_mat_b = make_b_k_major(torch.stack(l1_fp4_list)) - l1_scale_b = assemble_scales_3d_side(l1_sf_list) - l1_gs = torch.tensor(l1_gs_list, dtype=torch.float32, device=device) - - gs_val = x.abs().max().item() / (6.0 * 448.0) - x_fp4, x_sf = quantize_activation_nvfp4(x, gs_val) - tokens_per_expert = [num_tokens // num_experts] * num_experts - scale_a = assemble_scales_2d_side([x_sf[i*tpe:(i+1)*tpe] for i, tpe in enumerate(tokens_per_expert)]) - expert_offsets = torch.tensor( - [sum(tokens_per_expert[:e+1]) for e in range(num_experts)], - dtype=torch.int32, device=device, - ) - global_scale_a = torch.full((num_experts,), gs_val, dtype=torch.float32, device=device) - - warmup_compilation(num_experts, hidden // 2, (2 * intermediate) // 2, device) - out_a = run_nvfp4_grouped_gemm( - mat_a=x_fp4, mat_b=l1_mat_b, - scale_a=scale_a, scale_b=l1_scale_b, - expert_offsets=expert_offsets, - global_scale_a=global_scale_a, global_scale_b=l1_gs, - ) - - # === Path B: Interleaved (quantize together, interleave after) === - # Use the SAME quantized weights, just interleave the N dimension - l1_stacked = torch.stack(l1_fp4_list) # (E, K, N) - l1_interleaved = interleave_l1_weights(l1_stacked) - l1_mat_b_int = make_b_k_major(l1_interleaved) - - # Interleave scales to match - l1_sf_interleaved = interleave_sfb(l1_sf_list) - l1_scale_b_int = assemble_scales_3d_side(l1_sf_interleaved) - # Global scales are the same (quantized together) - - out_b = run_nvfp4_grouped_gemm( - mat_a=x_fp4, mat_b=l1_mat_b_int, - scale_a=scale_a, scale_b=l1_scale_b_int, - expert_offsets=expert_offsets, - global_scale_a=global_scale_a, global_scale_b=l1_gs, - ) - - # De-interleave out_b BF16 to match out_a layout - N = out_b.shape[1] - N_half = N // 2 - g = 8 # granularity in BF16 - out_b_reshaped = out_b.reshape(num_tokens, N // (2 * g), 2, g) - gate_b = out_b_reshaped[:, :, 0, :].reshape(num_tokens, N_half) - up_b = out_b_reshaped[:, :, 1, :].reshape(num_tokens, N_half) - out_b_deint = torch.cat([gate_b, up_b], dim=1) - - diff = (out_a - out_b_deint).float() - rel_err = diff.norm() / out_a.float().norm() - max_err = diff.abs().max() - - print(f"Non-interleaved vs interleaved+deinterleaved:") - print(f" Relative error: {rel_err.item():.6f}") - print(f" Max abs error: {max_err.item():.6f}") - print(f" PASS" if rel_err.item() < 0.01 else " FAIL") - - # Apply SiLU and compare - gate_a = out_a[:, :intermediate] - up_a = out_a[:, intermediate:] - result_a = torch.nn.functional.silu(gate_a) * up_a - result_b = torch.nn.functional.silu(gate_b) * up_b - - diff2 = (result_a - result_b).float() - rel_err2 = diff2.norm() / result_a.float().norm() - print(f" SiLU result error: {rel_err2.item():.6f}") - print(f" SiLU PASS" if rel_err2.item() < 0.01 else " SiLU FAIL") - - -if __name__ == "__main__": - test_interleave_gemm() diff --git a/tests/archive/test_inv_rope.py b/tests/archive/test_inv_rope.py deleted file mode 100644 index 212a8dae..00000000 --- a/tests/archive/test_inv_rope.py +++ /dev/null @@ -1,126 +0,0 @@ -"""Test _apply_inv_rope_bf16: inverse RoPE should undo forward RoPE.""" -import torch -import math - -def apply_rope_bf16(x, positions, cos_sin_cache, nope_dim, rope_dim): - """Forward GPT-J style RoPE.""" - if rope_dim == 0 or x.numel() == 0: - return x - half_rot = rope_dim // 2 - x_f32 = x.to(torch.float32) - cache = cos_sin_cache.index_select(0, positions.to(torch.long)) - cos = cache[:, :half_rot].to(torch.float32) - sin = cache[:, half_rot:2*half_rot].to(torch.float32) - view_shape = (positions.shape[0], 1, half_rot) - cos = cos.view(view_shape) - sin = sin.view(view_shape) - rope = x_f32[..., nope_dim:] - y_even = rope[..., 0::2] - y_odd = rope[..., 1::2] - rope_out = torch.stack( - (y_even * cos - y_odd * sin, y_odd * cos + y_even * sin), - dim=-1, - ).flatten(-2) - x_f32 = x_f32.clone() - x_f32[..., nope_dim:] = rope_out - return x_f32.to(x.dtype) - -def apply_inv_rope_bf16(o, positions, cos_sin_cache, nope_dim, rope_dim): - """Inverse GPT-J style RoPE (sin -> -sin).""" - if rope_dim == 0 or o.numel() == 0: - return o - half_rot = rope_dim // 2 - o_f32 = o.to(torch.float32) - cache = cos_sin_cache.index_select(0, positions.to(torch.long)) - cos = cache[:, :half_rot].to(torch.float32) - sin = cache[:, half_rot:2*half_rot].to(torch.float32) - view_shape = (positions.shape[0], 1, half_rot) - cos = cos.view(view_shape) - sin = sin.view(view_shape) - rope = o_f32[..., nope_dim:] - y_even = rope[..., 0::2] - y_odd = rope[..., 1::2] - rope_out = torch.stack( - (y_even * cos + y_odd * sin, y_odd * cos - y_even * sin), - dim=-1, - ).flatten(-2) - o_f32 = o_f32.clone() - o_f32[..., nope_dim:] = rope_out - return o_f32.to(o.dtype) - - -def test_inv_rope_roundtrip(): - """Forward RoPE then inverse RoPE should be identity.""" - torch.manual_seed(42) - num_tokens = 8 - num_heads = 16 - head_dim = 512 - nope_dim = 448 - rope_dim = 64 - max_pos = 1024 - - # Build cos/sin cache (like RotaryEmbedding) - half_rot = rope_dim // 2 - inv_freq = 1.0 / (10000.0 ** (torch.arange(0, half_rot, dtype=torch.float32) / half_rot)) - positions = torch.randint(0, max_pos, (num_tokens,)) - freqs = positions.float().unsqueeze(1) * inv_freq.unsqueeze(0) - cos_cache_full = torch.zeros(max_pos, half_rot, dtype=torch.float32) - sin_cache_full = torch.zeros(max_pos, half_rot, dtype=torch.float32) - cos_vals = torch.cos(freqs) - sin_vals = torch.sin(freqs) - for i, p in enumerate(positions): - cos_cache_full[p] = cos_vals[i] - sin_cache_full[p] = sin_vals[i] - cos_sin_cache = torch.cat([cos_cache_full, sin_cache_full], dim=1) - - x = torch.randn(num_tokens, num_heads, head_dim, dtype=torch.bfloat16) - - # Forward RoPE - x_rope = apply_rope_bf16(x, positions, cos_sin_cache, nope_dim, rope_dim) - - # Inverse RoPE - x_recovered = apply_inv_rope_bf16(x_rope, positions, cos_sin_cache, nope_dim, rope_dim) - - # Should be identity (within BF16 precision) - diff = (x.to(torch.float32) - x_recovered.to(torch.float32)).abs().max().item() - print(f"Max abs diff: {diff:.6f}") - assert diff < 0.05, f"Roundtrip error too large: {diff}" - print("PASS: inverse RoPE roundtrip within tolerance") - - -def test_nope_dim_unchanged(): - """NoPE dimensions should be unchanged by inverse RoPE.""" - torch.manual_seed(42) - num_tokens = 4 - num_heads = 4 - head_dim = 128 - nope_dim = 96 - rope_dim = 32 - max_pos = 512 - - half_rot = rope_dim // 2 - inv_freq = 1.0 / (10000.0 ** (torch.arange(0, half_rot, dtype=torch.float32) / half_rot)) - positions = torch.randint(0, max_pos, (num_tokens,)) - freqs = positions.float().unsqueeze(1) * inv_freq.unsqueeze(0) - cos_cache_full = torch.zeros(max_pos, half_rot, dtype=torch.float32) - sin_cache_full = torch.zeros(max_pos, half_rot, dtype=torch.float32) - cos_vals = torch.cos(freqs) - sin_vals = torch.sin(freqs) - for i, p in enumerate(positions): - cos_cache_full[p] = cos_vals[i] - sin_cache_full[p] = sin_vals[i] - cos_sin_cache = torch.cat([cos_cache_full, sin_cache_full], dim=1) - - x = torch.randn(num_tokens, num_heads, head_dim, dtype=torch.bfloat16) - x_inv = apply_inv_rope_bf16(x, positions, cos_sin_cache, nope_dim, rope_dim) - - # NoPE dims should be unchanged - nope_diff = (x[..., :nope_dim].to(torch.float32) - x_inv[..., :nope_dim].to(torch.float32)).abs().max().item() - print(f"NoPE max diff: {nope_diff:.6f}") - assert nope_diff == 0.0, "NoPE dimensions should be unchanged" - print("PASS: NoPE dimensions unchanged") - - -if __name__ == "__main__": - test_inv_rope_roundtrip() - test_nope_dim_unchanged() diff --git a/tests/archive/test_kv_cache_b200.py b/tests/archive/test_kv_cache_b200.py deleted file mode 100644 index c827d4ff..00000000 --- a/tests/archive/test_kv_cache_b200.py +++ /dev/null @@ -1,358 +0,0 @@ -#!/usr/bin/env python3 -""" -DeepSeek-V4 KV Cache Kernel — NVFP4 Compressed Storage - -Architecture: -- SWA cache: (T, HD=512) per token, stored as fp8_e4m3 (512 bytes per token) -- CSA cache (C4A): every 4th token stored, (T//4, HD) fp8 (128 bytes per token) -- HCA cache (C128A): every 128th token stored, (T//128, HD) fp8 (4 bytes per token) - -The KV latent is (1, HD=512) — single KV head. After kv_norm + RoPE, -it's quantized to fp8_e4m3 and stored in the paged KV cache. - -For CSA/HCA layers, the compressor further reduces the cache: -- The indexer finds top-k positions in the compressed cache -- Attention only attends to those positions - -This kernel tests: -1. KV quantization: BF16 → fp8_e4m3 (with per-token scale) -2. KV dequantization: fp8_e4m3 → BF16 -3. RoPE on dequantized KV (applied after dequant) -4. Full attention using the cache -5. Compressed cache (CSA/HCA) storage and retrieval - -Usage (on B200): - cd /root/nvfp4-megamoe-kernel - PYTHONPATH=/root/nvfp4-megamoe-kernel tests/venv/bin/python tests/test_kv_cache_b200.py -""" - -import sys, os, json, torch, torch.nn.functional as F, math -from safetensors import safe_open - -REPO = "/root/nvfp4-megamoe-kernel" -sys.path.insert(0, REPO) -MODEL = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4" -DEV = "cuda:0" - -H = 7168; NH = 128; HD = 512; NOPE = 448; ROPE = 64 -QL = 1536; OL = 1024; OG = 16; HPG = NH // OG -EPS = 1e-6; WINDOW = 128; SCALE = HD ** -0.5 - -E2M1 = torch.tensor([0,.5,1.,1.5,2.,3.,4.,6.,-0,-.5,-1.,-1.5,-2.,-3.,-4.,-6.], dtype=torch.float32) - -_cache = {} -def P(k, wm, md): - if k in _cache: return _cache[k] - with safe_open(os.path.join(md, wm[k]), framework="pt") as f: - t = f.get_tensor(k) - _cache[k] = t - return t - -def dequant(w, sf, gs): - d = w.device; lut = E2M1.to(d) - lo = lut[(w & 0xF).long()]; hi = lut[((w >> 4) & 0xF).long()] - O, I2 = w.shape; I = I2*2 - u = torch.empty(O, I, dtype=torch.float32, device=d) - u[:,0::2] = lo; u[:,1::2] = hi - bs = sf.float().repeat_interleave(16, dim=1)[:O,:I] - return (u * bs * gs).to(torch.bfloat16) - -def rms(x, w, eps=1e-6): - v = x.float().pow(2).mean(-1, keepdim=True) - return (w.float() * (x * torch.rsqrt(v+eps)).float()).to(x.dtype) - -def make_runner(w, sf, gs_t, inf, outf, fused=False, lw=None): - from dsv4.layers.linear import Nvfp4Linear - fp4 = w.view(torch.float4_e2m1fn_x2).permute(1,0).contiguous() - s = sf.to(torch.float8_e4m3fn) if sf.dtype != torch.float8_e4m3fn else sf - s = s.permute(1,0).contiguous() - if fused and gs_t.numel() == 2: - g1,g2 = gs_t[0].item(), gs_t[1].item(); gs = max(g1,g2) - if g1 != g2: - s32 = s.float(); sp = lw[0] if lw else outf//2 - s32[:sp] *= g1/gs; s32[sp:] *= g2/gs; s = s32.to(torch.float8_e4m3fn) - else: - gs = gs_t.max().item() if gs_t.numel() > 1 else gs_t.item() - r = Nvfp4Linear(in_features=inf, out_features=outf, max_num_tokens=8192, device=str(w.device)) - r.fp4 = [fp4]; r.sf = [s]; r.gs = [gs] - r.finalize_weights(); r._ensure_initialized() - return r - -def build_cos_sin(max_pos=4096, rope_dim=ROPE): - half = rope_dim // 2 - inv_freq = 1.0 / (10000.0 ** (torch.arange(0, half, dtype=torch.float32) / half)) - freqs = torch.outer(torch.arange(max_pos, dtype=torch.float32), inv_freq) - return torch.cat([freqs.cos(), freqs.sin()], dim=-1) - -def apply_gptj_rope(x, positions, cos_sin, nope, rope): - if rope == 0 or x.numel() == 0: return x - half = rope // 2 - cos = cos_sin[positions, :half].to(x.dtype) - sin = cos_sin[positions, half:2*half].to(x.dtype) - if x.dim() == 3: cos = cos.unsqueeze(1); sin = sin.unsqueeze(1) - x_rope = x[..., nope:].clone() - even = x_rope[..., 0::2]; odd = x_rope[..., 1::2] - out = x.clone() - out[..., nope:][..., 0::2] = even * cos - odd * sin - out[..., nope:][..., 1::2] = even * sin + odd * cos - return out - - -# ── KV Cache Kernels ──────────────────────────────────────────────── - -def kv_quantize_fp8(kv_bf16): - """Quantize KV latent to fp8_e4m3 with per-token scale. - - kv_bf16: (T, HD) BF16 - Returns: (T, HD) fp8, (T, 1) per-token scale (BF16) - """ - # Per-token absmax - amax = kv_bf16.float().abs().amax(dim=-1, keepdim=True).clamp(min=1e-12) - fp8_max = torch.tensor(448.0, dtype=torch.float32, device=kv_bf16.device) # e4m3 max - scale = fp8_max / amax # (T, 1) - kv_scaled = kv_bf16.float() * scale - kv_fp8 = kv_scaled.to(torch.float8_e4m3fn) - # Store inverse scale for dequant - inv_scale = amax / fp8_max # (T, 1) — multiply by this to recover - return kv_fp8, inv_scale.to(torch.bfloat16) - - -def kv_dequantize_fp8(kv_fp8, inv_scale): - """Dequantize fp8 KV back to BF16. - - kv_fp8: (T, HD) fp8_e4m3 - inv_scale: (T, 1) per-token scale - Returns: (T, HD) BF16 - """ - return (kv_fp8.to(torch.bfloat16) * inv_scale).to(torch.bfloat16) - - -def kv_quantize_nvfp4(kv_bf16): - """Quantize KV latent to NVFP4 using CuTeDSL quantize_to_nvfp4. - - More aggressive compression: 2x smaller than fp8 (4 bits vs 8 bits per element). - - kv_bf16: (T, HD) BF16 - Returns: (T, HD//2) fp4, (T, HD//16) sf, scalar gs - """ -from dsv4.ops.quantize import ( - quantize_to_nvfp4, -) - return quantize_to_nvfp4(kv_bf16) - - -def kv_dequantize_nvfp4(kv_fp4, kv_sf, kv_gs, head_dim=HD): - """Dequantize NVFP4 KV back to BF16. - - kv_fp4: (T, HD//2) fp4 (as float4_e2m1fn_x2 viewed as uint8) - kv_sf: (T, HD//16) fp8 block scales - kv_gs: scalar global scale - """ - device = kv_fp4.device - lut = E2M1.to(device) - packed = kv_fp4.view(torch.uint8) - lo = lut[(packed & 0xF).long()] - hi = lut[((packed >> 4) & 0xF).long()] - T = kv_fp4.shape[0] - u = torch.empty(T, head_dim, dtype=torch.float32, device=device) - u[:, 0::2] = lo - u[:, 1::2] = hi - sf_exp = kv_sf.float().repeat_interleave(16, dim=1)[:, :head_dim] - return (u * sf_exp * kv_gs).to(torch.bfloat16) - - -def paged_kv_write(kv_fp8, slot_mapping, cache, block_size): - """Write KV into paged cache. - - kv_fp8: (T, HD) fp8 to write - slot_mapping: (T,) slot indices (position in flat cache) - cache: (num_blocks, block_size, HD) fp8 cache tensor - block_size: tokens per block - """ - for t in range(kv_fp8.shape[0]): - slot = slot_mapping[t].item() - block_idx = slot // block_size - offset = slot % block_size - if block_idx < cache.shape[0]: - cache[block_idx, offset] = kv_fp8[t] - - -def paged_kv_read(slot_mapping, cache, block_size, num_tokens): - """Read KV from paged cache. - - Returns: (num_tokens, HD) fp8 - """ - device = cache.device - HD = cache.shape[-1] - kv = torch.zeros(num_tokens, HD, dtype=cache.dtype, device=device) - for t in range(num_tokens): - slot = slot_mapping[t].item() - block_idx = slot // block_size - offset = slot % block_size - if block_idx < cache.shape[0]: - kv[t] = cache[block_idx, offset] - return kv - - -def main(): - torch.cuda.set_device(0) - torch.manual_seed(42) - - print("=" * 70) - print(" DeepSeek-V4 KV Cache Kernel Test") - print(" fp8 and NVFP4 quantization for paged KV cache") - print("=" * 70) - - # Load real weights - with open(os.path.join(MODEL, "model.safetensors.index.json")) as f: - wm = json.load(f)["weight_map"] - G = lambda k: P(k, wm, MODEL).to(DEV) - - p = "model.layers.0"; a = f"{p}.self_attn" - emb = G("model.embed_tokens.weight") - anorm = G(f"{p}.input_layernorm.weight") - qn = G(f"{a}.q_a_norm.weight"); kvn = G(f"{a}.kv_norm.weight") - qa_w = G(f"{a}.q_a_proj.weight"); qa_sf = G(f"{a}.q_a_proj.weight_scale"); qa_gs = G(f"{a}.q_a_proj.weight_scale_2") - kv_w = G(f"{a}.kv_proj.weight"); kv_sf = G(f"{a}.kv_proj.weight_scale"); kv_gs = G(f"{a}.kv_proj.weight_scale_2") - - r_qa = make_runner(qa_w, qa_sf, qa_gs, H, qa_w.shape[0]) - r_kv = make_runner(kv_w, kv_sf, kv_gs, H, kv_w.shape[0]) - - cos_sin = build_cos_sin(max_pos=4096).to(DEV) - - token_ids = torch.tensor([1, 450, 8403, 315, 5413, 374], dtype=torch.long, device=DEV) - NT = len(token_ids) - positions = torch.arange(NT, dtype=torch.int64, device=DEV) - - with torch.no_grad(): - hidden = emb[token_ids] - normed = rms(hidden, anorm, EPS) - kv_bf16 = r_kv.run(normed) - kv_bf16 = rms(kv_bf16, kvn, EPS) - - # ── Test 1: FP8 KV quantize/dequant roundtrip ──────────────── - print("\n--- Test 1: FP8 KV quantize/dequant ---") - kv_fp8, inv_scale = kv_quantize_fp8(kv_bf16) - kv_recovered = kv_dequantize_fp8(kv_fp8, inv_scale) - c = F.cosine_similarity(kv_bf16.flatten().unsqueeze(0).float(), kv_recovered.flatten().unsqueeze(0).float()).item() - print(f" FP8 roundtrip cosine: {c:.6f} {'✅' if c>=0.99 else '❌'}") - print(f" FP8 cache size: {kv_fp8.numel()} bytes (vs {kv_bf16.numel()*2} BF16)") - - # ── Test 2: NVFP4 KV quantize/dequant roundtrip ────────────── - print("\n--- Test 2: NVFP4 KV quantize/dequant ---") - try: - kv_nfp4, kv_nsf, kv_ngs = kv_quantize_nvfp4(kv_bf16) - kv_n_recovered = kv_dequantize_nvfp4(kv_nfp4, kv_nsf, kv_ngs) - c = F.cosine_similarity(kv_bf16.flatten().unsqueeze(0).float(), kv_n_recovered.flatten().unsqueeze(0).float()).item() - print(f" NVFP4 roundtrip cosine: {c:.6f} {'✅' if c>=0.98 else '❌'}") - print(f" NVFP4 cache size: {kv_nfp4.view(torch.uint8).numel()} bytes (vs {kv_bf16.numel()*2} BF16, {kv_fp8.numel()} FP8)") - except Exception as e: - print(f" NVFP4 quantize failed: {e}") - - # ── Test 3: Paged KV cache write/read with FP8 ─────────────── - print("\n--- Test 3: Paged KV cache (FP8) ---") - block_size = 256 - num_blocks = 64 - cache = torch.zeros(num_blocks, block_size, HD, dtype=torch.float8_e4m3fn, device=DEV) - # Slot mapping: position → flat slot (simplified: slot = position) - slot_mapping = positions # (NT,) - - # Write KV into cache - paged_kv_write(kv_fp8, slot_mapping, cache, block_size) - - # Read back - kv_read = paged_kv_read(slot_mapping, cache, block_size, NT) - c = F.cosine_similarity(kv_fp8.flatten().unsqueeze(0).float(), kv_read.flatten().unsqueeze(0).float()).item() - print(f" Paged read back cosine: {c:.6f} {'✅' if c>=0.999 else '❌'}") - - # ── Test 4: Apply RoPE after dequant ───────────────────────── - print("\n--- Test 4: RoPE on dequantized KV ---") - # KV needs RoPE applied at the positions it was stored at - kv_with_rope = apply_gptj_rope(kv_recovered.unsqueeze(1), positions, cos_sin, NOPE, ROPE).squeeze(1) - print(f" KV+RoPE: amax={kv_with_rope.amax():.4f} NaN={torch.isnan(kv_with_rope).any()}") - - # ── Test 5: Full attention with FP8 KV cache ───────────────── - print("\n--- Test 5: Full attention pipeline with FP8 KV cache ---") - qa_bf16_ref = dequant(qa_w, qa_sf, qa_gs.item()) - qb_bf16_ref = dequant( - G(f"{a}.q_b_proj.weight"), - G(f"{a}.q_b_proj.weight_scale"), - G(f"{a}.q_b_proj.weight_scale_2").item() - ) - kv_bf16_ref = dequant(kv_w, kv_sf, kv_gs.item()) - - r_qb = make_runner( - G(f"{a}.q_b_proj.weight"), - G(f"{a}.q_b_proj.weight_scale"), - G(f"{a}.q_b_proj.weight_scale_2"), - QL, G(f"{a}.q_b_proj.weight").shape[0] - ) - - # Full BF16 reference - qa_ref = normed @ qa_bf16_ref.T - kv_ref = normed @ kv_bf16_ref.T - qa_n_ref = rms(qa_ref, qn, EPS) - kv_n_ref = rms(kv_ref, kvn, EPS) - q_ref = (qa_n_ref @ qb_bf16_ref.T).view(NT, NH, HD) - q_rope_ref = apply_gptj_rope(q_ref, positions, cos_sin, NOPE, ROPE) - - # BF16 causal attention using dequantized FP8 KV cache - kv_from_cache = kv_dequantize_fp8(kv_read, inv_scale) - kv_from_cache_rope = apply_gptj_rope(kv_from_cache.unsqueeze(1), positions, cos_sin, NOPE, ROPE).squeeze(1) - - # Full attention with cached KV - T, NH_t, HD_t = q_rope_ref.shape - q_2d = q_rope_ref.reshape(T * NH_t, HD_t) - kv_exp = kv_from_cache_rope.unsqueeze(1).expand(-1, NH_t, -1).contiguous() - k_2d = kv_exp.permute(1, 0, 2).unsqueeze(1).expand(NH_t, T, T, -1).contiguous().reshape(T * NH_t, T, HD_t) - scores = torch.matmul(q_2d.unsqueeze(1), k_2d.transpose(-1, -2)) * SCALE - qpos = torch.arange(T, device=DEV).unsqueeze(1).repeat(1, NH_t).reshape(T * NH_t) - kpos = torch.arange(T, device=DEV).unsqueeze(0) - causal = kpos <= qpos.unsqueeze(1) - scores = scores.squeeze(1).masked_fill(~causal, float('-inf')) - weights = F.softmax(scores.float(), dim=-1).to(q_rope_ref.dtype) - v_2d = k_2d.clone() - out = torch.matmul(weights.unsqueeze(1), v_2d).squeeze(1).reshape(T, NH_t, HD_t) - - # BF16 attention with original (no cache) KV - kv_exp2 = kv_n_ref.unsqueeze(1).expand(-1, NH_t, -1).contiguous() - k_2d2 = kv_exp2.permute(1, 0, 2).unsqueeze(1).expand(NH_t, T, T, -1).contiguous().reshape(T * NH_t, T, HD_t) - scores2 = torch.matmul(q_2d.unsqueeze(1), k_2d2.transpose(-1, -2)) * SCALE - scores2 = scores2.squeeze(1).masked_fill(~causal, float('-inf')) - weights2 = F.softmax(scores2.float(), dim=-1).to(q_rope_ref.dtype) - out2 = torch.matmul(weights2.unsqueeze(1), v_2d).squeeze(1).reshape(T, NH_t, HD_t) - - c = F.cosine_similarity(out.flatten().unsqueeze(0).float(), out2.flatten().unsqueeze(0).float()).item() - print(f" FP8 cached KV vs BF16 KV attention cosine: {c:.6f} {'✅' if c>=0.98 else '❌'}") - - # ── Test 6: CSA compressed cache (cr=4) ────────────────────── - print("\n--- Test 6: CSA compressed cache (cr=4) ---") - cr = 4 - # Store every 4th token in the compressed cache - compressed_positions = positions[::cr] # every 4th position - compressed_kv = kv_fp8[::cr] # (T//4, HD) fp8 - compressed_inv_scale = inv_scale[::cr] - print(f" Compressed KV shape: {compressed_kv.shape} (from {kv_fp8.shape})") - print(f" Compression ratio: {kv_fp8.shape[0] / compressed_kv.shape[0]:.0f}x") - - # Dequant compressed KV - compressed_kv_bf16 = kv_dequantize_fp8(compressed_kv, compressed_inv_scale) - c = F.cosine_similarity(kv_bf16[::cr].flatten().unsqueeze(0).float(), compressed_kv_bf16.flatten().unsqueeze(0).float()).item() - print(f" Compressed KV dequant cosine: {c:.6f} {'✅' if c>=0.99 else '❌'}") - - # ── Test 7: HCA compressed cache (cr=128) ──────────────────── - print("\n--- Test 7: HCA compressed cache (cr=128) ---") - cr = 128 - compressed_positions_128 = positions[::cr] - compressed_kv_128 = kv_fp8[::cr] if len(kv_fp8) >= cr else kv_fp8[:1] - compressed_inv_scale_128 = inv_scale[::cr] if len(inv_scale) >= cr else inv_scale[:1] - print(f" HCA compressed KV shape: {compressed_kv_128.shape}") - print(f" Tokens in HCA cache: {compressed_kv_128.shape[0]} (from {NT})") - - print(f"\n{'='*70}") - print(f" DONE — KV cache kernels tested") - print(f"{'='*70}") - - -if __name__ == "__main__": - main() diff --git a/tests/archive/test_layer_schedule.py b/tests/archive/test_layer_schedule.py deleted file mode 100644 index 00189470..00000000 --- a/tests/archive/test_layer_schedule.py +++ /dev/null @@ -1,85 +0,0 @@ -"""Tests for layer schedule — pure data, no kernels, no tensors.""" - -from dsv4.model.config import DSV4Config -from dsv4.model.layer_schedule import ( - AttentionType, FFNType, RouterMode, - LayerSpec, build_schedule, validate_schedule, -) - - -def test_flash_schedule(): - config = DSV4Config.flash() - schedule = build_schedule(config) - validate_schedule(schedule, config) - - assert len(schedule) == 43 - - # First two layers: SWA + hash routing - assert schedule[0].attn == AttentionType.SWA - assert schedule[1].attn == AttentionType.SWA - assert schedule[0].router_mode == RouterMode.HASH - assert schedule[1].router_mode == RouterMode.HASH - - # Layer 2: CSA + hash routing (last hash layer) - assert schedule[2].attn == AttentionType.CSA - assert schedule[2].router_mode == RouterMode.HASH - - # Layer 3: HCA + dense routing (first dense layer) - assert schedule[3].attn == AttentionType.HCA - assert schedule[3].router_mode == RouterMode.DENSE - - # Alternation continues - assert schedule[4].attn == AttentionType.CSA - assert schedule[5].attn == AttentionType.HCA - - # All layers are MoE - for spec in schedule: - assert spec.ffn == FFNType.MOE - - -def test_pro_schedule(): - config = DSV4Config.pro() - schedule = build_schedule(config) - validate_schedule(schedule, config) - - assert len(schedule) == 61 - - # First two layers: HCA + hash routing - assert schedule[0].attn == AttentionType.HCA - assert schedule[1].attn == AttentionType.HCA - assert schedule[0].router_mode == RouterMode.HASH - - # Layer 2: CSA + hash routing - assert schedule[2].attn == AttentionType.CSA - assert schedule[2].router_mode == RouterMode.HASH - - # Layer 3: HCA + dense routing - assert schedule[3].attn == AttentionType.HCA - assert schedule[3].router_mode == RouterMode.DENSE - - -def test_layer_spec_frozen(): - """LayerSpec is frozen — mutation should raise.""" - config = DSV4Config.flash() - spec = build_schedule(config)[0] - try: - spec.attn = AttentionType.HCA - assert False, "should have raised" - except AttributeError: - pass - - -def test_schedule_indices_match(): - """Each LayerSpec.layer_idx matches its position in the list.""" - config = DSV4Config.flash() - schedule = build_schedule(config) - for i, spec in enumerate(schedule): - assert spec.layer_idx == i - - -if __name__ == "__main__": - test_flash_schedule() - test_pro_schedule() - test_layer_spec_frozen() - test_schedule_indices_match() - print("All schedule tests passed") diff --git a/tests/archive/test_layout_compare.py b/tests/archive/test_layout_compare.py deleted file mode 100644 index 71a94aee..00000000 --- a/tests/archive/test_layout_compare.py +++ /dev/null @@ -1,95 +0,0 @@ -"""Compare C-fragment composition layout vs A-fragment layout for PV P operand.""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils -from cutlass.cute.nvgpu import tcgen05 -from cutlass import Float32, BFloat16 -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class LayoutCompareKernel: - def __init__(self): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = (128, 128) - self.cta_group = tcgen05.CtaGroup.ONE - self.threads_per_cta = 64 # minimal - - @cute.jit - def __call__(self, q, k, v, c, stream): - a_major = LayoutEnum.from_tensor(q).mma_major_mode() - b_major = LayoutEnum.from_tensor(k).mma_major_mode() - v_major = LayoutEnum.from_tensor(v).mma_major_mode() - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, a_major, b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, v_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - pv_mma_tiler = (qk_mma_tiler[0], qk_mma_tiler[2], qk_mma_tiler[1]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - - qk_acc_shape = qk_thr.partition_shape_C(qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - - # P A-fragment - p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, pv_mma_tiler, self.q_dtype, 1) - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - - # C-fragment composition layout - tilePlikeFP32 = qk_mma_tiler[1] // Float32.width * self.o_dtype.width - tStS_P_layout = cute.composition(tStS.layout, cute.make_layout((128, tilePlikeFP32))) - tStS_P = cute.make_tensor(tStS.iterator + 32, tStS_P_layout) # offset 32 FP32 columns - - # With scaled offset for A-fragment - p_offset_in_a_elements = self.qk_acc_dtype.width // self.q_dtype.width * 32 # = 64 - tOrP0 = cute.make_tensor(tOrP.iterator + p_offset_in_a_elements, tOrP.layout) - - # Print layouts - cute.printf("tStS layout: {}", tStS.layout) - cute.printf("tOrP layout: {}", tOrP.layout) - cute.printf("tStS_P layout: {}", tStS_P_layout) - cute.printf("tOrP0 layout: {}", tOrP0.layout) - cute.printf("tOrP shape: {}", tOrP.shape) - cute.printf("tStS_P shape: {}", tStS_P.shape) - cute.printf("tOtO layout: {}", tOtO.layout) - cute.printf("pv_mma_tiler: {}", pv_mma_tiler) - cute.printf("qk_mma_tiler: {}", qk_mma_tiler) - - -def test(): - m, n, head_dim = 128, 128, 64 - q = torch.randn(m, head_dim, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(n, head_dim, 1, dtype=torch.bfloat16, device='cuda') - v_base = torch.randn(head_dim, n, dtype=torch.bfloat16, device='cuda') - v = v_base.as_strided((head_dim, n), (1, head_dim)).unsqueeze(-1) - c = torch.zeros(m, head_dim, 1, dtype=torch.bfloat16, device='cuda') - - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - kernel = LayoutCompareKernel() - print('Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - print('Running...', flush=True) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_mma_si_only.py b/tests/archive/test_mma_si_only.py deleted file mode 100644 index 83a86319..00000000 --- a/tests/archive/test_mma_si_only.py +++ /dev/null @@ -1,247 +0,0 @@ -""" -Minimal test: Stage A + mma_si pipeline (no PV, no V). -If this deadlocks, the mma_si pipeline is broken. -If this passes, the deadlock is caused by adding V/PV. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - - -class MmaSiTest: - def __init__(self, mma_tiler_mn, use_2cta_instrs=False, use_tma_store=True): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = use_2cta_instrs; self.use_tma_store = use_tma_store - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.epilog_sync_bar_id = 1; self.tmem_alloc_sync_bar_id = 2; self.tmem_dealloc_sync_bar_id = 3 - self.num_c_stage = 2 - - def _setup(self, tiled_mma): - mma_inst_k = cute.size(tiled_mma.shape_mnk, mode=[2]) - self.mma_tiler = (*self.mma_tiler_mn, mma_inst_k * 4) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (tiled_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), - self.mma_tiler[1], self.mma_tiler[2]) - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, self.use_2cta_instrs, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(tiled_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(tiled_mma, self.mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) - tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(tCtAcc_fake, arch="sm_100") - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) - ) * cute.size(tiled_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - self.q_dtype = a.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(a).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(b).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - tiled_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - self._setup(tiled_mma) - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - tma_a, tma_ta = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma.thr_id), - a, a_smem, self.mma_tiler, tiled_mma, self.cluster_layout_vmnk.shape) - tma_b, tma_tb = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, tiled_mma.thr_id), - b, b_smem, self.mma_tiler, tiled_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(tiled_mma, tma_a, tma_ta, tma_b, tma_tb, tma_c, tma_tc, - self.cluster_layout_vmnk, self.a_smem_s, self.b_smem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, tiled_mma, tma_a, mA, tma_b, mB, tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(tiled_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_a); cpasync.prefetch_descriptor(tma_b); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] # ADDED: mma_si - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - # ADDED: mma_si pipeline (same as v27) - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=self.tmem_alloc_sync_bar_id, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sA = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sB = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gA = cute.local_tile(mA, cute.slice_(self.mma_tiler, (None,0,None)), (None,None,None)) - gB = cute.local_tile(mB, cute.slice_(self.mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gA, mode=[3]) - - thr_mma = tiled_mma.get_slice(0) - tCgA = thr_mma.partition_A(gA); tCgB = thr_mma.partition_B(gB); tCgC = thr_mma.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsA, tAgA = cpasync.tma_partition(tma_a, 0, a_lay, cute.group_modes(sA,0,3), cute.group_modes(tCgA,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsB, tBgB = cpasync.tma_partition(tma_b, 0, b_lay, cute.group_modes(sB,0,3), cute.group_modes(tCgB,0,3)) - tAgA = tAgA[(None,0,None,0)]; tBgB = tBgB[(None,0,None,0)] - - tCrA = tiled_mma.make_fragment_A(sA); tCrB = tiled_mma.make_fragment_B(sB) - acc_shape = thr_mma.partition_shape_C(self.mma_tiler[:2]) - tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # TMA WARP - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_a, tAgA[(None,h.count)], tAsA[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_b, tBgB[(None,h.count)], tBsB[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_mma_si_pv.py b/tests/archive/test_mma_si_pv.py deleted file mode 100644 index 182fd69f..00000000 --- a/tests/archive/test_mma_si_pv.py +++ /dev/null @@ -1,355 +0,0 @@ -""" -Stage B test: MMA + mma_si + V TMA + PV MMA. -Built incrementally from working test_mma_si_only. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - - -class MmaSiPvTest: - def __init__(self, mma_tiler_mn, head_dim, use_2cta_instrs=False, use_tma_store=True): - self.head_dim = head_dim - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = use_2cta_instrs; self.use_tma_store = use_tma_store - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.epilog_sync_bar_id = 1; self.tmem_alloc_sync_bar_id = 2; self.tmem_dealloc_sync_bar_id = 3 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - self.pv_mma_tiler = (self.qk_mma_tiler[0], self.qk_mma_tiler[2], self.qk_mma_tiler[1]) - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, self.use_2cta_instrs, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q: cute.Tensor, k: cute.Tensor, v: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - # Compute PV tiler: swap N and K from QK tiler (FMHA convention) - # QK: (M=128, N=128, K=64) -> PV: (M=128, N=64, K=128) - # PV mma_tiler_mn is (M, N_pv) = (128, head_dim=64), NOT (128, 128) - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - # BUG FIX: pv_mma_tiler_mn must be (M, head_dim), not (M, N_qk) - # Passing mma_tiler_mn=(128,128) creates a (128,128) MMA that expects 128-column output - # but PV output is (128,64). This caused cosine ~0.01. - pv_mma_tiler_mn = (self.mma_tiler_mn[0], self.head_dim) # (128, 64) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, pv_mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=self.tmem_alloc_sync_bar_id, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.pv_mma_tiler, (None,0,None)), (None,None,None)) # Use PV tiler for output - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = pv_thr.partition_C(gC) # PV output: partition with pv_thr - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_model_forward_b200.py b/tests/archive/test_model_forward_b200.py deleted file mode 100644 index 05a62d23..00000000 --- a/tests/archive/test_model_forward_b200.py +++ /dev/null @@ -1,238 +0,0 @@ -#!/usr/bin/env python3 -""" -Reproduce the vLLM empty-output bug outside the container. - -Strategy: Run the model in FULL BF16 (dequantized weights) and compare -against CuTeDSL at each projection. Also check: does the warmup gs -cause issues at inference time? - -Key diagnostic: inspect CuTeDSL runner.run() to see if it uses the -fixed warmup gs or recomputes per-call. - -Usage (on B200): - source /root/nvfp4-megamoe-kernel/tests/.venv/bin/activate - python3 tests/test_model_forward_b200.py -""" - -import sys, os, json, torch, torch.nn.functional as F, inspect -from safetensors import safe_open - -REPO = "/root/nvfp4-megamoe-kernel" -sys.path.insert(0, REPO) -MODEL = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4" -DEV = "cuda:0" - -H = 7168; NH = 128; HD = 512; NOPE = 448; ROPE = 64 -QL = 1536; OL = 1024; OG = 16; HPG = NH // OG -EPS = 1e-6 - -E2M1 = torch.tensor([0,.5,1.,1.5,2.,3.,4.,6.,-0,-.5,-1.,-1.5,-2.,-3.,-4.,-6.], dtype=torch.float32) - -_cache = {} -def P(k, wm, md): - if k in _cache: return _cache[k] - with safe_open(os.path.join(md, wm[k]), framework="pt") as f: - t = f.get_tensor(k) - _cache[k] = t - return t - -def dequant(w, sf, gs): - d = w.device; lut = E2M1.to(d) - lo = lut[(w & 0xF).long()]; hi = lut[((w >> 4) & 0xF).long()] - O, I2 = w.shape; I = I2*2 - u = torch.empty(O, I, dtype=torch.float32, device=d) - u[:,0::2] = lo; u[:,1::2] = hi - bs = sf.float().repeat_interleave(16, dim=1)[:O,:I] - return (u * bs * gs).to(torch.bfloat16) - -def rms(x, w, eps=1e-6): - v = x.float().pow(2).mean(-1, keepdim=True) - return (w.float() * (x * torch.rsqrt(v+eps)).float()).to(x.dtype) - -def make_runner(w, sf, gs_t, inf, outf, fused=False, lw=None): - from dsv4.layers.linear import Nvfp4Linear - fp4 = w.view(torch.float4_e2m1fn_x2).permute(1,0).contiguous() - s = sf.to(torch.float8_e4m3fn) if sf.dtype != torch.float8_e4m3fn else sf - s = s.permute(1,0).contiguous() - if fused and gs_t.numel() == 2: - g1,g2 = gs_t[0].item(), gs_t[1].item(); gs = max(g1,g2) - if g1 != g2: - s32 = s.float(); sp = lw[0] if lw else outf//2 - s32[:sp] *= g1/gs; s32[sp:] *= g2/gs; s = s32.to(torch.float8_e4m3fn) - else: - gs = gs_t.max().item() if gs_t.numel() > 1 else gs_t.item() - r = Nvfp4Linear(in_features=inf, out_features=outf, max_num_tokens=8192, device=str(w.device)) - r.fp4 = [fp4]; r.sf = [s]; r.gs = [gs] - r.finalize_weights(); r._ensure_initialized() - return r - - -def main(): - torch.cuda.set_device(0) - - print("=" * 70) - print(" Diagnose: Why does vLLM produce empty output?") - print("=" * 70) - - with open(os.path.join(MODEL, "model.safetensors.index.json")) as f: - wm = json.load(f)["weight_map"] - G = lambda k: P(k, wm, MODEL).to(DEV) - - # ── INSPECT: How does CuTeDSL runner.run() use gs? ──────────────── - print("\n--- INSPECTING CuTeDSL runner internals ---") - from dsv4.layers.linear import Nvfp4Linear -from dsv4.ops.quantize import ( - quantize_activation_nvfp4, -) - - print("\n quantize_activation_nvfp4 signature:") - sig = inspect.signature(quantize_activation_nvfp4) - print(f" {sig}") - - print("\n Nvfp4Linear._run_impl source (key lines):") - src = inspect.getsource(Nvfp4Linear._run_impl) - for i, line in enumerate(src.split('\n')): - stripped = line.strip() - if any(kw in stripped for kw in ['global_scale', '_activation', 'quantize', 'return', 'def ']): - print(f" L{i}: {stripped}") - - # ── CRITICAL TEST: warmup gs vs per-input gs ────────────────────── - print("\n--- CRITICAL TEST: warmup gs vs per-input gs ---") - - p = "model.layers.0"; a = f"{p}.self_attn" - qa_w = G(f"{a}.q_a_proj.weight"); qa_sf = G(f"{a}.q_a_proj.weight_scale"); qa_gs = G(f"{a}.q_a_proj.weight_scale_2") - - # Load embedding + norm weight - emb = G("model.embed_tokens.weight") - anorm = G(f"{p}.input_layernorm.weight") - - # Create runner with warmup gs - r = make_runner(qa_w, qa_sf, qa_gs, qa_w.shape[1]*2, qa_w.shape[0]) - torch.manual_seed(42) - warmup = torch.randn(1, H, dtype=torch.bfloat16, device=DEV)*2.0 - with torch.no_grad(): - r.compute_activation_global_scale(warmup) - print(f" Warmup gs (random input amax={warmup.amax():.4f}): {r._activation_global_scale:.8f}") - - # Get REAL input (embedding output) - token_ids = torch.tensor([1, 450, 8403, 315, 5413, 374], dtype=torch.long, device=DEV) - with torch.no_grad(): - hidden = emb[token_ids] - normed = rms(hidden, anorm, EPS) - print(f" Real input (after RMS norm) amax: {normed.amax():.4f}") - - # What gs would the real input need? - real_gs = normed.amax().item() / (6.0 * 448.0) - print(f" Correct gs for real input: {real_gs:.8f}") - print(f" Ratio warmup/correct: {r._activation_global_scale / real_gs:.4f}" if real_gs > 0 else " real_gs is 0!") - - # Run with warmup gs - with torch.no_grad(): - out_warmup = r.run(normed) - - # Run with dynamic gs (recompute for this input) - r2 = make_runner(qa_w, qa_sf, qa_gs, qa_w.shape[1]*2, qa_w.shape[0]) - with torch.no_grad(): - r2.compute_activation_global_scale(normed) - out_dynamic = r2.run(normed) - - # BF16 reference - qa_bf16 = dequant(qa_w, qa_sf, qa_gs.item()) - with torch.no_grad(): - ref = normed @ qa_bf16.T - - c_warmup = F.cosine_similarity(out_warmup.flatten().unsqueeze(0).float(), ref.flatten().unsqueeze(0).float()).item() - c_dynamic = F.cosine_similarity(out_dynamic.flatten().unsqueeze(0).float(), ref.flatten().unsqueeze(0).float()).item() - - print(f"\n q_a cosine vs BF16 (warmup gs): {c_warmup:.6f} {'✅' if c_warmup>=0.98 else '❌'}") - print(f" q_a cosine vs BF16 (dynamic gs): {c_dynamic:.6f} {'✅' if c_dynamic>=0.98 else '❌'}") - print(f" amax warmup: {out_warmup.amax():.4f} amax dynamic: {out_dynamic.amax():.4f} amax ref: {ref.amax():.4f}") - - # ── Test: run FULL model in BF16 (1 layer) then check logits ────── - print("\n--- FULL BF16 model: 1 layer → LM head ---") - lm_head = G("lm_head.weight") - fnorm_w = G("model.norm.weight") - qn = G(f"{a}.q_a_norm.weight") - kvn = G(f"{a}.kv_norm.weight") - fnorm_l0 = G(f"{p}.post_attention_layernorm.weight") - - # Dequantize all layer 0 attention weights - qa_bf16 = dequant(qa_w, qa_sf, qa_gs.item()) - qb_w = G(f"{a}.q_b_proj.weight"); qb_sf = G(f"{a}.q_b_proj.weight_scale"); qb_gs = G(f"{a}.q_b_proj.weight_scale_2") - kv_bf16 = dequant(G(f"{a}.kv_proj.weight"), G(f"{a}.kv_proj.weight_scale"), G(f"{a}.kv_proj.weight_scale_2").item()) - qb_bf16 = dequant(qb_w, qb_sf, qb_gs.item()) - woa = G(f"{a}.o_a_proj.weight") # already BF16 - wob_bf16 = dequant(G(f"{a}.o_b_proj.weight"), G(f"{a}.o_b_proj.weight_scale"), G(f"{a}.o_b_proj.weight_scale_2").item()) - - with torch.no_grad(): - x = hidden.clone() - print(f" Input: amax={x.amax():.4f}") - - # RMS norm - x = rms(x, anorm, EPS) - - # Attention projections (BF16) - qa = x @ qa_bf16.T - kv = x @ kv_bf16.T - qa_n = rms(qa, qn, EPS) - qb = qa_n @ qb_bf16.T - print(f" q_a: amax={qa.amax():.4f}, kv: amax={kv.amax():.4f}, q_b: amax={qb.amax():.4f}") - - # Skip attention, use random output - o = torch.randn(len(token_ids), NH, HD, dtype=torch.bfloat16, device=DEV) * 0.1 - - # wo_a: BMM (o_a_proj is (OG*OL, HPG*HD)) - o_grouped = o.view(len(token_ids), OG, HPG * HD).permute(1, 0, 2) - woa_3d = woa.view(OG, OL, HPG * HD) - z = torch.bmm(o_grouped, woa_3d.transpose(1, 2)).permute(1, 0, 2).reshape(len(token_ids), OG * OL) - - # wo_b - attn_out = z @ wob_bf16.T - print(f" attn_out (BF16): amax={attn_out.amax():.4f}") - - # Skip MoE, just add residual - x = hidden + attn_out - - # Final norm + LM head - x_normed = rms(x, fnorm_w, EPS) - logits = x_normed @ lm_head.T - print(f" logits: amax={logits.amax():.4f} NaN={torch.isnan(logits).any()}") - - top5 = torch.topk(logits[-1], 5) - print(f" top5 IDs: {top5.indices.tolist()}") - print(f" top5 logits: {[f'{v:.2f}' for v in top5.values.tolist()]}") - log_std = logits[-1].float().std().item() - print(f" logit std: {log_std:.4f}") - - # ── KEY INSIGHT: check if the runner re-reads gs at inference time ─ - print("\n" + "=" * 70) - print(" KEY: Does runner.run() use FIXED warmup gs or RECOMPUTE?") - print("=" * 70) - - # Monkey-patch the gs and see if output changes - r3 = make_runner(qa_w, qa_sf, qa_gs, qa_w.shape[1]*2, qa_w.shape[0]) - with torch.no_grad(): - r3.compute_activation_global_scale(normed) - gs_original = r3._activation_global_scale - out_original = r3.run(normed).clone() - - # Change gs by 10x - r3._activation_global_scale = gs_original * 10.0 - out_changed = r3.run(normed).clone() - - c_changed = F.cosine_similarity(out_original.flatten().unsqueeze(0).float(), out_changed.flatten().unsqueeze(0).float()).item() - print(f" Original gs: {gs_original:.8f}") - print(f" Changed gs: {gs_original * 10:.8f}") - print(f" Cosine sim after 10x gs change: {c_changed:.6f}") - if abs(c_changed - 1.0) < 0.001: - print(" ➡️ Changing gs has NO effect on output!") - print(" ➡️ The runner recomputes gs internally at inference time.") - print(" ➡️ Warmup gs is IRRELEVANT — the bug is elsewhere.") - else: - print(" ➡️ Changing gs DOES change the output!") - print(" ➡️ The runner uses the warmup gs at inference time.") - print(" ➡️ Wrong warmup gs would cause wrong quantization → garbage.") - - -if __name__ == "__main__": - main() diff --git a/tests/archive/test_moe_nan_b200.py b/tests/archive/test_moe_nan_b200.py deleted file mode 100644 index c2a5b8aa..00000000 --- a/tests/archive/test_moe_nan_b200.py +++ /dev/null @@ -1,225 +0,0 @@ -#!/usr/bin/env python3 -""" -DeepSeek-V4 MoE NaN Reproduction Test - -Finds where NaN originates in the MoE forward pass. -Tests individual experts (gate+up+down) with the CuTeDSL NVFP4 linear runner. -Then tests the grouped GEMM MoE runner with stacked weights. - -Key insight: DeepSeek-V4 is a MegaMoE with 384 experts. -The NaN might come from: -1. Weight loading / quantization -2. Activation quantization (quantize_activation_nvfp4) -3. The grouped GEMM kernel -4. The combine/scatter step - -Usage (on B200): - cd /root/nvfp4-megamoe-kernel - PYTHONPATH=/root/nvfp4-megamoe-kernel tests/venv/bin/python tests/test_moe_nan_b200.py -""" - -import sys, os, json, torch, torch.nn.functional as F -from safetensors import safe_open - -REPO = "/root/nvfp4-megamoe-kernel" -sys.path.insert(0, REPO) -MODEL = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4" -DEV = "cuda:0" - -H = 7168 -INTERMEDIATE = 3072 # DeepSeek-V4 MoE intermediate -NUM_EXPERTS = 384 -TOPK = 6 -EPS = 1e-6 - -E2M1 = torch.tensor([0,.5,1.,1.5,2.,3.,4.,6.,-0,-.5,-1.,-1.5,-2.,-3.,-4.,-6.], dtype=torch.float32) - -_cache = {} -def P(k, wm, md): - if k in _cache: return _cache[k] - with safe_open(os.path.join(md, wm[k]), framework="pt") as f: - t = f.get_tensor(k) - _cache[k] = t - return t - -def rms(x, w, eps=1e-6): - v = x.float().pow(2).mean(-1, keepdim=True) - return (w.float() * (x * torch.rsqrt(v+eps)).float()).to(x.dtype) - -def make_runner(w, sf, gs_t, inf, outf): - from dsv4.layers.linear import Nvfp4Linear - fp4 = w.view(torch.float4_e2m1fn_x2).permute(1,0).contiguous() - s = sf.to(torch.float8_e4m3fn) if sf.dtype != torch.float8_e4m3fn else sf - s = s.permute(1,0).contiguous() - gs = gs_t.max().item() if gs_t.numel() > 1 else gs_t.item() - r = Nvfp4Linear(in_features=inf, out_features=outf, max_num_tokens=8192, device=str(w.device)) - r.fp4 = [fp4]; r.sf = [s]; r.gs = [gs] - r.finalize_weights(); r._ensure_initialized() - return r - - -def test_single_expert(layer_id=2, expert_id=0): - """Test a single expert's gate+up+down with CuTeDSL NVFP4 linear.""" - torch.cuda.set_device(0) - torch.cuda.empty_cache() - _cache.clear() - - with open(os.path.join(MODEL, "model.safetensors.index.json")) as f: - wm = json.load(f)["weight_map"] - G = lambda k: P(k, wm, MODEL).to(DEV) - - p = f"model.layers.{layer_id}" - m = f"{p}.mlp" - e = f"{m}.experts.{expert_id}" - - emb = G("model.embed_tokens.weight") - fnorm = G(f"{p}.post_attention_layernorm.weight") - - # Load expert weights - gate_w = G(f"{e}.gate_proj.weight"); gate_sf = G(f"{e}.gate_proj.weight_scale"); gate_gs = G(f"{e}.gate_proj.weight_scale_2") - up_w = G(f"{e}.up_proj.weight"); up_sf = G(f"{e}.up_proj.weight_scale"); up_gs = G(f"{e}.up_proj.weight_scale_2") - down_w = G(f"{e}.down_proj.weight"); down_sf = G(f"{e}.down_proj.weight_scale"); down_gs = G(f"{e}.down_proj.weight_scale_2") - - print(f" Expert {expert_id}:") - print(f" gate: shape={gate_w.shape} dtype={gate_w.dtype} sf_shape={gate_sf.shape} gs={gate_gs.tolist()}") - print(f" up: shape={up_w.shape} dtype={up_w.dtype}") - print(f" down: shape={down_w.shape} dtype={down_w.dtype}") - print(f" gate NaN: {torch.isnan(gate_w.float()).any()}") - print(f" gate_gs NaN: {torch.isnan(gate_gs).any()}") - print(f" gate input_scale: exists={f'{e}.gate_proj.input_scale' in wm}") - - # Check for zero or extreme gs values - for name, gs in [("gate", gate_gs), ("up", up_gs), ("down", down_gs)]: - if gs.numel() > 0: - print(f" {name} gs: min={gs.min().item():.6f} max={gs.max().item():.6f}") - if gs.min().item() == 0: - print(f" WARNING: {name} gs has zero value — will cause division by zero!") - - r_gate = make_runner(gate_w, gate_sf, gate_gs, H, gate_w.shape[0]) - r_up = make_runner(up_w, up_sf, up_gs, H, up_w.shape[0]) - r_down = make_runner(down_w, down_sf, down_gs, INTERMEDIATE, down_w.shape[0]) - - # Test with various token counts - for num_tokens in [1, 4, 8, 16]: - token_ids = torch.randint(1, 1000, (num_tokens,), dtype=torch.long, device=DEV) - hidden = emb[token_ids] - normed = rms(hidden, fnorm, EPS) - - with torch.no_grad(): - gate_out = r_gate.run(normed) - up_out = r_up.run(normed) - - # Check gate and up - gate_nan = torch.isnan(gate_out).any().item() - up_nan = torch.isnan(up_out).any().item() - - if gate_nan or up_nan: - print(f" {num_tokens} tokens: gate NaN={gate_nan} up NaN={up_nan}") - # Find which row has NaN - gate_nan_rows = torch.isnan(gate_out).any(dim=1).nonzero().flatten().tolist() - print(f" Gate NaN rows: {gate_nan_rows}") - continue - - # SiLU activation - activated = F.silu(gate_out) * up_out - act_nan = torch.isnan(activated).any().item() - - if act_nan: - print(f" {num_tokens} tokens: NaN after SiLU activation!") - continue - - down_out = r_down.run(activated) - down_nan = torch.isnan(down_out).any().item() - - if down_nan: - print(f" {num_tokens} tokens: down NaN={down_nan}") - continue - - print(f" {num_tokens} tokens: amax={down_out.amax():.4f} OK") - - del r_gate, r_up, r_down - torch.cuda.empty_cache() - - -def test_quantize_activation(): - """Test the activation quantization used by the MoE grouped GEMM.""" - from cutedsl.nvfp4_linear import quantize_activation_nvfp4 - - torch.cuda.set_device(0) - - for num_tokens in [1, 4, 8, 16]: - # Create realistic input (after SiLU * up) - x = torch.randn(num_tokens, INTERMEDIATE, dtype=torch.bfloat16, device=DEV) - - # quantize_activation_nvfp4 returns (x_sf, x_gs) or similar - # The grouped GEMM needs quantized activation as input - try: - result = quantize_activation_nvfp4(x, num_tokens) - if isinstance(result, tuple): - for i, r in enumerate(result): - if r is not None and r.is_floating_point(): - print(f" {num_tokens} tokens: quantize result[{i}] NaN={torch.isnan(r).any()}") - else: - print(f" {num_tokens} tokens: quantize result NaN={torch.isnan(result).any()}") - except Exception as e: - print(f" {num_tokens} tokens: quantize failed: {e}") - - print() - - -def test_grouped_gemm_shapes(): - """Test the CuTeDSL grouped GEMM with MegaMoE-like shapes.""" - from cutedsl.moe import run_nvfp4_grouped_gemm - - torch.cuda.set_device(0) - - # Create simple test: 4 experts, 8 tokens, top-2 - num_experts = 4 - num_tokens = 8 - hidden_size = 512 # Small for testing - intermediate_size = 256 - - # Allocate weight tensors (random) - # L1: (num_experts, 2*intermediate_size, hidden_size//2) fp4 - # L2: (num_experts, hidden_size, intermediate_size//2) fp4 - l1_shape = (num_experts, 2 * intermediate_size, hidden_size // 2) - l2_shape = (num_experts, hidden_size, intermediate_size // 2) - - print(f" Testing grouped GEMM with:") - print(f" num_experts={num_experts}, num_tokens={num_tokens}") - print(f" l1_shape={l1_shape}, l2_shape={l2_shape}") - - # This test just checks if the kernel can handle various expert distributions - # without NaN. The actual weight values don't matter for NaN detection. - print(f" (Skipping — requires proper weight packing from vLLM model loader)") - print(f" The CuTeDSL grouped GEMM needs weights in a specific packed format") - print(f" that the vLLM model loader creates during model initialization.") - print() - - -def main(): - print("=" * 70) - print(" DeepSeek-V4 MoE NaN Reproduction Test") - print(" Finds where NaN originates in the MoE forward pass") - print("=" * 70) - - print("\n=== Test 1: Single expert gate+up+down ===") - for expert_id in [0, 1, 100, 383]: - test_single_expert(layer_id=2, expert_id=expert_id) - _cache.clear() - - print("\n=== Test 2: Activation quantization ===") - test_quantize_activation() - - print("\n=== Test 3: Grouped GEMM shapes ===") - test_grouped_gemm_shapes() - - print(f"\n{'='*70}") - print(f" Summary: If single experts produce NaN, the issue is in weight") - print(f" loading or the CuTeDSL NVFP4 linear kernel. If they're fine,") - print(f" the NaN comes from the grouped GEMM or the combine step.") - print(f"{'='*70}") - - -if __name__ == "__main__": - main() diff --git a/tests/archive/test_moe_runner_nan_b200.py b/tests/archive/test_moe_runner_nan_b200.py deleted file mode 100644 index ae447a9b..00000000 --- a/tests/archive/test_moe_runner_nan_b200.py +++ /dev/null @@ -1,190 +0,0 @@ -#!/usr/bin/env python3 -""" -DeepSeek-V4 MoE Runner NaN Test - -Tests the Nvfp4MoE (grouped GEMM path) with real weights. -The single-expert tests pass — this test exercises the FULL MoE runner -with routing, padding, grouped GEMM, and combine. - -Usage (on B200): - cd /root/nvfp4-megamoe-kernel - PYTHONPATH=/root/nvfp4-megamoe-kernel tests/venv/bin/python tests/test_moe_runner_nan_b200.py -""" - -import sys, os, json, torch, torch.nn.functional as F -from safetensors import safe_open - -REPO = "/root/nvfp4-megamoe-kernel" -sys.path.insert(0, REPO) -MODEL = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4" -DEV = "cuda:0" - -H = 7168 -INTERMEDIATE = 3072 -NUM_EXPERTS = 384 -TOPK = 6 -EPS = 1e-6 - -_cache = {} -def P(k, wm, md): - if k in _cache: return _cache[k] - with safe_open(os.path.join(md, wm[k]), framework="pt") as f: - t = f.get_tensor(k) - _cache[k] = t - return t - -def rms(x, w, eps=1e-6): - v = x.float().pow(2).mean(-1, keepdim=True) - return (w.float() * (x * torch.rsqrt(v+eps)).float()).to(x.dtype) - - -def pack_expert_weights(wm, G, layer_id=2, num_local_experts=16): - """Pack per-expert weights into stacked format for Nvfp4MoE. - Only loads the first num_local_experts to fit in memory. - """ - m = f"model.layers.{layer_id}.mlp" - - # Load expert weights and stack (only first num_local_experts) - gate_ws, gate_sfs, gate_gss = [], [], [] - up_ws, up_sfs, up_gss = [], [], [] - down_ws, down_sfs, down_gss = [], [], [] - - for i in range(num_local_experts): - e = f"{m}.experts.{i}" - gate_ws.append(G(f"{e}.gate_proj.weight")) - gate_sfs.append(G(f"{e}.gate_proj.weight_scale")) - gate_gs = G(f"{e}.gate_proj.weight_scale_2") - gate_gss.append(gate_gs) - - up_ws.append(G(f"{e}.up_proj.weight")) - up_sfs.append(G(f"{e}.up_proj.weight_scale")) - up_gs = G(f"{e}.up_proj.weight_scale_2") - up_gss.append(up_gs) - - down_ws.append(G(f"{e}.down_proj.weight")) - down_sfs.append(G(f"{e}.down_proj.weight_scale")) - down_gs = G(f"{e}.down_proj.weight_scale_2") - down_gss.append(down_gs) - - if i % 50 == 0: - print(f" Loaded expert {i}/{num_local_experts}") - - # Stack into (E, ...) tensors - w13_w = torch.stack(gate_ws) # (E, 3072, 3584) - w13_sf = torch.stack(gate_sfs) - w13_gs = torch.stack(gate_gss) if gate_gss[0].dim() > 0 else torch.tensor([g.item() for g in gate_gss], device=DEV) - - # Actually w13 = stacked gate+up, w2 = down - # But our runner expects separate L1 (gate+up) and L2 (down) - # The w13 format is (E, 2*intermediate, hidden//2) with gate and up interleaved - # For Nvfp4MoE, we stack gate and up side-by-side - - # Stack gate and up into w13 format: (E, 2*intermediate, hidden//2) - w13_w = torch.cat([torch.stack(gate_ws), torch.stack(up_ws)], dim=1) # (E, 6144, 3584) - w13_sf = torch.cat([torch.stack(gate_sfs), torch.stack(up_sfs)], dim=1) - w13_gs = torch.cat([torch.stack(gate_gss), torch.stack(up_gss)], dim=0) - - w2_w = torch.stack(down_ws) - w2_sf = torch.stack(down_sfs) - w2_gs = torch.stack(down_gss) - - return w13_w, w13_sf, w13_gs, w2_w, w2_sf, w2_gs - - -def test_moe_runner(layer_id=2): - """Test the Nvfp4MoE with real weights.""" - from dsv4.layers.moe import Nvfp4MoE - - torch.cuda.set_device(0) - torch.manual_seed(42) - torch.cuda.empty_cache() - _cache.clear() - - with open(os.path.join(MODEL, "model.safetensors.index.json")) as f: - wm = json.load(f)["weight_map"] - G = lambda k: P(k, wm, MODEL).to(DEV) - - p = f"model.layers.{layer_id}" - m = f"{p}.mlp" - - emb = G("model.embed_tokens.weight") - fnorm = G(f"{p}.post_attention_layernorm.weight") - - print(f" Packing expert weights (384 experts)...") - # Test with fewer experts to fit in memory - num_local_experts = 16 # Use 16 experts (out of 384) for testing - - # Create the runner first, then prepare weights - intermediate_size = INTERMEDIATE # 3072 - hidden_size = H # 7168 - - runner = Nvfp4MoE( - num_experts=num_local_experts, - hidden_size=hidden_size, - intermediate_size=intermediate_size, - max_num_tokens=8192, - top_k=TOPK, - device=str(DEV), - ) - - # Load and pack weights - print(f" Loading expert weights...") - w13_w, w13_sf, w13_gs, w2_w, w2_sf, w2_gs = pack_expert_weights(wm, G, layer_id, num_local_experts) - - print(f" w13_w: {w13_w.shape}, w2_w: {w2_w.shape}") - print(f" w13_gs: {w13_gs.shape}, w2_gs: {w2_gs.shape}") - print(f" w13 NaN: {torch.isnan(w13_w.float()).any()}") - print(f" w2 NaN: {torch.isnan(w2_w.float()).any()}") - - # Prepare weights for the runner - l1_fp4 = w13_w.view(torch.float4_e2m1fn_x2) - l2_fp4 = w2_w.view(torch.float4_e2m1fn_x2) - l1_sf = w13_sf.to(torch.float8_e4m3fn) if w13_sf.dtype != torch.float8_e4m3fn else w13_sf - l2_sf = w2_sf.to(torch.float8_e4m3fn) if w2_sf.dtype != torch.float8_e4m3fn else w2_sf - - runner.prepare_weights_from_stacked( - l1_fp4, l1_sf, w13_gs.tolist() if w13_gs.dim() == 1 else w13_gs.flatten().tolist(), - l2_fp4, l2_sf, w2_gs.tolist() if w2_gs.dim() == 1 else w2_gs.flatten().tolist(), - ) - - # Test with various token counts - for num_tokens in [1, 4, 8, 16]: - token_ids = torch.randint(1, 1000, (num_tokens,), dtype=torch.long, device=DEV) - hidden = emb[token_ids] - normed = rms(hidden, fnorm, EPS) - - topk_ids = torch.randint(0, num_local_experts, (num_tokens, TOPK), device=DEV) - print(f" {num_tokens} tokens: input amax={normed.amax():.4f} NaN={torch.isnan(normed).any()}") - topk_weights = torch.softmax(torch.randn(num_tokens, TOPK, device=DEV), dim=-1) - - print(f" {num_tokens} tokens: input amax={normed.amax():.4f} NaN={torch.isnan(normed).any()}") - - with torch.no_grad(): - result = runner.run(normed, topk_weights, topk_ids) - - result_nan = torch.isnan(result).any().item() - result_amax = result.amax().item() if not result_nan else -1 - print(f" {num_tokens} tokens: output amax={result_amax:.4f} NaN={result_nan}") - - if result_nan: - nan_rows = torch.isnan(result).any(dim=1).sum().item() - print(f" {num_tokens} tokens: {nan_rows}/{num_tokens} rows have NaN") - - del runner, w13_w, w13_sf, w13_gs, w2_w, w2_sf, w2_gs - torch.cuda.empty_cache() - _cache.clear() - - -def main(): - print("=" * 70) - print(" DeepSeek-V4 MoE Runner NaN Test") - print(" Tests Nvfp4MoE (grouped GEMM) with real weights") - print("=" * 70) - - test_moe_runner(layer_id=2) - - print(f"\n{'='*70}") - - -if __name__ == "__main__": - main() diff --git a/tests/archive/test_multilayer.py b/tests/archive/test_multilayer.py deleted file mode 100644 index 829256ff..00000000 --- a/tests/archive/test_multilayer.py +++ /dev/null @@ -1,161 +0,0 @@ -"""Extended pipeline test: simulate multi-layer MoE to check for error accumulation. -Uses same config as vLLM: max_num_tokens=8192, max_chunks=8, 48 experts.""" -import torch -import torch.nn.functional as F -import sys, os, glob - -sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), '..')) - -MODEL_PATH = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4" -LAYER_IDX = 0 # Use layer 0 weights for all layers (just testing accumulation) -NUM_EXPERTS = 48 -HIDDEN_SIZE = 7168 -INTERMEDIATE_SIZE = 3072 -NUM_TOKENS = 5 # "The capital of France is" -TOP_K = 6 -SWIGLU_LIMIT = 10.0 -DEVICE = "cuda" -NUM_LAYERS = 3 # Test error accumulation over multiple layers - - -def load_layer_tensors(model_dir, layer_idx): - tensors = {} - for sf in glob.glob(os.path.join(model_dir, "*.safetensors")): - from safetensors.torch import load_file - data = load_file(sf) - for k, v in data.items(): - if f"layers.{layer_idx}." in k and "mlp.experts" in k: - tensors[k.removeprefix("model.")] = v - return tensors - - -def dequantize_nvfp4_weight(packed_uint8, scale_e4m3, global_scale): - lut = torch.tensor([0.,0.5,1.,1.5,2.,3.,4.,6.,-0.,-0.5,-1.,-1.5,-2.,-3.,-4.,-6.], - dtype=torch.float32, device=packed_uint8.device) - lower = lut[(packed_uint8 & 0x0F).long()] - upper = lut[((packed_uint8 >> 4) & 0x0F).long()] - N, K = packed_uint8.shape[0], packed_uint8.shape[1] * 2 - bf16 = torch.stack([lower, upper], dim=-1).reshape(N, K) - K_sf = scale_e4m3.shape[1] - scale_2d = scale_e4m3.float().repeat_interleave(K // K_sf, dim=1) - return (bf16 * scale_2d * global_scale).to(torch.bfloat16) - - -def main(): - torch.cuda.set_device(0) - torch.manual_seed(42) - - print(f"=== Multi-Layer Pipeline Test ({NUM_LAYERS} layers) ===") - nvfp4_tensors = load_layer_tensors(MODEL_PATH, LAYER_IDX) - expert_indices = list(range(NUM_EXPERTS)) - - # Start with random hidden states (like after embedding + first attention) - hidden = torch.randn(NUM_TOKENS, HIDDEN_SIZE, dtype=torch.bfloat16, device=DEVICE) * 2.0 - - topk_ids = torch.zeros(NUM_TOKENS, TOP_K, dtype=torch.int64, device=DEVICE) - for i in range(NUM_TOKENS): - topk_ids[i] = torch.randperm(NUM_EXPERTS)[:TOP_K] - topk_weights = torch.ones(NUM_TOKENS, TOP_K, dtype=torch.float32, device=DEVICE) / TOP_K - - # Setup runner - from vllm.nvfp4_cutedsl import Nvfp4MoE -from dsv4.ops.layouts import ( - assemble_scales_3d_side, - make_b_k_major, -) - - l1_fp4, l1_sf, l1_gs_list = [], [], [] - l2_fp4, l2_sf, l2_gs_list = [], [], [] - for e in expert_indices: - gw = nvfp4_tensors[f"layers.{LAYER_IDX}.mlp.experts.{e}.gate_proj.weight"].to(DEVICE) - uw = nvfp4_tensors[f"layers.{LAYER_IDX}.mlp.experts.{e}.up_proj.weight"].to(DEVICE) - gsf = nvfp4_tensors[f"layers.{LAYER_IDX}.mlp.experts.{e}.gate_proj.weight_scale"].to(DEVICE) - usf = nvfp4_tensors[f"layers.{LAYER_IDX}.mlp.experts.{e}.up_proj.weight_scale"].to(DEVICE) - ggs = nvfp4_tensors[f"layers.{LAYER_IDX}.mlp.experts.{e}.gate_proj.weight_scale_2"].item() - ugs = nvfp4_tensors[f"layers.{LAYER_IDX}.mlp.experts.{e}.up_proj.weight_scale_2"].item() - fw = torch.cat([gw, uw], dim=0).view(torch.float4_e2m1fn_x2).permute(1,0).contiguous() - fsf = torch.cat([gsf, usf], dim=0).permute(1,0).contiguous() - mgs = max(ggs, ugs) - if ggs != ugs: - sf32 = fsf.float() - sf32[:, :INTERMEDIATE_SIZE] *= (ggs / mgs) - sf32[:, INTERMEDIATE_SIZE:] *= (ugs / mgs) - fsf = sf32.to(torch.float8_e4m3fn) - l1_fp4.append(fw); l1_sf.append(fsf); l1_gs_list.append(mgs) - dk = f"layers.{LAYER_IDX}.mlp.experts.{e}.down_proj.weight" - if dk in nvfp4_tensors: - dw = nvfp4_tensors[dk].to(DEVICE) - dsf = nvfp4_tensors[f"layers.{LAYER_IDX}.mlp.experts.{e}.down_proj.weight_scale"].to(DEVICE) - dgs = nvfp4_tensors[f"layers.{LAYER_IDX}.mlp.experts.{e}.down_proj.weight_scale_2"].item() - l2_fp4.append(dw.view(torch.float4_e2m1fn_x2).permute(1,0).contiguous()) - l2_sf.append(dsf.permute(1,0).contiguous()); l2_gs_list.append(dgs) - else: - l2_fp4.append(torch.zeros(INTERMEDIATE_SIZE//2, HIDDEN_SIZE, dtype=torch.float4_e2m1fn_x2, device=DEVICE)) - l2_sf.append(torch.ones(INTERMEDIATE_SIZE//16, HIDDEN_SIZE, dtype=torch.float8_e4m3fn, device=DEVICE)) - l2_gs_list.append(1.0) - - runner = Nvfp4MoE( - num_experts=NUM_EXPERTS, hidden_size=HIDDEN_SIZE, - intermediate_size=INTERMEDIATE_SIZE, max_num_tokens=NUM_TOKENS, - top_k=TOP_K, device=DEVICE, - ) - runner.l1_fp4 = l1_fp4; runner.l1_sf = l1_sf; runner.l1_gs = l1_gs_list - runner.l2_fp4 = l2_fp4; runner.l2_sf = l2_sf; runner.l2_gs = l2_gs_list - runner.set_swiglu_limit(SWIGLU_LIMIT) - - # Warmup - with torch.no_grad(): - runner.compute_activation_global_scales(hidden, topk_weights, topk_ids) - - # Run multiple layers (using same weights, but hidden evolves) - ref_hidden = hidden.clone() - run_hidden = hidden.clone() - - for layer in range(NUM_LAYERS): - with torch.no_grad(): - # Runner - run_hidden_saved = run_hidden.clone() - runner.compute_activation_global_scales(run_hidden, topk_weights, topk_ids) - run_out = runner.run(run_hidden, topk_weights, topk_ids) - run_hidden = run_hidden + run_hidden_saved # Residual connection - - # BF16 reference - ref_hidden_saved = ref_hidden.clone() - ref_out = torch.zeros(NUM_TOKENS, HIDDEN_SIZE, dtype=torch.bfloat16, device=DEVICE) - for i, e in enumerate(expert_indices): - dk = f"layers.{LAYER_IDX}.mlp.experts.{e}.down_proj.weight" - gk = f"layers.{LAYER_IDX}.mlp.experts.{e}.gate_proj.weight" - uk = f"layers.{LAYER_IDX}.mlp.experts.{e}.up_proj.weight" - if dk not in nvfp4_tensors: - continue - gate_bf16 = dequantize_nvfp4_weight(nvfp4_tensors[gk].to(DEVICE), nvfp4_tensors[gk.replace('.weight', '.weight_scale')].to(DEVICE), nvfp4_tensors[gk.replace('.weight', '.weight_scale_2')].item()) - up_bf16 = dequantize_nvfp4_weight(nvfp4_tensors[uk].to(DEVICE), nvfp4_tensors[uk.replace('.weight', '.weight_scale')].to(DEVICE), nvfp4_tensors[uk.replace('.weight', '.weight_scale_2')].item()) - down_bf16 = dequantize_nvfp4_weight(nvfp4_tensors[dk].to(DEVICE), nvfp4_tensors[dk.replace('.weight', '.weight_scale')].to(DEVICE), nvfp4_tensors[dk.replace('.weight', '.weight_scale_2')].item()) - for t in range(NUM_TOKENS): - for k in range(TOP_K): - if topk_ids[t, k].item() != i: - continue - w = topk_weights[t, k].item() - x = ref_hidden[t] - gate = x @ gate_bf16.T - up = x @ up_bf16.T - gate_silu = F.silu(gate).clamp(max=SWIGLU_LIMIT) - up = up.clamp(min=-SWIGLU_LIMIT, max=SWIGLU_LIMIT) - act = gate_silu * up - ref_out[t] += w * (act @ down_bf16.T) - ref_hidden = ref_out + ref_hidden_saved # Residual - - cos_moe = F.cosine_similarity(ref_out.flatten().unsqueeze(0), run_out.flatten().unsqueeze(0)).item() - cos = F.cosine_similarity(ref_hidden.flatten().unsqueeze(0), run_hidden.flatten().unsqueeze(0)).item() - has_nan = torch.isnan(run_hidden).any().item() - has_inf = torch.isinf(run_hidden).any().item() - moe_scale = run_out.abs().mean().item() / max(ref_out.abs().mean().item(), 1e-8) - print(f"Layer {layer}: MoE_cosine={cos_moe:.6f} MoE_scale={moe_scale:.4f} ref_moe_amax={ref_out.amax().item():.4f} run_moe_amax={run_out.amax().item():.4f} NaN={has_nan}") - - if has_nan: - print(f" ❌ NaN detected after layer {layer}! Stopping.") - break - - -if __name__ == "__main__": - main() diff --git a/tests/archive/test_nvfp4_attention_b200.py b/tests/archive/test_nvfp4_attention_b200.py deleted file mode 100644 index 0fed2f2d..00000000 --- a/tests/archive/test_nvfp4_attention_b200.py +++ /dev/null @@ -1,255 +0,0 @@ -#!/usr/bin/env python3 -""" -Test NVFP4 attention: quantize Q and K, GEMM in NVFP4, softmax in BF16. - -Step 1: Verify NVFP4 quantize/dequant roundtrip for attention -Step 2: Q×K^T using CuTeDSL NVFP4 GEMM -Step 3: Softmax + attn×V -Step 4: Full pipeline with real weights, compare to BF16 SDPA - -Usage (on B200): - cd /root/nvfp4-megamoe-kernel - PYTHONPATH=/root/nvfp4-megamoe-kernel tests/venv/bin/python tests/test_nvfp4_attention_b200.py -""" - -import sys, os, json, torch, torch.nn.functional as F, math -from safetensors import safe_open - -REPO = "/root/nvfp4-megamoe-kernel" -sys.path.insert(0, REPO) -MODEL = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4" -DEV = "cuda:0" - -H = 7168; NH = 128; HD = 512; NOPE = 448; ROPE = 64 -QL = 1536; OL = 1024; OG = 16; HPG = NH // OG -EPS = 1e-6; WINDOW = 8192; SCALE = HD ** -0.5 - -E2M1 = torch.tensor([0,.5,1.,1.5,2.,3.,4.,6.,-0,-.5,-1.,-1.5,-2.,-3.,-4.,-6.], dtype=torch.float32) - -_cache = {} -def P(k, wm, md): - if k in _cache: return _cache[k] - with safe_open(os.path.join(md, wm[k]), framework="pt") as f: - t = f.get_tensor(k) - _cache[k] = t - return t - -def dequant(w, sf, gs): - d = w.device; lut = E2M1.to(d) - lo = lut[(w & 0xF).long()]; hi = lut[((w >> 4) & 0xF).long()] - O, I2 = w.shape; I = I2*2 - u = torch.empty(O, I, dtype=torch.float32, device=d) - u[:,0::2] = lo; u[:,1::2] = hi - bs = sf.float().repeat_interleave(16, dim=1)[:O,:I] - return (u * bs * gs).to(torch.bfloat16) - -def rms(x, w, eps=1e-6): - v = x.float().pow(2).mean(-1, keepdim=True) - return (w.float() * (x * torch.rsqrt(v+eps)).float()).to(x.dtype) - -def make_runner(w, sf, gs_t, inf, outf, fused=False, lw=None): - from dsv4.layers.linear import Nvfp4Linear - fp4 = w.view(torch.float4_e2m1fn_x2).permute(1,0).contiguous() - s = sf.to(torch.float8_e4m3fn) if sf.dtype != torch.float8_e4m3fn else sf - s = s.permute(1,0).contiguous() - if fused and gs_t.numel() == 2: - g1,g2 = gs_t[0].item(), gs_t[1].item(); gs = max(g1,g2) - if g1 != g2: - s32 = s.float(); sp = lw[0] if lw else outf//2 - s32[:sp] *= g1/gs; s32[sp:] *= g2/gs; s = s32.to(torch.float8_e4m3fn) - else: - gs = gs_t.max().item() if gs_t.numel() > 1 else gs_t.item() - r = Nvfp4Linear(in_features=inf, out_features=outf, max_num_tokens=8192, device=str(w.device)) - r.fp4 = [fp4]; r.sf = [s]; r.gs = [gs] - r.finalize_weights(); r._ensure_initialized() - return r - -def apply_gptj_rope(x, positions, cos_sin, nope, rope): - if rope == 0 or x.numel() == 0: return x - half = rope // 2 - cos = cos_sin[positions, :half].to(x.dtype) - sin = cos_sin[positions, half:].to(x.dtype) - if x.dim() == 3: cos = cos.unsqueeze(1); sin = sin.unsqueeze(1) - x_rope = x[..., nope:].clone() - even = x_rope[..., 0::2]; odd = x_rope[..., 1::2] - out = x.clone() - out[..., nope:][..., 0::2] = even * cos - odd * sin - out[..., nope:][..., 1::2] = even * sin + odd * cos - return out - -def build_cos_sin(max_pos=4096, rope_dim=ROPE): - half = rope_dim // 2 - inv_freq = 1.0 / (10000.0 ** (torch.arange(0, half, dtype=torch.float32) / half)) - freqs = torch.outer(torch.arange(max_pos, dtype=torch.float32), inv_freq) - return torch.cat([freqs.cos(), freqs.sin()], dim=-1) - - -def bf16_full_attention(q, kv, scale): - """BF16 reference: full self-attention with causal mask.""" - T, NH, HD = q.shape - q_2d = q.reshape(T * NH, HD) - kv_expanded = kv.unsqueeze(1).expand(-1, NH, -1).contiguous() - k_2d = kv_expanded.permute(1, 0, 2).unsqueeze(1).expand(NH, T, T, -1).contiguous().reshape(T * NH, T, HD) - v_2d = k_2d.clone() - scores = torch.matmul(q_2d.unsqueeze(1), k_2d.transpose(-1, -2)) * scale - query_pos = torch.arange(T, device=q.device).unsqueeze(1).repeat(1, NH).reshape(T * NH) - kv_pos = torch.arange(T, device=q.device).unsqueeze(0) - causal = kv_pos <= query_pos.unsqueeze(1) - scores = scores.squeeze(1).masked_fill(~causal, float('-inf')) - weights = F.softmax(scores.float(), dim=-1).to(q.dtype) - out = torch.matmul(weights.unsqueeze(1), v_2d).squeeze(1) - return out.reshape(T, NH, HD) - - -def nvfp4_qk_attention(q, kv, scale): - """NVFP4 attention: quantize Q and K for Q×K^T, then BF16 softmax + attn×V. - - Key insight: Q×K^T is (T*NH, HD) × (HD, T) = (T*NH, T). - This is a standard GEMM that CuTeDSL can handle. - We quantize Q as the "activation" and K^T as the "weight". - """ -from dsv4.ops.quantize import ( - quantize_to_nvfp4, - quantize_activation_nvfp4, -) - from dsv4.layers.linear import Nvfp4Linear - - T, NH, HD = q.shape - device = q.device - - # Q as activation: (T*NH, HD) → NVFP4 - q_2d = q.reshape(T * NH, HD) - q_fp4, q_sf, q_gs = quantize_to_nvfp4(q_2d) # (T*NH, HD//2), (T*NH, HD//16), scalar - - # K as weight: (T, HD) → transpose to (HD, T), quantize as weight - # In our framework, "weight" means quantized along K dim - kv_T = kv.T.contiguous() # (HD, T) - w_fp4, w_sf, w_gs = quantize_to_nvfp4(kv_T) # (HD//2, T), (HD//16, T), scalar - - # Use Nvfp4Linear runner for Q×K^T GEMM - # in_features=HD, out_features=T - # Q is "activation" side, K^T is "weight" side - M = T * NH - K = HD - N = T - - # Create runner for this specific (M, K, N) combination - runner = Nvfp4Linear( - in_features=K, out_features=N, max_num_tokens=M, device=str(device) - ) - - # Weight is kv_T: set up as (N, K//2) in N-major (standard row-major) - # runner expects: weight fp4 is (N, K//2), weight sf is (N, K//16) - # Our w_fp4 from quantize_to_nvfp4(kv_T) is (K//2, T) — that's (K_packed, N) - # Need to transpose to (N, K_packed) - w_fp4_loaded = w_fp4.T.contiguous() # (T, HD//2) = (N, K_packed) - w_sf_loaded = w_sf.T.contiguous() # (T, HD//16) = (N, K_sf) - - runner.fp4 = [w_fp4_loaded] - runner.sf = [w_sf_loaded] - runner.gs = [w_gs] - runner.finalize_weights() - runner._ensure_initialized() - - # Run: Q×K^T - # q_2d is (M, K) BF16, runner produces (M, N) BF16 - scores = runner.run(q_2d) * scale # (T*NH, T) - - # Causal mask - query_pos = torch.arange(T, device=device).unsqueeze(1).repeat(1, NH).reshape(T * NH) - kv_pos = torch.arange(T, device=device).unsqueeze(0) - causal = kv_pos <= query_pos.unsqueeze(1) - scores = scores.masked_fill(~causal, float('-inf')) - - # Softmax in BF16 (must be full precision for numerical stability) - weights = F.softmax(scores.float(), dim=-1).to(q.dtype) # (T*NH, T) - - # attn×V: (T*NH, T) × (T, HD) → (T*NH, HD) - # V = kv (shared, BF16) — no quantization needed here since attn weights are already BF16 - out = torch.matmul(weights, kv) # (T*NH, HD) - - return out.reshape(T, NH, HD) - - -def main(): - torch.cuda.set_device(0) - torch.manual_seed(42) - - print("=" * 70) - print(" NVFP4 Attention Kernel Test") - print("=" * 70) - - with open(os.path.join(MODEL, "model.safetensors.index.json")) as f: - wm = json.load(f)["weight_map"] - G = lambda k: P(k, wm, MODEL).to(DEV) - - p = "model.layers.0"; a = f"{p}.self_attn" - - # Load weights - emb = G("model.embed_tokens.weight") - anorm = G(f"{p}.input_layernorm.weight") - qn = G(f"{a}.q_a_norm.weight"); kvn = G(f"{a}.kv_norm.weight") - woa = G(f"{a}.o_a_proj.weight") - - qa_w = G(f"{a}.q_a_proj.weight"); qa_sf = G(f"{a}.q_a_proj.weight_scale"); qa_gs = G(f"{a}.q_a_proj.weight_scale_2") - qb_w = G(f"{a}.q_b_proj.weight"); qb_sf = G(f"{a}.q_b_proj.weight_scale"); qb_gs = G(f"{a}.q_b_proj.weight_scale_2") - kv_w = G(f"{a}.kv_proj.weight"); kv_sf = G(f"{a}.kv_proj.weight_scale"); kv_gs = G(f"{a}.kv_proj.weight_scale_2") - wob_w = G(f"{a}.o_b_proj.weight"); wob_sf = G(f"{a}.o_b_proj.weight_scale"); wob_gs = G(f"{a}.o_b_proj.weight_scale_2") - sinks = G(f"{a}.sinks") - - # BF16 references - qa_bf16 = dequant(qa_w, qa_sf, qa_gs.item()) - qb_bf16 = dequant(qb_w, qb_sf, qb_gs.item()) - kv_bf16 = dequant(kv_w, kv_sf, kv_gs.item()) - wob_bf16 = dequant(wob_w, wob_sf, wob_gs.item()) - - # CuTeDSL runners - r_qa = make_runner(qa_w, qa_sf, qa_gs, H, qa_w.shape[0]) - r_qb = make_runner(qb_w, qb_sf, qb_gs, QL, qb_w.shape[0]) - r_kv = make_runner(kv_w, kv_sf, kv_gs, H, kv_w.shape[0]) - r_wob = make_runner(wob_w, wob_sf, wob_gs, OG*OL, wob_w.shape[0]) - - # Input - token_ids = torch.tensor([1, 450, 8403, 315, 5413, 374], dtype=torch.long, device=DEV) - NT = len(token_ids) - cos_sin = build_cos_sin(max_pos=WINDOW + 256).to(DEV) - positions = torch.arange(NT, dtype=torch.int64, device=DEV) - - print(f" Input: {NT} tokens, {NH} heads, HD={HD}") - - with torch.no_grad(): - hidden = emb[token_ids] - normed = rms(hidden, anorm, EPS) - - # Projections - qa_cute = r_qa.run(normed) - kv_cute = r_kv.run(normed) - qa_n = rms(qa_cute, qn, EPS) - kv_n = rms(kv_cute, kvn, EPS) - q_cute = r_qb.run(qa_n).view(NT, NH, HD) - q_rope = apply_gptj_rope(q_cute, positions, cos_sin, NOPE, ROPE) - - # ── BF16 reference ──────────────────────────────────────────── - print("\n--- Step 1: BF16 reference attention ---") - o_bf16 = bf16_full_attention(q_rope, kv_n, SCALE) - print(f" BF16 attention output: amax={o_bf16.amax():.4f} NaN={torch.isnan(o_bf16).any()}") - - # ── NVFP4 Q×K^T attention ──────────────────────────────────── - print("\n--- Step 2: NVFP4 Q×K^T attention ---") - try: - o_nvfp4 = nvfp4_qk_attention(q_rope, kv_n, SCALE) - print(f" NVFP4 attention output: amax={o_nvfp4.amax():.4f} NaN={torch.isnan(o_nvfp4).any()}") - - c = F.cosine_similarity(o_nvfp4.flatten().unsqueeze(0).float(), o_bf16.flatten().unsqueeze(0).float()).item() - print(f" NVFP4 vs BF16 cosine: {c:.6f} {'✅' if c>=0.98 else '❌'}") - except Exception as e: - print(f" ERROR: {e}") - import traceback; traceback.print_exc() - - print("\n" + "=" * 70) - print(" Done") - print("=" * 70) - - -if __name__ == "__main__": - main() diff --git a/tests/archive/test_nvfp4_attn_gemm_b200.py b/tests/archive/test_nvfp4_attn_gemm_b200.py deleted file mode 100644 index a47b6263..00000000 --- a/tests/archive/test_nvfp4_attn_gemm_b200.py +++ /dev/null @@ -1,373 +0,0 @@ -#!/usr/bin/env python3 -""" -CuTeDSL NVFP4 Attention Kernel — Q×K^T GEMM - -DeepSeek-V4 attention is CSA/HCA (NOT MLA): -- KV latent: (T, 512) shared across all 128 heads -- Q: (T, 128, 512) — 128 heads, 512 dim each -- Q×K^T: (T, 128, 512) × (512, T) → (T, 128, T) per head -- softmax → (T, 128, T) -- attn×V: (T, 128, T) × (T, 512) → (T, 128, 512) - -The Q×K^T step is the expensive one. For T=8192 tokens, NH=128: -- M = T*NH = 1,048,576 -- K = HD = 512 -- N = T = 8192 -- FLOPs: 2 * M * K * N ≈ 8.8 TFLOPS - -NVFP4 quantization cuts the data movement by 4x (BF16→FP4). - -This test: -1. Build a CuTeDSL NVFP4 GEMM runner for Q×K^T -2. Compare output against BF16 reference -3. Test with real model weights (full attention pipeline) - -Usage (on B200): - cd /root/nvfp4-megamoe-kernel - PYTHONPATH=/root/nvfp4-megamoe-kernel tests/venv/bin/python tests/test_nvfp4_attn_gemm_b200.py -""" - -import sys, os, json, torch, torch.nn.functional as F, math, time -from safetensors import safe_open - -REPO = "/root/nvfp4-megamoe-kernel" -sys.path.insert(0, REPO) -MODEL = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4" -DEV = "cuda:0" - -# Model config -H = 7168; NH = 128; HD = 512; NOPE = 448; ROPE = 64 -QL = 1536; OL = 1024; OG = 16; HPG = NH // OG -EPS = 1e-6; WINDOW = 128; SCALE = HD ** -0.5 - -E2M1 = torch.tensor([0,.5,1.,1.5,2.,3.,4.,6.,-0,-.5,-1.,-1.5,-2.,-3.,-4.,-6.], dtype=torch.float32) - -_cache = {} -def P(k, wm, md): - if k in _cache: return _cache[k] - with safe_open(os.path.join(md, wm[k]), framework="pt") as f: - t = f.get_tensor(k) - _cache[k] = t - return t - -def dequant(w, sf, gs): - d = w.device; lut = E2M1.to(d) - lo = lut[(w & 0xF).long()]; hi = lut[((w >> 4) & 0xF).long()] - O, I2 = w.shape; I = I2*2 - u = torch.empty(O, I, dtype=torch.float32, device=d) - u[:,0::2] = lo; u[:,1::2] = hi - bs = sf.float().repeat_interleave(16, dim=1)[:O,:I] - return (u * bs * gs).to(torch.bfloat16) - -def rms(x, w, eps=1e-6): - v = x.float().pow(2).mean(-1, keepdim=True) - return (w.float() * (x * torch.rsqrt(v+eps)).float()).to(x.dtype) - -def make_runner(w, sf, gs_t, inf, outf, fused=False, lw=None): - from dsv4.layers.linear import Nvfp4Linear - fp4 = w.view(torch.float4_e2m1fn_x2).permute(1,0).contiguous() - s = sf.to(torch.float8_e4m3fn) if sf.dtype != torch.float8_e4m3fn else sf - s = s.permute(1,0).contiguous() - if fused and gs_t.numel() == 2: - g1,g2 = gs_t[0].item(), gs_t[1].item(); gs = max(g1,g2) - if g1 != g2: - s32 = s.float(); sp = lw[0] if lw else outf//2 - s32[:sp] *= g1/gs; s32[sp:] *= g2/gs; s = s32.to(torch.float8_e4m3fn) - else: - gs = gs_t.max().item() if gs_t.numel() > 1 else gs_t.item() - r = Nvfp4Linear(in_features=inf, out_features=outf, max_num_tokens=8192, device=str(w.device)) - r.fp4 = [fp4]; r.sf = [s]; r.gs = [gs] - r.finalize_weights(); r._ensure_initialized() - return r - -def apply_gptj_rope(x, positions, cos_sin, nope, rope): - if rope == 0 or x.numel() == 0: return x - half = rope // 2 - cos = cos_sin[positions, :half].to(x.dtype) - sin = cos_sin[positions, half:2*half].to(x.dtype) - if x.dim() == 3: cos = cos.unsqueeze(1); sin = sin.unsqueeze(1) - x_rope = x[..., nope:].clone() - even = x_rope[..., 0::2]; odd = x_rope[..., 1::2] - out = x.clone() - out[..., nope:][..., 0::2] = even * cos - odd * sin - out[..., nope:][..., 1::2] = even * sin + odd * cos - return out - -def apply_inv_gptj_rope(x, positions, cos_sin, nope, rope): - if rope == 0 or x.numel() == 0: return x - half = rope // 2 - cos = cos_sin[positions, :half].to(x.dtype) - sin = cos_sin[positions, half:2*half].to(x.dtype) - if x.dim() == 3: cos = cos.unsqueeze(1); sin = sin.unsqueeze(1) - x_rope = x[..., nope:].clone() - even = x_rope[..., 0::2]; odd = x_rope[..., 1::2] - out = x.clone() - out[..., nope:][..., 0::2] = even * cos + odd * sin - out[..., nope:][..., 1::2] = -even * sin + odd * cos - return out - -def build_cos_sin(max_pos=4096, rope_dim=ROPE): - half = rope_dim // 2 - inv_freq = 1.0 / (10000.0 ** (torch.arange(0, half, dtype=torch.float32) / half)) - freqs = torch.outer(torch.arange(max_pos, dtype=torch.float32), inv_freq) - return torch.cat([freqs.cos(), freqs.sin()], dim=-1) - - -def bf16_causal_attention(q, kv, scale): - """BF16 reference: full causal self-attention.""" - T, NH, HD = q.shape - q_2d = q.reshape(T * NH, HD) - kv_exp = kv.unsqueeze(1).expand(-1, NH, -1).contiguous() - k_2d = kv_exp.permute(1, 0, 2).unsqueeze(1).expand(NH, T, T, -1).contiguous().reshape(T * NH, T, HD) - v_2d = k_2d.clone() - scores = torch.matmul(q_2d.unsqueeze(1), k_2d.transpose(-1, -2)) * scale - query_pos = torch.arange(T, device=q.device).unsqueeze(1).repeat(1, NH).reshape(T * NH) - kv_pos = torch.arange(T, device=q.device).unsqueeze(0) - causal = kv_pos <= query_pos.unsqueeze(1) - scores = scores.squeeze(1).masked_fill(~causal, float('-inf')) - weights = F.softmax(scores.float(), dim=-1).to(q.dtype) - out = torch.matmul(weights.unsqueeze(1), v_2d).squeeze(1) - return out.reshape(T, NH, HD) - - -class NVFP4Attention: - """CuTeDSL NVFP4 attention kernel. - - Q×K^T via NVFP4 GEMM, softmax in BF16, attn×V in BF16. - - The Q×K^T GEMM: (T*NH, HD) × (HD, T) → (T*NH, T) - - Q is the "activation": quantized per-row (dynamic) - - K^T is the "weight": quantized from (T, HD) KV latent - - For decode (M=1 per head), the GEMM is tiny — NVFP4 overhead isn't worth it. - For prefill (M=chunk_size), the GEMM is large — NVFP4 saves 4x memory bandwidth. - - This kernel targets the prefill case where T is large. - """ - - def __init__(self, head_dim: int, num_heads: int, max_seq_len: int, device: str = "cuda"): - self.head_dim = head_dim - self.num_heads = num_heads - self.max_seq_len = max_seq_len - self.device = device - self._runner = None # Compiled on first call - - def forward(self, q_bf16, kv_bf16, scale): - """Forward pass. - - Args: - q_bf16: (T, NH, HD) with RoPE applied - kv_bf16: (T, HD) shared KV latent (BF16) - scale: 1/sqrt(HD) - - Returns: - (T, NH, HD) attention output - """ - from dsv4.layers.linear import Nvfp4Linear - - T, NH, HD = q_bf16.shape - device = q_bf16.device - - # Reshape Q: (T, NH, HD) → (T*NH, HD) — treat as 2D for GEMM - q_2d = q_bf16.reshape(T * NH, HD) - - # ── Q×K^T via NVFP4 GEMM ──────────────────────────────────── - # Q is "activation" (T*NH, HD), K^T is "weight" (T, HD) - # GEMM: (T*NH, HD) × (HD, T) → (T*NH, T) - # - # We use Nvfp4Linear with in_features=HD, out_features=T - # Q is the "hidden_states", K (kv) is the "weight" matrix - - # Create or get cached runner - cache_key = (T, HD, NH) - if self._runner is None or getattr(self, '_cache_key', None) != cache_key: - runner = Nvfp4Linear( - in_features=HD, - out_features=T, - max_num_tokens=T * NH, - device=str(device), - ) - - # Set K as the weight: kv (T, HD) → treat as weight (N=T, K=HD) - # quantize_to_nvfp4 quantizes along last dim (D=HD) as activation - # For weight, we need (K, N) layout — but kv is (T, HD) = (N, K) - # Nvfp4Linear expects weight in (N, K//2) after permute - -from dsv4.ops.quantize import ( - quantize_to_nvfp4, -) - # Quantize KV as a 2D tensor: (T, HD) - # quantize_to_nvfp4 works on last dim (D=HD), returns: - # (T, HD//2) fp4, (T, HD//16) sf, scalar gs - kv_fp4, kv_sf, kv_gs = quantize_to_nvfp4(kv_bf16) - - # For Nvfp4Linear, weight is (N, K_packed) = (T, HD//2) - # Our kv_fp4 is already (T, HD//2) — perfect! - # sf needs to be (N, K_sf) = (T, HD//16) — already correct - - w_fp4 = kv_fp4 # (T, HD//2) — already in row-major (N, K_packed) - w_sf = kv_sf # (T, HD//16) - - # Set up the runner with K^T as weight - # The runner expects fp4 as list of (N, K_packed), sf as list of (N, K_sf) - # after finalize_weights, it does permute(1,0) internally - runner.fp4 = [w_fp4] - runner.sf = [w_sf] - runner.gs = [kv_gs] - runner.finalize_weights() - runner._ensure_initialized() - - self._runner = runner - self._cache_key = cache_key - - # Run Q×K^T GEMM - scores = self._runner.run(q_2d) # (T*NH, N_padded) - scores = scores[:, :T] # Slice to actual N=T (runner pads to 128) - scores = scores * scale - - # Causal mask - query_pos = torch.arange(T, device=device).unsqueeze(1).repeat(1, NH).reshape(T * NH) - kv_pos = torch.arange(T, device=device).unsqueeze(0) - causal = kv_pos <= query_pos.unsqueeze(1) - scores = scores.masked_fill(~causal, float('-inf')) - - # Softmax (BF16 for numerical stability, actually float32) - weights = F.softmax(scores.float(), dim=-1).to(q_bf16.dtype) - - # attn×V: (T*NH, T) × (T, HD) → (T*NH, HD) - # V = K = kv (shared latent) — BF16, no quantization - out = torch.matmul(weights, kv_bf16) - - return out.reshape(T, NH, HD) - - -def main(): - torch.cuda.set_device(0) - torch.manual_seed(42) - - print("=" * 70) - print(" CuTeDSL NVFP4 Attention Kernel Test") - print(" Q×K^T via NVFP4 GEMM, softmax BF16, attn×V BF16") - print("=" * 70) - - # ── Step 1: Synthetic test with random data ────────────────────── - print("\n--- Step 1: Synthetic random test ---") - T = 8 - q_rand = torch.randn(T, NH, HD, dtype=torch.bfloat16, device=DEV) - kv_rand = torch.randn(T, HD, dtype=torch.bfloat16, device=DEV) - - with torch.no_grad(): - ref = bf16_causal_attention(q_rand, kv_rand, SCALE) - print(f" BF16 reference: amax={ref.amax():.4f}") - - kernel = NVFP4Attention(HD, NH, max_seq_len=8192, device=DEV) - out = kernel.forward(q_rand, kv_rand, SCALE) - print(f" NVFP4 kernel: amax={out.amax():.4f}") - - c = F.cosine_similarity(ref.flatten().unsqueeze(0).float(), out.flatten().unsqueeze(0).float()).item() - print(f" Cosine: {c:.6f} {'✅' if c>=0.95 else '❌'}") - - # ── Step 2: Real model weights, full attention pipeline ────────── - print("\n--- Step 2: Real model weights (layer 0, C128A) ---") - with open(os.path.join(MODEL, "model.safetensors.index.json")) as f: - wm = json.load(f)["weight_map"] - G = lambda k: P(k, wm, MODEL).to(DEV) - - p = "model.layers.0"; a = f"{p}.self_attn" - - emb = G("model.embed_tokens.weight") - anorm = G(f"{p}.input_layernorm.weight") - qn = G(f"{a}.q_a_norm.weight"); kvn = G(f"{a}.kv_norm.weight") - woa = G(f"{a}.o_a_proj.weight") - - qa_w = G(f"{a}.q_a_proj.weight"); qa_sf = G(f"{a}.q_a_proj.weight_scale"); qa_gs = G(f"{a}.q_a_proj.weight_scale_2") - qb_w = G(f"{a}.q_b_proj.weight"); qb_sf = G(f"{a}.q_b_proj.weight_scale"); qb_gs = G(f"{a}.q_b_proj.weight_scale_2") - kv_w = G(f"{a}.kv_proj.weight"); kv_sf = G(f"{a}.kv_proj.weight_scale"); kv_gs = G(f"{a}.kv_proj.weight_scale_2") - wob_w = G(f"{a}.o_b_proj.weight"); wob_sf = G(f"{a}.o_b_proj.weight_scale"); wob_gs = G(f"{a}.o_b_proj.weight_scale_2") - - qa_bf16 = dequant(qa_w, qa_sf, qa_gs.item()) - qb_bf16 = dequant(qb_w, qb_sf, qb_gs.item()) - kv_bf16 = dequant(kv_w, kv_sf, kv_gs.item()) - wob_bf16 = dequant(wob_w, wob_sf, wob_gs.item()) - - r_qa = make_runner(qa_w, qa_sf, qa_gs, H, qa_w.shape[0]) - r_qb = make_runner(qb_w, qb_sf, qb_gs, QL, qb_w.shape[0]) - r_kv = make_runner(kv_w, kv_sf, kv_gs, H, kv_w.shape[0]) - r_wob = make_runner(wob_w, wob_sf, wob_gs, OG*OL, wob_w.shape[0]) - - token_ids = torch.tensor([1, 450, 8403, 315, 5413, 374], dtype=torch.long, device=DEV) - NT = len(token_ids) - cos_sin = build_cos_sin(max_pos=WINDOW + 256).to(DEV) - positions = torch.arange(NT, dtype=torch.int64, device=DEV) - - with torch.no_grad(): - hidden = emb[token_ids] - normed = rms(hidden, anorm, EPS) - - # Projections (CuTeDSL) - qa_cute = r_qa.run(normed) - kv_cute = r_kv.run(normed) - qa_n = rms(qa_cute, qn, EPS) - kv_n = rms(kv_cute, kvn, EPS) - q_cute = r_qb.run(qa_n).view(NT, NH, HD) - q_rope = apply_gptj_rope(q_cute, positions, cos_sin, NOPE, ROPE) - - # ── NVFP4 Attention ────────────────────────────────────── - attn_kernel = NVFP4Attention(HD, NH, max_seq_len=8192, device=DEV) - o_nvfp4 = attn_kernel.forward(q_rope, kv_n, SCALE) - print(f" NVFP4 attention: amax={o_nvfp4.amax():.4f}") - - # ── BF16 reference ─────────────────────────────────────── - o_bf16 = bf16_causal_attention(q_rope, kv_n, SCALE) - print(f" BF16 attention: amax={o_bf16.amax():.4f}") - - c = F.cosine_similarity(o_nvfp4.flatten().unsqueeze(0).float(), o_bf16.flatten().unsqueeze(0).float()).item() - print(f" NVFP4 vs BF16 cosine: {c:.6f} {'✅' if c>=0.95 else '❌'}") - - # ── Full pipeline: attention → o_a → o_b ───────────────── - o_inv = apply_inv_gptj_rope(o_nvfp4, positions, cos_sin, NOPE, ROPE) - o_grouped = o_inv.view(NT, OG, HPG * HD).permute(1, 0, 2) - woa_3d = woa.view(OG, OL, HPG * HD) - z = torch.bmm(o_grouped, woa_3d.transpose(1, 2)).permute(1, 0, 2).reshape(NT, OG * OL) - attn_out = r_wob.run(z) - - # BF16 reference pipeline - o_inv_bf = apply_inv_gptj_rope(o_bf16, positions, cos_sin, NOPE, ROPE) - o_grouped_bf = o_inv_bf.view(NT, OG, HPG * HD).permute(1, 0, 2) - z_bf = torch.bmm(o_grouped_bf, woa_3d.transpose(1, 2)).permute(1, 0, 2).reshape(NT, OG * OL) - attn_bf = z_bf @ wob_bf16.T - - c_full = F.cosine_similarity(attn_out.flatten().unsqueeze(0).float(), attn_bf.flatten().unsqueeze(0).float()).item() - print(f" Full pipeline cosine: {c_full:.6f} {'✅' if c_full>=0.95 else '❌'}") - - # Logits - fnorm_w = G("model.norm.weight") - lm_head = G("lm_head.weight") - x = hidden + attn_out - x_n = rms(x, fnorm_w, EPS) - logits = x_n @ lm_head.T - log_std = logits[-1].float().std().item() - print(f" logits: std={log_std:.4f} {'✅' if 0.5 < log_std < 50 else '❌'}") - - # ── Step 3: Larger sequence test ───────────────────────────────── - print("\n--- Step 3: Larger sequence (T=64) ---") - torch.cuda.empty_cache() - T64 = 64 - with torch.no_grad(): - q64 = torch.randn(T64, NH, HD, dtype=torch.bfloat16, device=DEV) - kv64 = torch.randn(T64, HD, dtype=torch.bfloat16, device=DEV) - - ref64 = bf16_causal_attention(q64, kv64, SCALE) - kernel64 = NVFP4Attention(HD, NH, max_seq_len=8192, device=DEV) - out64 = kernel64.forward(q64, kv64, SCALE) - - c64 = F.cosine_similarity(ref64.flatten().unsqueeze(0).float(), out64.flatten().unsqueeze(0).float()).item() - print(f" T=64 NVFP4 vs BF16 cosine: {c64:.6f} {'✅' if c64>=0.95 else '❌'}") - - print(f"\n{'='*70}") - print(f" DONE") - print(f"{'='*70}") - - -if __name__ == "__main__": - main() diff --git a/tests/archive/test_nvfp4_mapper.py b/tests/archive/test_nvfp4_mapper.py deleted file mode 100644 index b7610587..00000000 --- a/tests/archive/test_nvfp4_mapper.py +++ /dev/null @@ -1,186 +0,0 @@ -"""Test that the NVFP4 weight mapper correctly maps checkpoint keys to model parameter names.""" -import re -import sys - -# Inline WeightsMapper (from vllm.model_executor.models.utils) -from dataclasses import dataclass, field -from typing import Any, Iterable, Mapping - -@dataclass -class WeightsMapper: - orig_to_new_regex: Mapping[re.Pattern, str | None] = field(default_factory=dict) - orig_to_new_substr: Mapping[str, str | None] = field(default_factory=dict) - orig_to_new_prefix: Mapping[str, str | None] = field(default_factory=dict) - orig_to_new_suffix: Mapping[str, str | None] = field(default_factory=dict) - - def _map_name(self, key: str) -> str | None: - for pattern, new_key in self.orig_to_new_regex.items(): - if pattern.search(key): - if new_key is None: - return None - key = pattern.sub(new_key, key) - for substr, new_key in self.orig_to_new_substr.items(): - if substr in key: - if new_key is None: - return None - key = key.replace(substr, new_key, 1) - for prefix, new_key in self.orig_to_new_prefix.items(): - if key.startswith(prefix): - if new_key is None: - return None - key = key.replace(prefix, new_key, 1) - for suffix, new_key in self.orig_to_new_suffix.items(): - if key.endswith(suffix): - if new_key is None: - return None - key = new_key.join(key.rsplit(suffix, 1)) - return key - - -def make_nvfp4_mapper() -> WeightsMapper: - expert_rename_regex = { - re.compile(r"(\.experts\.\d+\.)gate_proj\."): r"\1w1.", - re.compile(r"(\.experts\.\d+\.)up_proj\."): r"\1w3.", - re.compile(r"(\.experts\.\d+\.)down_proj\."): r"\1w2.", - } - return WeightsMapper( - orig_to_new_prefix={ - "layers.": "model.layers.", - "embed.": "model.embed.", - "norm.": "model.norm.", - # hc_head NOT mapped — checkpoint already has model.hc_head.* - # and model params are flat (hc_head_fn, not hc_head.fn) - "mtp.": "model.mtp.", - }, - orig_to_new_regex=expert_rename_regex, - orig_to_new_suffix={ - "head.weight": "lm_head.weight", - "embed.weight": "embed_tokens.weight", - ".ffn_norm.weight": ".ffn.norm_gate.norm.weight", - ".ffn.gate.weight": ".ffn.norm_gate.gate.weight", - ".ffn.gate.bias": ".ffn.norm_gate.e_score_correction_bias", - ".ffn.gate.tid2eid": ".ffn.norm_gate.tid2eid", - }, - orig_to_new_substr={ - ".self_attn.compressor.indexer.q_b_proj.": ".attn.indexer.wq_b.", - ".self_attn.compressor.indexer.weights_proj.": ".attn.indexer.weights_proj.", - ".self_attn.compressor.indexer.kv_norm.": ".attn.indexer.k_norm.", - ".self_attn.compressor.indexer.kv_proj.": ".attn.indexer.compressor.wkv.", - ".self_attn.compressor.indexer.gate_proj.": ".attn.indexer.compressor.wgate.", - ".self_attn.compressor.indexer.position_bias": ".attn.indexer.compressor.ape", - "compressor.kv_proj.": "compressor.wkv.", - "compressor.gate_proj.": "compressor.wgate.", - "compressor.kv_norm.": "compressor.norm.", - "compressor.position_bias": "compressor.ape", - ".self_attn.compressor.": ".attn.mla_attn.compressor.", - ".self_attn.q_a_proj.": ".attn.wq_a.", - ".self_attn.kv_proj.": ".attn.wkv.", - ".self_attn.q_b_proj.": ".attn.wq_b.", - ".self_attn.o_a_proj.": ".attn.wo_a.", - ".self_attn.o_b_proj.": ".attn.wo_b.", - ".self_attn.q_a_norm.": ".attn.q_norm.", - ".self_attn.kv_norm.": ".attn.kv_norm.", - ".self_attn.sinks": ".attn.attn_sink", - ".mlp.shared_experts.gate_proj.": ".ffn.shared_experts.w1.", - ".mlp.shared_experts.up_proj.": ".ffn.shared_experts.w3.", - ".mlp.shared_experts.down_proj.": ".ffn.shared_experts.down_proj.", - ".mlp.": ".ffn.", - ".self_attn.": ".attn.", - "input_layernorm.": "attn_norm.", - "post_attention_layernorm.": "ffn_norm.", - ".attn_hc.fn": ".hc_attn_fn", - ".attn_hc.base": ".hc_attn_base", - ".attn_hc.scale": ".hc_attn_scale", - ".ffn_hc.fn": ".hc_ffn_fn", - ".ffn_hc.base": ".hc_ffn_base", - ".ffn_hc.scale": ".hc_ffn_scale", - "hc_head.hc_fn": "hc_head_fn", - "hc_head.hc_base": "hc_head_base", - "hc_head.hc_scale": "hc_head_scale", - }, - ) - - -def test_mapper(): - mapper = make_nvfp4_mapper() - - test_cases = [ - # Attention projections - ("layers.0.self_attn.q_a_proj.weight", "model.layers.0.attn.wq_a.weight"), - ("layers.0.self_attn.q_a_proj.weight_scale", "model.layers.0.attn.wq_a.weight_scale"), - ("layers.0.self_attn.q_a_proj.weight_scale_2", "model.layers.0.attn.wq_a.weight_scale_2"), - ("layers.0.self_attn.q_a_proj.input_scale", "model.layers.0.attn.wq_a.input_scale"), - ("layers.0.self_attn.kv_proj.weight", "model.layers.0.attn.wkv.weight"), - ("layers.0.self_attn.q_b_proj.weight", "model.layers.0.attn.wq_b.weight"), - ("layers.0.self_attn.o_a_proj.weight", "model.layers.0.attn.wo_a.weight"), - ("layers.0.self_attn.o_b_proj.weight", "model.layers.0.attn.wo_b.weight"), - ("layers.0.self_attn.o_b_proj.weight_scale", "model.layers.0.attn.wo_b.weight_scale"), - ("layers.0.self_attn.q_a_norm.weight", "model.layers.0.attn.q_norm.weight"), - ("layers.0.self_attn.kv_norm.weight", "model.layers.0.attn.kv_norm.weight"), - ("layers.0.self_attn.sinks", "model.layers.0.attn.attn_sink"), - - # Compressor (non-indexer) - ("layers.0.self_attn.compressor.kv_proj.weight", "model.layers.0.attn.mla_attn.compressor.wkv.weight"), - ("layers.0.self_attn.compressor.kv_proj.input_scale", "model.layers.0.attn.mla_attn.compressor.wkv.input_scale"), - ("layers.0.self_attn.compressor.gate_proj.weight", "model.layers.0.attn.mla_attn.compressor.wgate.weight"), - ("layers.0.self_attn.compressor.kv_norm.weight", "model.layers.0.attn.mla_attn.compressor.norm.weight"), - ("layers.0.self_attn.compressor.position_bias", "model.layers.0.attn.mla_attn.compressor.ape"), - - # Indexer - ("layers.2.self_attn.compressor.indexer.q_b_proj.weight", "model.layers.2.attn.indexer.wq_b.weight"), - ("layers.2.self_attn.compressor.indexer.weights_proj.weight", "model.layers.2.attn.indexer.weights_proj.weight"), - ("layers.2.self_attn.compressor.indexer.kv_norm.weight", "model.layers.2.attn.indexer.k_norm.weight"), - ("layers.2.self_attn.compressor.indexer.kv_proj.weight", "model.layers.2.attn.indexer.compressor.wkv.weight"), - ("layers.2.self_attn.compressor.indexer.gate_proj.weight", "model.layers.2.attn.indexer.compressor.wgate.weight"), - ("layers.2.self_attn.compressor.indexer.position_bias", "model.layers.2.attn.indexer.compressor.ape"), - - # Expert weights - ("layers.0.mlp.experts.0.gate_proj.weight", "model.layers.0.ffn.experts.0.w1.weight"), - ("layers.0.mlp.experts.0.gate_proj.weight_scale", "model.layers.0.ffn.experts.0.w1.weight_scale"), - ("layers.0.mlp.experts.0.gate_proj.weight_scale_2", "model.layers.0.ffn.experts.0.w1.weight_scale_2"), - ("layers.0.mlp.experts.0.gate_proj.input_scale", "model.layers.0.ffn.experts.0.w1.input_scale"), - ("layers.0.mlp.experts.0.up_proj.weight", "model.layers.0.ffn.experts.0.w3.weight"), - ("layers.0.mlp.experts.0.down_proj.weight", "model.layers.0.ffn.experts.0.w2.weight"), - - # Shared experts - ("layers.0.mlp.shared_experts.gate_proj.weight", "model.layers.0.ffn.shared_experts.w1.weight"), - ("layers.0.mlp.shared_experts.up_proj.weight", "model.layers.0.ffn.shared_experts.w3.weight"), - ("layers.0.mlp.shared_experts.down_proj.weight", "model.layers.0.ffn.shared_experts.down_proj.weight"), - - # MoE gate + norm - ("layers.0.mlp.gate.weight", "model.layers.0.ffn.norm_gate.gate.weight"), - ("layers.0.mlp.gate.tid2eid", "model.layers.0.ffn.norm_gate.tid2eid"), - ("layers.0.input_layernorm.weight", "model.layers.0.attn_norm.weight"), - ("layers.0.post_attention_layernorm.weight", "model.layers.0.ffn.norm_gate.norm.weight"), - - # HC params - ("layers.0.attn_hc.fn", "model.layers.0.hc_attn_fn"), - ("layers.0.ffn_hc.scale", "model.layers.0.hc_ffn_scale"), - - # HC head (checkpoint has model.hc_head.hc_fn, model params are flat hc_head_fn) - ("hc_head.hc_fn", "hc_head_fn"), - - # MTP (already uses ffn prefix in checkpoint) - ("mtp.0.ffn.experts.0.w1.weight", "model.mtp.0.ffn.experts.0.w1.weight"), - ("mtp.0.ffn_norm.weight", "model.mtp.0.ffn.norm_gate.norm.weight"), - ] - - passed = 0 - failed = 0 - for ckpt_key, expected in test_cases: - result = mapper._map_name(ckpt_key) - if result == expected: - passed += 1 - else: - print(f"FAIL: {ckpt_key}") - print(f" Expected: {expected}") - print(f" Got: {result}") - failed += 1 - - print(f"\n{passed}/{passed+failed} tests passed") - if failed > 0: - sys.exit(1) - - -if __name__ == "__main__": - test_mapper() diff --git a/tests/archive/test_o_projection.py b/tests/archive/test_o_projection.py deleted file mode 100644 index ada2bed8..00000000 --- a/tests/archive/test_o_projection.py +++ /dev/null @@ -1,159 +0,0 @@ -"""Test BF16 inverse RoPE + wo_a BMM (no GPU needed). - -Validates the O projection path we patched into the attention forward. -""" - -import torch -import math - - -def apply_inv_rope_bf16( - o: torch.Tensor, - positions: torch.Tensor, - cos_sin_cache: torch.Tensor, - nope_dim: int = 448, - rope_dim: int = 64, -) -> torch.Tensor: - """Same as the patched version in deepseek_v4_attention.py.""" - if rope_dim == 0 or o.numel() == 0: - return o - half_rope = rope_dim // 2 - - cos_all = cos_sin_cache[positions, :half_rope].unsqueeze(1).to(o.dtype) - sin_all = cos_sin_cache[positions, half_rope:].unsqueeze(1).to(o.dtype) - - o_rope = o[:, :, nope_dim:] - o_even = o_rope[:, :, 0::2] - o_odd = o_rope[:, :, 1::2] - - inv_even = o_even * cos_all + o_odd * sin_all - inv_odd = -o_even * sin_all + o_odd * cos_all - - result = o.clone() - result[:, :, nope_dim:][:, :, 0::2] = inv_even - result[:, :, nope_dim:][:, :, 1::2] = inv_odd - return result - - -def apply_gptj_rope( - x: torch.Tensor, - positions: torch.Tensor, - cos_sin_cache: torch.Tensor, - nope_dim: int = 448, - rope_dim: int = 64, -) -> torch.Tensor: - """Apply forward GPT-J style RoPE (for testing roundtrip).""" - half_rope = rope_dim // 2 - cos_all = cos_sin_cache[positions, :half_rope].unsqueeze(1).to(x.dtype) - sin_all = cos_sin_cache[positions, half_rope:].unsqueeze(1).to(x.dtype) - - x_rope = x[:, :, nope_dim:] - x_even = x_rope[:, :, 0::2] - x_odd = x_rope[:, :, 1::2] - - rot_even = x_even * cos_all - x_odd * sin_all - rot_odd = x_even * sin_all + x_odd * cos_all - - result = x.clone() - result[:, :, nope_dim:][:, :, 0::2] = rot_even - result[:, :, nope_dim:][:, :, 1::2] = rot_odd - return result - - -def test_inv_rope_roundtrip(): - """inv_rope(forward_rope(x)) should recover x.""" - torch.manual_seed(42) - T, H, D = 4, 8, 512 # tokens, heads, head_dim - nope_dim, rope_dim = 448, 64 - max_pos = 100 - - # Build cos_sin_cache for positions 0..max_pos - inv_freq = 1.0 / (10000.0 ** (torch.arange(0, rope_dim, 2).float() / rope_dim)) - t = torch.arange(max_pos, dtype=torch.float32) - freqs = torch.einsum("i,j -> ij", t, inv_freq) # (max_pos, half_rope) - cos_sin_cache = torch.cat([freqs.cos(), freqs.sin()], dim=-1) # (max_pos, rope_dim) - - x = torch.randn(T, H, D, dtype=torch.bfloat16) * 0.1 - positions = torch.tensor([0, 5, 10, 50], dtype=torch.int64) - - # Apply forward RoPE, then inverse - rotated = apply_gptj_rope(x, positions, cos_sin_cache, nope_dim, rope_dim) - recovered = apply_inv_rope_bf16(rotated, positions, cos_sin_cache, nope_dim, rope_dim) - - # NoPE portion unchanged - nope_diff = (recovered[:, :, :nope_dim] - x[:, :, :nope_dim]).abs().max().item() - assert nope_diff == 0, f"NoPE should be unchanged, max diff: {nope_diff}" - - # RoPE portion should roundtrip within BF16 precision - rope_diff = (recovered[:, :, nope_dim:] - x[:, :, nope_dim:]).abs().max().item() - assert rope_diff < 0.02, f"RoPE roundtrip error too high: {rope_diff}" - print(f"✅ inv_rope roundtrip: NoPE diff={nope_diff}, RoPE diff={rope_diff:.6f}") - - -def test_wo_a_bmm(): - """wo_a BMM should match einsum 'tgd,grd->tgr'.""" - torch.manual_seed(42) - T = 3 - n_local_groups = 4 - heads_per_group = 2 - head_dim = 512 - o_lora_rank = 128 - n_local_heads = n_local_groups * heads_per_group - - # wo_a weight: (n_groups * o_lora_rank, heads_per_group * head_dim) - wo_a_weight = torch.randn(n_local_groups * o_lora_rank, heads_per_group * head_dim, dtype=torch.bfloat16) - - # Attention output (after inv RoPE): (T, n_local_heads, head_dim) - o_inv = torch.randn(T, n_local_heads, head_dim, dtype=torch.bfloat16) - - # BMM path (our implementation) - hidden_dim = heads_per_group * head_dim - o_grouped = o_inv.view(T, n_local_groups, hidden_dim) - wo_a_w = wo_a_weight.view(n_local_groups, o_lora_rank, hidden_dim) - z_bmm = torch.bmm( - o_grouped.permute(1, 0, 2), - wo_a_w.transpose(1, 2), - ).permute(1, 0, 2) - - # Reference: einsum - o_for_einsum = o_inv.view(T, n_local_groups, hidden_dim).float() - wo_a_for_einsum = wo_a_w.float() - z_einsum = torch.einsum("tgd,grd->tgr", o_for_einsum, wo_a_for_einsum).bfloat16() - - diff = (z_bmm - z_einsum).abs().max().item() - assert diff < 0.01, f"wo_a BMM vs einsum diff: {diff}" - print(f"✅ wo_a BMM matches einsum: max diff={diff:.6f}") - - -def test_inv_rope_at_zero(): - """At position 0, cos=1, sin=0, so inv_rope should be identity on RoPE dims.""" - torch.manual_seed(42) - T, H, D = 2, 4, 512 - nope_dim, rope_dim = 448, 64 - - inv_freq = 1.0 / (10000.0 ** (torch.arange(0, rope_dim, 2).float() / rope_dim)) - t = torch.arange(10, dtype=torch.float32) - freqs = torch.einsum("i,j -> ij", t, inv_freq) - cos_sin_cache = torch.cat([freqs.cos(), freqs.sin()], dim=-1) # (10, rope_dim) - # At pos 0, cos=1, sin=0 - - x = torch.randn(T, H, D, dtype=torch.bfloat16) * 0.1 - positions = torch.zeros(T, dtype=torch.int64) - - # Forward RoPE at pos 0 should be identity (cos=1, sin=0) - rotated = apply_gptj_rope(x, positions, cos_sin_cache, nope_dim, rope_dim) - diff = (rotated - x).abs().max().item() - assert diff < 1e-5, f"RoPE at pos=0 should be identity, diff={diff}" - - # Inverse RoPE on unrotated input at pos 0 should also be identity - inv = apply_inv_rope_bf16(x, positions, cos_sin_cache, nope_dim, rope_dim) - diff2 = (inv - x).abs().max().item() - assert diff2 < 1e-5, f"inv RoPE at pos=0 should be identity, diff={diff2}" - print(f"✅ inv_rope at pos=0 is identity (diff={diff2:.8f})") - - -if __name__ == "__main__": - test_inv_rope_roundtrip() - test_wo_a_bmm() - test_inv_rope_at_zero() - print("\n✅ All attention O-projection tests passed") diff --git a/tests/archive/test_o_projection_b200.py b/tests/archive/test_o_projection_b200.py deleted file mode 100644 index f65bcd83..00000000 --- a/tests/archive/test_o_projection_b200.py +++ /dev/null @@ -1,306 +0,0 @@ -#!/usr/bin/env python3 -""" -B200 test: Proves the attention O-projection root cause and fix. - -Loads real model weights from the NVFP4 checkpoint and tests: - 1. OLD path: fused_inv_rope_fp8_quant + FP8 einsum (crashes with BF16 wo_a) - 2. NEW path: BF16 inv RoPE + BMM wo_a + NVFP4 wo_b (should work) - -Also tests the NVFP4 linear kernel (wo_b) with real weights. - -Usage (on B200): - python3 tests/test_o_projection_b200.py - -Requires: Real model weights at /root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4 - CuTeDSL, CUDA, Blackwell GPU -""" - -import sys -import os -import json -import torch -import torch.nn.functional as F -from safetensors import safe_open - -MODEL_PATH = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4" -DEVICE = "cuda:0" -LAYER_IDX = 0 - -# DeepSeek V4 Pro dimensions -HIDDEN_SIZE = 7168 -NUM_HEADS = 128 -HEAD_DIM = 512 -NOPE_DIM = 448 -ROPE_DIM = 64 -Q_LORA_RANK = 1536 -O_LORA_RANK = 1024 -O_GROUPS = 16 # from config (not TP-sharded) -HEADS_PER_GROUP = NUM_HEADS // O_GROUPS # 8 -NUM_TOKENS = 4 - -_cache = {} - - -def load_tensor(key, wm, model_dir): - if key in _cache: - return _cache[key] - shard_path = os.path.join(model_dir, wm[key]) - with safe_open(shard_path, framework="pt") as f: - t = f.get_tensor(key) - _cache[key] = t - return t - - -# ── OLD PATH: What the unpatched vLLM forward does ────────────────── - -def old_path_o_projection(o, positions, cos_sin_cache, wo_a_weight_bf16): - """Simulates the OLD (broken) attention forward. - - The old path does: - o_fp8, o_scale = fused_inv_rope_fp8_quant(o, ...) - wo_a_scale = self.wo_a.weight_scale_inv ← DOESN'T EXIST on BF16 wo_a - deepseek_v4_fp8_einsum(o_fp8, o_scale, wo_a_fp8, wo_a_scale, z, ...) - - Since wo_a is BF16 (no weight_scale_inv), this crashes. - We simulate the crash by showing weight_scale_inv doesn't exist. - """ - has_scale_inv = hasattr(wo_a_weight_bf16, 'weight_scale_inv') or \ - (isinstance(wo_a_weight_bf16, torch.Tensor) and False) - - # The weight is BF16 — FP8 einsum can't use it - # The old code does: wo_a_fp8 = self.wo_a.weight (BF16!) - # Then: wo_a_scale = self.wo_a.weight_scale_inv (AttributeError!) - - # Simulate what would happen: FP8 einsum with BF16 weight - # If we naively try to quantize BF16 to FP8 without proper scales... - print(" OLD PATH: wo_a.weight is BF16 (shape={}, dtype={})".format( - wo_a_weight_bf16.shape, wo_a_weight_bf16.dtype)) - - # Try to access weight_scale_inv (this is what the old code does) - if isinstance(wo_a_weight_bf16, torch.Tensor): - # This is what vLLM does: self.wo_a.weight_scale_inv - # Since wo_a is a plain BF16 tensor, this AttributeError crashes the worker - try: - _ = wo_a_weight_bf16.weight_scale_inv - print(" ❌ UNEXPECTED: weight_scale_inv exists (shouldn't for BF16)") - except AttributeError: - print(" ✅ CONFIRMED: weight_scale_inv does NOT exist → AttributeError in vLLM") - print(" This is the root cause: the FP8 einsum path crashes because") - print(" wo_a has quant_config=None (BF16) but the forward expects FP8.") - - return None # Can't produce valid output - - -# ── NEW PATH: Our patched forward ─────────────────────────────────── - -def apply_inv_rope_bf16(o, positions, cos_sin_cache, nope_dim=NOPE_DIM, rope_dim=ROPE_DIM): - """BF16 inverse RoPE (pure PyTorch).""" - if rope_dim == 0 or o.numel() == 0: - return o - half_rope = rope_dim // 2 - cos_all = cos_sin_cache[positions, :half_rope].unsqueeze(1).to(o.dtype) - sin_all = cos_sin_cache[positions, half_rope:].unsqueeze(1).to(o.dtype) - o_rope = o[:, :, nope_dim:] - o_even = o_rope[:, :, 0::2] - o_odd = o_rope[:, :, 1::2] - inv_even = o_even * cos_all + o_odd * sin_all - inv_odd = -o_even * sin_all + o_odd * cos_all - result = o.clone() - result[:, :, nope_dim:][:, :, 0::2] = inv_even - result[:, :, nope_dim:][:, :, 1::2] = inv_odd - return result - - -def new_path_o_projection(o, positions, cos_sin_cache, wo_a_weight_bf16): - """NEW path: BF16 inv RoPE + BMM wo_a. - - Returns z of shape (T, n_local_groups, o_lora_rank) ready for wo_b. - """ - # Step 1: Inverse RoPE (BF16) - o_inv = apply_inv_rope_bf16(o, positions, cos_sin_cache) - print(f" Inverse RoPE: shape={o_inv.shape} amax={o_inv.amax():.4f} NaN={torch.isnan(o_inv).any()}") - - # Step 2: wo_a BMM - num_tokens = o_inv.shape[0] - # wo_a weight: (O_GROUPS * O_LORA_RANK, HEADS_PER_GROUP * HEAD_DIM) - hidden_dim = wo_a_weight_bf16.shape[1] # 4096 = HEADS_PER_GROUP * HEAD_DIM - out_dim = wo_a_weight_bf16.shape[0] # 16384 = O_GROUPS * O_LORA_RANK - o_grouped = o_inv.view(num_tokens, O_GROUPS, hidden_dim) - wo_a_w = wo_a_weight_bf16.view(O_GROUPS, O_LORA_RANK, hidden_dim) - z = torch.bmm( - o_grouped.permute(1, 0, 2), - wo_a_w.transpose(1, 2), - ).permute(1, 0, 2) - print(f" wo_a BMM: shape={z.shape} amax={z.amax():.4f} NaN={torch.isnan(z).any()}") - - return z - - -# ── NVFP4 wo_b test ───────────────────────────────────────────────── - -E2M1_LUT = torch.tensor([ - 0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, - -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0, -], dtype=torch.float32) - - -def dequant_nvfp4(packed_uint8, scale_e4m3, global_scale): - """Dequantize NVFP4 weight to BF16 for reference.""" - device = packed_uint8.device - lut = E2M1_LUT.to(device) - lower = lut[(packed_uint8 & 0x0F).long()] - upper = lut[((packed_uint8 >> 4) & 0x0F).long()] - out_features = packed_uint8.shape[0] - in_features = packed_uint8.shape[1] * 2 - unpacked = torch.empty(out_features, in_features, dtype=torch.float32, device=device) - unpacked[:, 0::2] = lower - unpacked[:, 1::2] = upper - block_scale = scale_e4m3.float() - block_expanded = block_scale.repeat_interleave(16, dim=1)[:out_features, :in_features] - return (unpacked * block_expanded * global_scale).to(torch.bfloat16) - - -def test_wo_b_nvfp4(z, wo_b_weight, wo_b_sf, wo_b_gs): - """Test wo_b NVFP4 GEMM against BF16 reference.""" - sys.path.insert(0, "/root/nvfp4-megamoe-kernel") - from dsv4.layers.linear import Nvfp4Linear - - in_features = wo_b_weight.shape[1] * 2 - out_features = wo_b_weight.shape[0] - - # Convert to CuTeDSL format - fp4 = [wo_b_weight.view(torch.float4_e2m1fn_x2).permute(1, 0).contiguous()] - sf = [wo_b_sf.permute(1, 0).contiguous()] - gs = [wo_b_gs] - - runner = Nvfp4Linear( - in_features=in_features, - out_features=out_features, - max_num_tokens=8192, - device=DEVICE, - ) - runner.fp4 = fp4 - runner.sf = sf - runner.gs = gs - runner.finalize_weights() - runner._ensure_initialized() - - # Warmup: compute activation global scale - z_flat = z.flatten(1) # (T, O_GROUPS * O_LORA_RANK) - runner.compute_activation_global_scale(z_flat) - - # Run CuTeDSL - with torch.no_grad(): - output = runner.run(z_flat) - print(f" wo_b CuTeDSL: shape={output.shape} amax={output.amax():.4f} NaN={torch.isnan(output).any()}") - - # BF16 reference - bf16_w = dequant_nvfp4(wo_b_weight, wo_b_sf, wo_b_gs) - with torch.no_grad(): - ref = z_flat @ bf16_w.T - print(f" wo_b BF16 ref: shape={ref.shape} amax={ref.amax():.4f}") - - cos = F.cosine_similarity(ref.flatten().unsqueeze(0), - output.flatten().unsqueeze(0)).item() - mse = (ref - output).pow(2).mean().item() - status = "✅" if cos >= 0.98 else "❌" - print(f" wo_b cosine={cos:.6f} MSE={mse:.6e} {status}") - return cos - - -def build_cos_sin_cache(max_pos=4096, rope_dim=ROPE_DIM): - """Build cos_sin_cache in the same format as vLLM's RotaryEmbedding.""" - half_rope = rope_dim // 2 - base = 10000.0 - inv_freq = 1.0 / (base ** (torch.arange(0, half_rope, dtype=torch.float32) / half_rope)) - t = torch.arange(max_pos, dtype=torch.float32) - freqs = torch.outer(t, inv_freq) # (max_pos, half_rope) - return torch.cat([freqs.cos(), freqs.sin()], dim=-1) # (max_pos, rope_dim) - - -def main(): - torch.cuda.set_device(0) - torch.manual_seed(42) - - print("=" * 70) - print(" B200 Test: O-Projection Root Cause + Fix") - print("=" * 70) - - # Load weight map - with open(os.path.join(MODEL_PATH, "model.safetensors.index.json")) as f: - wm = json.load(f)["weight_map"] - P = lambda key: load_tensor(key, wm, MODEL_PATH).to(DEVICE) - - prefix = f"model.layers.{LAYER_IDX}.self_attn" - - # Load wo_a (BF16) and wo_b (NVFP4) weights - print("\n--- Loading weights ---") - wo_a_w = P(f"{prefix}.o_a_proj.weight") - print(f" wo_a: shape={wo_a_w.shape} dtype={wo_a_w.dtype}") - - wo_b_w = P(f"{prefix}.o_b_proj.weight") - wo_b_sf = P(f"{prefix}.o_b_proj.weight_scale") - wo_b_gs = P(f"{prefix}.o_b_proj.weight_scale_2").item() - print(f" wo_b: shape={wo_b_w.shape} dtype={wo_b_w.dtype} gs={wo_b_gs:.8f}") - - # Check: wo_a should NOT have weight_scale_inv - # (it's a plain BF16 tensor, not a quantized layer) - - # Build cos_sin_cache - cos_sin_cache = build_cos_sin_cache().to(DEVICE) - - # Simulate attention output (what FlashMLA would produce) - print("\n--- Simulating attention output ---") - positions = torch.tensor([0, 1, 2, 3], dtype=torch.int64, device=DEVICE) - o = torch.randn(NUM_TOKENS, NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device=DEVICE) * 0.1 - print(f" Attention output: shape={o.shape} amax={o.amax():.4f}") - - # ═══════════════════════════════════════════════════════════════════ - # TEST 1: OLD PATH (should show crash/AttributeError) - # ═══════════════════════════════════════════════════════════════════ - print("\n" + "=" * 70) - print(" TEST 1: OLD PATH (FP8 einsum — should crash)") - print("=" * 70) - old_path_o_projection(o, positions, cos_sin_cache, wo_a_w) - - # ═══════════════════════════════════════════════════════════════════ - # TEST 2: NEW PATH (BF16 inv RoPE + BMM wo_a — should work) - # ═══════════════════════════════════════════════════════════════════ - print("\n" + "=" * 70) - print(" TEST 2: NEW PATH (BF16 inv RoPE + BMM wo_a)") - print("=" * 70) - z = new_path_o_projection(o, positions, cos_sin_cache, wo_a_w) - - if z is not None and not torch.isnan(z).any(): - # ═══════════════════════════════════════════════════════════════ - # TEST 3: wo_b NVFP4 GEMM - # ═══════════════════════════════════════════════════════════════ - print("\n" + "=" * 70) - print(" TEST 3: wo_b NVFP4 GEMM (CuTeDSL vs BF16 reference)") - print("=" * 70) - cos_wo_b = test_wo_b_nvfp4(z, wo_b_w, wo_b_sf, wo_b_gs) - else: - print("\n❌ z is invalid (NaN or None), skipping wo_b test") - cos_wo_b = 0.0 - - # ═══════════════════════════════════════════════════════════════════ - # SUMMARY - # ═══════════════════════════════════════════════════════════════════ - print("\n" + "=" * 70) - print(" SUMMARY") - print("=" * 70) - print(" OLD PATH (FP8 einsum): CRASHES — wo_a has no weight_scale_inv") - print(f" NEW PATH (BF16 inv RoPE + BMM): z amax={z.amax():.4f} NaN={torch.isnan(z).any()}") - print(f" wo_b NVFP4 cosine: {cos_wo_b:.6f} {'✅' if cos_wo_b >= 0.98 else '❌'}") - - if cos_wo_b >= 0.98 and z is not None and not torch.isnan(z).any(): - print("\n✅ ROOT CAUSE CONFIRMED + FIX VALIDATED") - print(" The attention forward was crashing because wo_a is BF16") - print(" but the FP8 einsum path expected weight_scale_inv.") - print(" Our patched forward (BF16 inv RoPE + BMM) fixes this.") - else: - print("\n❌ FIX INCOMPLETE — further investigation needed") - - -if __name__ == "__main__": - main() diff --git a/tests/archive/test_packing_diag.py b/tests/archive/test_packing_diag.py deleted file mode 100644 index d9db61e5..00000000 --- a/tests/archive/test_packing_diag.py +++ /dev/null @@ -1,133 +0,0 @@ -""" -BF16 Packing Diagnostic: Write specific F32 bit patterns to P TMEM via St32x32bOp. -Then run PV MMA with V=identity. Output = P (as MMA reads it). - -This reveals the BF16 packing order within F32 TMEM words. - -Strategy: -- Skip MMA1 entirely (no QK computation) -- Write packed F32 values where low BF16 ≠ high BF16 -- Use K=V=identity so output[i,j] = P[i,j] -- Compare output against both packing orders -""" -import torch, struct, sys -sys.path.insert(0, '/root/dsv4-nvfp4-workspace/kernel/tests') - -# Use the existing diag test that writes P=all-ones, but modify to write specific patterns -# The diag test is test_stage_b_diag.py - let me copy and modify v7 instead - -# Actually, let me just use the EXISTING kernel (test_stage_b_v7) with: -# 1. K = identity matrix -# 2. The identity softmax running normally (it writes softmax(Q@K^T) as P) -# 3. If Q is also identity, then Q@K^T = I@I = I, softmax(I) = softmax of identity -# 4. That's not what we want either. - -# Better: use the diag test (writes P=all-ones) with K=identity -# Output should be all-ones * I = all-ones. That gives no new info. - -# The REAL test: modify the kernel to write specific F32 values to the BF16 recast view -# Instead of writing 1.0 BF16, write alternating different values - -# Simplest approach: modify test_stage_b_diag.py to write j+1 in BF16 for each position -# Then with V=identity, output[i,j] = j+1 - -# But modifying the kernel requires JIT changes. Let me use the existing v7 kernel -# with a clever input choice instead. - -# Approach: Use the identity softmax (which writes Q@K^T scores as P) -# With Q=randn and K=identity: Q@K^T = Q (since K=I) -# Then P = softmax(Q) and output = softmax(Q) @ V -# With V=I: output = softmax(Q) -# This tests the FULL pipeline but doesn't isolate packing. - -# Let me just write the packing test kernel from scratch, minimally. -# It only needs: TMA load V, write F32 to P TMEM, PV MMA, epilogue. - -# Actually, the simplest isolation test: -# Use the existing test with P=all-ones and V=K (K=randn) -# We already know cos=0.08 for this. The 0.08 is not 0 or 1. -# If I use V where V[j] = 1 if j even, 0 if j odd, then: -# With correct packing: output = (number of ones in even K positions) per row -# This doesn't help because P=all-ones means all K positions are 1.0 - -# I think the key issue is that with P=all-ones (cos=0.08), the output should be -# EXACTLY sum(V, dim=0) for each row. Let me compare more carefully. - -# Let me run the EXISTING P=all-ones diag test and compare the output values -# against the reference. The PATTERN of errors will tell us about packing. - -print("Running existing P=all-ones diag with K=V=randn") -print("Comparing output vs reference to identify the error pattern") - -import cutlass.cute as cute -import cutlass.utils as utils -from cutlass.cute.nvgpu import tcgen05 -from cutlass import Float32, BFloat16 -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cutlass.torch as ct -import cuda.bindings.driver as cuda - -torch.manual_seed(42) -m, n, k = 128, 128, 128 -q = torch.randn(m, k, 1, dtype=torch.bfloat16, device='cuda') -kv = torch.randn(n, k, 1, dtype=torch.bfloat16, device='cuda') -c = torch.zeros(m, n, 1, dtype=torch.bfloat16, device='cuda') - -kvf = kv[:,:,0].float() -ref = torch.ones(128, 128, dtype=torch.float32, device='cuda') @ kvf - -mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) -mK = ct.from_dlpack(kv).mark_layout_dynamic(leading_dim=ct.get_leading_dim(kv)) -mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) -stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - -from test_stage_b_diag import StageBDiag -kernel = StageBDiag(mma_tiler_mn=(128, 128), use_2cta_instrs=False, use_tma_store=True) -print('Compiling diag test...', flush=True) -compiled = cute.compile(kernel, mQ, mK, mC, stream) -print('Running...', flush=True) -compiled(mQ, mK, mC, stream) -torch.cuda.synchronize() - -out = c[:,:,0].float() -cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - -print(f'\nCosine: {cos:.6f}') -print(f'Output row 0[:8]: {out[0,:8].tolist()}') -print(f'Ref row 0[:8]: {ref[0,:8].tolist()}') -print(f'Ratio out/ref row 0[:8]: {(out[0,:8] / ref[0,:8]).tolist()}') - -# Check if output is a scaled version of the reference -ratio = out[0] / ref[0] -ratio_clean = ratio[torch.isfinite(ratio)] -print(f'Ratio mean: {ratio_clean.mean().item():.6f}, std: {ratio_clean.std().item():.6f}') - -# Check if odd/even columns have different ratios -ratio_even = (out[0, 0::2] / ref[0, 0::2])[torch.isfinite(out[0, 0::2] / ref[0, 0::2])] -ratio_odd = (out[0, 1::2] / ref[0, 1::2])[torch.isfinite(out[0, 1::2] / ref[0, 1::2])] -print(f'Even col ratio mean: {ratio_even.mean().item():.6f}') -print(f'Odd col ratio mean: {ratio_odd.mean().item():.6f}') - -# Check if output matches a different V interpretation -# What if the V SMEM is being read in a different column order? -# Compute: what if V columns were permuted? -# With P=all-ones, output[i] = sum(V[:,j]) for each j -# This is the column sum of V, broadcast to all rows -# If V is read in a different order, the output would be the sum of -# differently-ordered V columns, but still all the same sum -# So output should be uniform across columns = sum of all V values per column -# But the output IS uniform (all rows identical). The VALUES are wrong. -# That means the SUM is wrong, which means the V values being read are wrong. - -# Let me check: does the output column structure match any known transform of V? -print(f'\nV (K) row 0[:8]: {kvf[0,:8].tolist()}') -print(f'V col sums (reference): {ref[0,:4].tolist()}') -print(f'Output col 0[:4]: {out[0,:4].tolist()}') - -# What if V is being read transposed? -# sum(V[j,:]) per column j = row sums of V -row_sums = kvf.sum(dim=1) -col_sums = kvf.sum(dim=0) -print(f'V row sums[:4]: {row_sums[:4].tolist()}') -print(f'V col sums[:4]: {col_sums[:4].tolist()}') diff --git a/tests/archive/test_pair_swap.py b/tests/archive/test_pair_swap.py deleted file mode 100644 index 68a2041a..00000000 --- a/tests/archive/test_pair_swap.py +++ /dev/null @@ -1,119 +0,0 @@ -""" -BF16 Pair-Swap Test: compare kernel output against reference with V rows -swapped in even/odd pairs (simulating BF16 packing swap within F32 words). -""" -import torch, sys -sys.path.insert(0, '/root/dsv4-nvfp4-workspace/kernel/tests') - -import cutlass.cute as cute -import cutlass.torch as ct -import cuda.bindings.driver as cuda -from test_stage_b_v7 import StageBIdentitySoftmax - -torch.manual_seed(42) -m, n, k = 128, 128, 128 -q = torch.randn(m, k, 1, dtype=torch.bfloat16, device='cuda') -kv = torch.randn(n, k, 1, dtype=torch.bfloat16, device='cuda') -c = torch.zeros(m, n, 1, dtype=torch.bfloat16, device='cuda') - -qf = q[:,:,0].float() -kvf = kv[:,:,0].float() - -# Standard reference: P @ V where P = Q @ K^T -ref = qf @ kvf.T @ kvf - -# Pair-swapped reference: if BF16 within each F32 are swapped, -# MMA reads K[2j] and K[2j+1] swapped, which means V rows are swapped in pairs -# Output = P @ V_with_row_pairs_swapped -V_swap = kvf.clone() -V_swap[0::2], V_swap[1::2] = kvf[1::2].clone(), kvf[0::2].clone() -ref_swap = qf @ kvf.T @ V_swap - -# Also try: just the P@K^T with V where every 2 consecutive rows are swapped -# but starting from different offsets -# Try 1-based offset: swap rows (1,2), (3,4), ... -V_swap1 = kvf.clone() -V_swap1[1::2], V_swap1[2::2] = kvf[2::2].clone(), kvf[1::2].clone() -V_swap1[0] = kvf[0] # row 0 unchanged -ref_swap1 = qf @ kvf.T @ V_swap1 - -mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) -mK = ct.from_dlpack(kv).mark_layout_dynamic(leading_dim=ct.get_leading_dim(kv)) -mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) -stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - -kernel = StageBIdentitySoftmax(mma_tiler_mn=(128, 128), use_2cta_instrs=False, use_tma_store=True) -print('Compiling...', flush=True) -compiled = cute.compile(kernel, mQ, mK, mC, stream) -print('Running...', flush=True) -compiled(mQ, mK, mC, stream) -torch.cuda.synchronize() - -out = c[:,:,0].float() -cos_ref = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() -cos_swap = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref_swap.flatten().unsqueeze(0)).item() -cos_swap1 = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref_swap1.flatten().unsqueeze(0)).item() - -print(f'\nCosine with standard ref: {cos_ref:.6f}') -print(f'Cosine with pair-swapped ref: {cos_swap:.6f}') -print(f'Cosine with offset-1 swapped: {cos_swap1:.6f}') - -# Also try: what if the entire P row is shifted by some amount? -# Or what if P columns are reordered according to the A-fragment's K partitioning? -# The A-fragment has K laid out as (16 inner, 4 outer chunks of 16) -# If the store writes columns sequentially but the MMA reads them in the -# chunked order, we'd get a specific permutation - -# The A-fragment K layout: (col, k0, k1) where col=0..15, k0=0..3, k1=0..1 -# K position = col + 16*k0 + 64*k1 -# If the store writes K sequentially (0,1,2,...,127) but the MMA reads -# in the fragment order, the mapping would be: -# Fragment index (col, k0, k1) -> K position (col + 16*k0 + 64*k1) -# But this is sequential (0,1,2,...), so no permutation. -# UNLESS the store doesn't fill the fragment's K blocks correctly. - -# Try: what if only the FIRST 64 K values (not 128) are filled? -# The store writes 64 F32 = 128 BF16. But what if the MMA's A-fragment -# K dimension is 64 (not 128)? Then only 64 K values have P, rest are garbage. -# But we already verified nblk_pv=4 and the fragment covers 128 BF16. - -# Try: what if the 64 F32 columns in the store map to 128 BF16 K positions -# but with the packing where BF16[K] and BF16[K+1] come from the same F32 word -# and K is determined by the A-fragment's partition order? - -# The A-fragment reads K in groups of 16 (inner K), then 4 outer groups, then 2 BF16 per F32 -# So K values are accessed as: [group0_bf16_0, group0_bf16_1, group1_bf16_0, group1_bf16_1, ...] -# where group = (col, k0, k1) and bf16_0/bf16_1 are the two BF16 in each F32 word - -# This is getting complex. Let me just check the specific permutation by -# matching output columns to reference columns for row 0. - -# For each output column j, find which reference column it matches -# (using the entire row as a signature) -print('\nTrying to identify the column permutation for row 0...') -# Since all values might not be unique, use the dot product with a known vector -# Actually, the simplest: for output[0, j], check if it equals ref[0, k] for any k -# But we need a unique signature. Let's use the output of the ENTIRE row. - -# Compare output row i with reference rows -for i in [0, 1, 2]: - best_cos = 0 - best_j = -1 - for j in range(128): - c = torch.nn.functional.cosine_similarity(out[i].unsqueeze(0), ref[j].unsqueeze(0)).item() - if c > best_cos: - best_cos = c - best_j = j - print(f' Output row {i} best matches ref row {best_j} (cos={best_cos:.4f})') - -# Also: compare output column j with reference columns -print('\nColumn matching (output col j vs ref col k):') -for j in [0, 1, 2, 3, 4, 5, 6, 7]: - best_cos = 0 - best_k = -1 - for k in range(128): - c = torch.nn.functional.cosine_similarity(out[:, j].unsqueeze(0), ref[:, k].unsqueeze(0)).item() - if c > best_cos: - best_cos = c - best_k = k - print(f' Output col {j} best matches ref col {best_k} (cos={best_cos:.4f})') diff --git a/tests/archive/test_pair_swap2.py b/tests/archive/test_pair_swap2.py deleted file mode 100644 index 2f4d7a5d..00000000 --- a/tests/archive/test_pair_swap2.py +++ /dev/null @@ -1,95 +0,0 @@ -""" -BF16 Pair-Swap Test: compare kernel output against reference with V rows -swapped in even/odd pairs (simulating BF16 packing swap within F32 words). -Also tries to identify the exact column permutation. -""" -import torch, sys -sys.path.insert(0, '/root/dsv4-nvfp4-workspace/kernel/tests') - -import cutlass.cute as cute -import cutlass.torch as ct -import cuda.bindings.driver as cuda -from test_stage_b_v7 import StageBIdentitySoftmax - -torch.manual_seed(42) -m, n, k = 128, 128, 128 -q = torch.randn(m, k, 1, dtype=torch.bfloat16, device='cuda') -kv = torch.randn(n, k, 1, dtype=torch.bfloat16, device='cuda') -c = torch.zeros(m, n, 1, dtype=torch.bfloat16, device='cuda') - -qf = q[:,:,0].float() -kvf = kv[:,:,0].float() -ref = qf @ kvf.T @ kvf - -# Pair-swapped: V rows 0↔1, 2↔3, 4↔5, ... -V_swap = kvf.clone() -V_swap[0::2], V_swap[1::2] = kvf[1::2].clone(), kvf[0::2].clone() -ref_swap = qf @ kvf.T @ V_swap - -mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) -mK = ct.from_dlpack(kv).mark_layout_dynamic(leading_dim=ct.get_leading_dim(kv)) -mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) -stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - -kernel = StageBIdentitySoftmax(mma_tiler_mn=(128, 128), use_2cta_instrs=False, use_tma_store=True) -print('Compiling...', flush=True) -compiled = cute.compile(kernel, mQ, mK, mC, stream) -print('Running...', flush=True) -compiled(mQ, mK, mC, stream) -torch.cuda.synchronize() - -out = c[:,:,0].float() -cos_ref = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() -cos_swap = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref_swap.flatten().unsqueeze(0)).item() - -print(f'\nCosine with standard ref: {cos_ref:.6f}') -print(f'Cosine with pair-swapped ref: {cos_swap:.6f}') - -# Column permutation identification -print('\nColumn permutation (output col j matches ref col k):') -perm = [] -for j in range(min(16, 128)): - best_cos = 0 - best_k = -1 - for k in range(128): - c_val = torch.nn.functional.cosine_similarity(out[:, j].unsqueeze(0), ref[:, k].unsqueeze(0)).item() - if c_val > best_cos: - best_cos = c_val - best_k = k - perm.append(best_k) - if j < 16: - print(f' out col {j} -> ref col {best_k} (cos={best_cos:.4f})') - -# Check if permutation has a pattern -print(f'\nPermutation (first 16): {perm}') -diffs = [perm[j+1] - perm[j] for j in range(len(perm)-1)] -print(f'Permutation diffs: {diffs}') - -# Check: is perm[j] = j with every pair swapped? (0→1, 1→0, 2→3, 3→2, ...) -expected_pair_swap = [j^1 for j in range(128)] # XOR with 1 swaps even/odd -matches_pair_swap = all(perm[j] == expected_pair_swap[j] for j in range(len(perm))) -print(f'Matches pair-swap (0↔1, 2↔3, ...): {matches_pair_swap}') - -# Also check the full permutation for all 128 columns -full_perm = [] -for j in range(128): - best_cos = 0 - best_k = -1 - for k in range(128): - c_val = torch.nn.functional.cosine_similarity(out[:, j].unsqueeze(0), ref[:, k].unsqueeze(0)).item() - if c_val > best_cos: - best_cos = c_val - best_k = k - full_perm.append(best_k) - -print(f'\nFull permutation: {full_perm}') -# Check if it's pair-swap for all columns -all_pair_swap = all(full_perm[j] == (j^1) for j in range(128)) -print(f'All columns match pair-swap: {all_pair_swap}') - -# If not pair-swap, what's the pattern? -if not all_pair_swap: - # Check if it's a 16-element block permutation - for block in range(8): - block_perm = full_perm[block*16:(block+1)*16] - print(f' Block {block}: {block_perm}') diff --git a/tests/archive/test_pipeline_real_weights.py b/tests/archive/test_pipeline_real_weights.py deleted file mode 100644 index 517db7cb..00000000 --- a/tests/archive/test_pipeline_real_weights.py +++ /dev/null @@ -1,164 +0,0 @@ -"""Full pipeline test: Fixed runner vs BF16 reference.""" -import torch -import torch.nn.functional as F -import sys, os, glob - -sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), '..')) - -MODEL_PATH = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4" -LAYER_IDX = 0 -NUM_EXPERTS = 48 -HIDDEN_SIZE = 7168 -INTERMEDIATE_SIZE = 3072 -NUM_TOKENS = 8 -TOP_K = 6 -SWIGLU_LIMIT = 10.0 -DEVICE = "cuda" -MAX_NUM_TOKENS = 8192 # match vLLM config - - -def load_layer_tensors(model_dir, layer_idx): - tensors = {} - for sf in glob.glob(os.path.join(model_dir, "*.safetensors")): - from safetensors.torch import load_file - data = load_file(sf) - for k, v in data.items(): - if f"layers.{layer_idx}." in k and "mlp.experts" in k: - tensors[k.removeprefix("model.")] = v - return tensors - - -def dequantize_nvfp4_weight(packed_uint8, scale_e4m3, global_scale): - lut = torch.tensor([0.,0.5,1.,1.5,2.,3.,4.,6.,-0.,-0.5,-1.,-1.5,-2.,-3.,-4.,-6.], - dtype=torch.float32, device=packed_uint8.device) - lower = lut[(packed_uint8 & 0x0F).long()] - upper = lut[((packed_uint8 >> 4) & 0x0F).long()] - N, K = packed_uint8.shape[0], packed_uint8.shape[1] * 2 - bf16 = torch.stack([lower, upper], dim=-1).reshape(N, K) - K_sf = scale_e4m3.shape[1] - scale_2d = scale_e4m3.float().repeat_interleave(K // K_sf, dim=1) - return (bf16 * scale_2d * global_scale).to(torch.bfloat16) - - -def main(): - torch.cuda.set_device(0) - torch.manual_seed(42) - - print("=== Full Pipeline Test (Fixed Runner) ===") - nvfp4_tensors = load_layer_tensors(MODEL_PATH, LAYER_IDX) - expert_indices = list(range(NUM_EXPERTS)) - - hidden_states = torch.randn(NUM_TOKENS, HIDDEN_SIZE, dtype=torch.bfloat16, device=DEVICE) * 2.0 - topk_ids = torch.zeros(NUM_TOKENS, TOP_K, dtype=torch.int64, device=DEVICE) - for i in range(NUM_TOKENS): - topk_ids[i] = torch.randperm(NUM_EXPERTS)[:TOP_K] - topk_weights = torch.ones(NUM_TOKENS, TOP_K, dtype=torch.float32, device=DEVICE) / TOP_K - - # BF16 reference - ref_out = torch.zeros(NUM_TOKENS, HIDDEN_SIZE, dtype=torch.bfloat16, device=DEVICE) - for i, e in enumerate(expert_indices): - dk = f"layers.{LAYER_IDX}.mlp.experts.{e}.down_proj.weight" - gk = f"layers.{LAYER_IDX}.mlp.experts.{e}.gate_proj.weight" - uk = f"layers.{LAYER_IDX}.mlp.experts.{e}.up_proj.weight" - if dk not in nvfp4_tensors: - continue - gate_bf16 = dequantize_nvfp4_weight( - nvfp4_tensors[gk].to(DEVICE), - nvfp4_tensors[gk.replace('.weight', '.weight_scale')].to(DEVICE), - nvfp4_tensors[gk.replace('.weight', '.weight_scale_2')].item()) - up_bf16 = dequantize_nvfp4_weight( - nvfp4_tensors[uk].to(DEVICE), - nvfp4_tensors[uk.replace('.weight', '.weight_scale')].to(DEVICE), - nvfp4_tensors[uk.replace('.weight', '.weight_scale_2')].item()) - down_bf16 = dequantize_nvfp4_weight( - nvfp4_tensors[dk].to(DEVICE), - nvfp4_tensors[dk.replace('.weight', '.weight_scale')].to(DEVICE), - nvfp4_tensors[dk.replace('.weight', '.weight_scale_2')].item()) - - for t in range(NUM_TOKENS): - for k in range(TOP_K): - if topk_ids[t, k].item() != i: - continue - w = topk_weights[t, k].item() - x = hidden_states[t] - gate = x @ gate_bf16.T - up = x @ up_bf16.T - gate_silu = F.silu(gate).clamp(max=SWIGLU_LIMIT) - up = up.clamp(min=-SWIGLU_LIMIT, max=SWIGLU_LIMIT) - act = gate_silu * up - ref_out[t] += w * (act @ down_bf16.T) - - print(f"BF16 ref: amax={ref_out.amax().item():.4f}") - - # CuTeDSL runner - from vllm.nvfp4_cutedsl import Nvfp4MoE -from dsv4.ops.layouts import ( - assemble_scales_3d_side, - make_b_k_major, -) - - l1_fp4, l1_sf, l1_gs = [], [], [] - l2_fp4, l2_sf, l2_gs = [], [], [] - for e in expert_indices: - gw = nvfp4_tensors[f"layers.{LAYER_IDX}.mlp.experts.{e}.gate_proj.weight"].to(DEVICE) - uw = nvfp4_tensors[f"layers.{LAYER_IDX}.mlp.experts.{e}.up_proj.weight"].to(DEVICE) - gsf = nvfp4_tensors[f"layers.{LAYER_IDX}.mlp.experts.{e}.gate_proj.weight_scale"].to(DEVICE) - usf = nvfp4_tensors[f"layers.{LAYER_IDX}.mlp.experts.{e}.up_proj.weight_scale"].to(DEVICE) - ggs = nvfp4_tensors[f"layers.{LAYER_IDX}.mlp.experts.{e}.gate_proj.weight_scale_2"].item() - ugs = nvfp4_tensors[f"layers.{LAYER_IDX}.mlp.experts.{e}.up_proj.weight_scale_2"].item() - fw = torch.cat([gw, uw], dim=0).view(torch.float4_e2m1fn_x2).permute(1,0).contiguous() - fsf = torch.cat([gsf, usf], dim=0).permute(1,0).contiguous() - mgs = max(ggs, ugs) - if ggs != ugs: - sf32 = fsf.float() - sf32[:, :INTERMEDIATE_SIZE] *= (ggs / mgs) - sf32[:, INTERMEDIATE_SIZE:] *= (ugs / mgs) - fsf = sf32.to(torch.float8_e4m3fn) - l1_fp4.append(fw); l1_sf.append(fsf); l1_gs.append(mgs) - - dk = f"layers.{LAYER_IDX}.mlp.experts.{e}.down_proj.weight" - if dk in nvfp4_tensors: - dw = nvfp4_tensors[dk].to(DEVICE) - dsf = nvfp4_tensors[f"layers.{LAYER_IDX}.mlp.experts.{e}.down_proj.weight_scale"].to(DEVICE) - dgs = nvfp4_tensors[f"layers.{LAYER_IDX}.mlp.experts.{e}.down_proj.weight_scale_2"].item() - l2_fp4.append(dw.view(torch.float4_e2m1fn_x2).permute(1,0).contiguous()) - l2_sf.append(dsf.permute(1,0).contiguous()); l2_gs.append(dgs) - else: - l2_fp4.append(torch.zeros(INTERMEDIATE_SIZE//2, HIDDEN_SIZE, dtype=torch.float4_e2m1fn_x2, device=DEVICE)) - l2_sf.append(torch.ones(INTERMEDIATE_SIZE//16, HIDDEN_SIZE, dtype=torch.float8_e4m3fn, device=DEVICE)) - l2_gs.append(1.0) - - runner = Nvfp4MoE( - num_experts=NUM_EXPERTS, hidden_size=HIDDEN_SIZE, - intermediate_size=INTERMEDIATE_SIZE, max_num_tokens=MAX_NUM_TOKENS, - top_k=TOP_K, device=DEVICE, - ) - runner.l1_fp4 = l1_fp4; runner.l1_sf = l1_sf; runner.l1_gs = l1_gs - runner.l2_fp4 = l2_fp4; runner.l2_sf = l2_sf; runner.l2_gs = l2_gs - runner.set_swiglu_limit(SWIGLU_LIMIT) - - with torch.no_grad(): - runner.compute_activation_global_scales(hidden_states, topk_weights, topk_ids) - runner_out = runner.run(hidden_states, topk_weights, topk_ids) - - print(f"Runner: amax={runner_out.amax().item():.4f}") - print(f"NaN: {torch.isnan(runner_out).any().item()}") - - cos = F.cosine_similarity(ref_out.flatten().unsqueeze(0), runner_out.flatten().unsqueeze(0)).item() - mse = (ref_out - runner_out).pow(2).mean().item() - print(f"\nCosine: {cos:.6f} MSE: {mse:.6e}") - - for t in range(NUM_TOKENS): - ct = F.cosine_similarity(ref_out[t].unsqueeze(0), runner_out[t].unsqueeze(0)).item() - print(f" Token {t}: cosine={ct:.4f}") - - if cos >= 0.98: - print(f"\n✅ PASS") - elif cos >= 0.90: - print(f"\n⚠️ MARGINAL") - else: - print(f"\n❌ FAIL") - - -if __name__ == "__main__": - main() diff --git a/tests/archive/test_pv64.py b/tests/archive/test_pv64.py deleted file mode 100644 index f2986d04..00000000 --- a/tests/archive/test_pv64.py +++ /dev/null @@ -1,244 +0,0 @@ -""" -Test (128,64) PV with separate V SMEM allocation. -Based on the working test_128_128_vdiag.py, adapted for head_dim=64. -Single ab pipeline, Q+K+V loaded together. -QK (all tiles) → softmax → PV (all tiles) → epilogue. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - -HEAD_DIM = 64 - -class Pv64Test: - def __init__(self): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0,1,2,3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.num_c_stage = 2 - self.num_ab_stage = 1; self.num_acc_stage = 1 - - def _setup(self, qk_mma, pv_mma): - qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (128, 128, qk_ik * 4) - pv_ik = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (128, HEAD_DIM, pv_ik * (128 // pv_ik)) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0]//cute.size(qk_mma.thr_id.shape), HEAD_DIM, self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - qk_thr = qk_mma.get_slice(0); qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - pv_thr = pv_mma.get_slice(0); pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0; self.tmem_p0_offset = 32 - self.tmem_o0_offset = find_tmem_tensor_col_offset(tOtO) - tCS = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage)) - tCO = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCS, tCO], arch="sm_100") - a_s = cute.slice_(self.a_smem_s,(None,None,None,0)); b_s = cute.slice_(self.b_smem_s,(None,None,None,0)) - v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - self.num_tma_load_bytes = (cute.size_in_bytes(self.q_dtype,a_s)+cute.size_in_bytes(self.q_dtype,b_s)+cute.size_in_bytes(self.q_dtype,v_s))*cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, self.a_major, self.b_major, self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, self.qk_acc_dtype, self.cta_group, (128,HEAD_DIM), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - q_s = cute.slice_(self.a_smem_s,(None,None,None,0)); k_s = cute.slice_(self.b_smem_s,(None,None,None,0)) - v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id),q,q_s,self.mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id),k,k_s,self.mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape) - epi_s = cute.select(self.c_smem_s,mode=[0,1]) - tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile) - self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.a_smem_s,self.b_smem_s,self.v_smem_s,self.p_tmem_s,self.c_smem_s,self.epi_tile).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx,_,_ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage*2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage*2] - tmem_dealloc: cutlass.Int64; holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - ab_p,ab_c = pipeline.PipelineTmaUmma.create(barrier_storage=st.ab_bar.data_ptr(),num_stages=self.num_ab_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.num_tma_load_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - mma_si_prod,mma_si_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.mma_si_bar.data_ptr(),num_stages=1,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.epilogue_warp_id))).make_participants() - acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(),num_stages=self.num_acc_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,len(self.epilogue_warp_id)),cta_layout_vmnk=cl_vmnk,defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2,num_threads=32*len((self.mma_warp_id,*self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr,barrier_for_retrieve=tmem_bar,allocator_warp_id=self.epilogue_warp_id[0],is_two_cta=cute.size(qk_mma.thr_id.shape)==2,two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk,is_relaxed=True) - sQ = smem.allocate_tensor(element_type=self.q_dtype,layout=a_smem_s.outer,byte_alignment=128,swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype,layout=b_smem_s.outer,byte_alignment=128,swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype,layout=v_smem_s.outer,byte_alignment=128,swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype,layout=c_smem_s.outer,byte_alignment=128,swizzle=c_smem_s.inner) - gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None)) - gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None)) - gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None)) - gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - tCgV = pv_thr.partition_B(gV); tCgC = pv_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,0,None,0)).shape) - tAsQ,tAgQ = cpasync.tma_partition(tma_q,0,a_lay,cute.group_modes(sQ,0,3),cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape) - tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3)) - tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)]; tVgV = tVgV[(None,0,None,0)] - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - tStS0 = cute.make_tensor(tStS.iterator+self.tmem_s0_offset,tStS.layout) - pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - tOtO0 = cute.make_tensor(tOtO.iterator+self.tmem_o0_offset,tOtO.layout) - tP = cute.make_tensor(tStS.iterator,p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None,None,None,0)] - tOrP0 = cute.make_tensor(tOrP.iterator+self.qk_acc_dtype.width//self.q_dtype.width*self.tmem_p0_offset,tOrP.layout) - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_as,self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_as,self.num_acc_stage)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt,unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q,tAgQ[(None,h.count)],tAsQ[(None,h.index)],tma_bar_ptr=h.barrier) - cute.copy(tma_k,tBgK[(None,h.count)],tBsK[(None,h.index)],tma_bar_ptr=h.barrier) - cute.copy(tma_v,tVgV[(None,h.count)],tVsV[(None,h.index)],tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else "FAIL"}') - if cos < 0.99: - print(f' out[0,:4]={out[0,:4].tolist()} ref[0,:4]={ref[0,:4].tolist()}') - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_pv64_fmha_v.py b/tests/archive/test_pv64_fmha_v.py deleted file mode 100644 index 8b7e61b9..00000000 --- a/tests/archive/test_pv64_fmha_v.py +++ /dev/null @@ -1,258 +0,0 @@ -""" -Test (128,64) PV with FMHA-style V reconstruction. -V storage is (n, hd) row-major. Inside CuTe, we reconstruct as (hd, n, 1) MN-major. -This makes PV B operand see N=hd=64, K=n=128 with all 128 K-values. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - -HEAD_DIM = 64 - -class Pv64FmhaV: - def __init__(self): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0,1,2,3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.num_c_stage = 2 - self.num_ab_stage = 1; self.num_acc_stage = 1 - - def _setup(self, qk_mma, pv_mma): - qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (128, 128, qk_ik * 4) - pv_ik = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (128, HEAD_DIM, pv_ik * (128 // pv_ik)) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0]//cute.size(qk_mma.thr_id.shape), HEAD_DIM, self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - qk_thr = qk_mma.get_slice(0); qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - pv_thr = pv_mma.get_slice(0); pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0; self.tmem_p0_offset = 32 - self.tmem_o0_offset = find_tmem_tensor_col_offset(tOtO) - tCS = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage)) - tCO = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCS, tCO], arch="sm_100") - a_s = cute.slice_(self.a_smem_s,(None,None,None,0)); b_s = cute.slice_(self.b_smem_s,(None,None,None,0)) - v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - self.num_tma_load_bytes = (cute.size_in_bytes(self.q_dtype,a_s)+cute.size_in_bytes(self.q_dtype,b_s)+cute.size_in_bytes(self.q_dtype,v_s))*cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - # FMHA-style V reconstruction: logical shape (HEAD_DIM, s_k, 1) MN-major - # Physical V storage is (n, hd) row-major. Reinterpret as (hd, n, 1) with stride (1, hd, hd*n). - s_k = 128 # sequence length (QK N dimension) - d = HEAD_DIM - v_fmha = cute.make_tensor( - v.iterator, - cute.make_layout( - (d, s_k, 1), - stride=(1, d, d * s_k), - ), - ) - self.v_major = LayoutEnum.from_tensor(v_fmha).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, self.a_major, self.b_major, self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, self.qk_acc_dtype, self.cta_group, (128,HEAD_DIM), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - q_s = cute.slice_(self.a_smem_s,(None,None,None,0)); k_s = cute.slice_(self.b_smem_s,(None,None,None,0)) - v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id),q,q_s,self.mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id),k,k_s,self.mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - # TMA for V uses v_fmha (reconstructed layout), not original v - tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v_fmha,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape) - epi_s = cute.select(self.c_smem_s,mode=[0,1]) - tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile) - self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.a_smem_s,self.b_smem_s,self.v_smem_s,self.p_tmem_s,self.c_smem_s,self.epi_tile).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx,_,_ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage*2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage*2] - tmem_dealloc: cutlass.Int64; holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - ab_p,ab_c = pipeline.PipelineTmaUmma.create(barrier_storage=st.ab_bar.data_ptr(),num_stages=self.num_ab_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.num_tma_load_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - mma_si_prod,mma_si_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.mma_si_bar.data_ptr(),num_stages=1,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.epilogue_warp_id))).make_participants() - softmax_done_bar = pipeline.NamedBarrier(barrier_id=3, num_threads=32 + 32*len(self.epilogue_warp_id)) - acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(),num_stages=self.num_acc_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,len(self.epilogue_warp_id)),cta_layout_vmnk=cl_vmnk,defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2,num_threads=32*len((self.mma_warp_id,*self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr,barrier_for_retrieve=tmem_bar,allocator_warp_id=self.epilogue_warp_id[0],is_two_cta=cute.size(qk_mma.thr_id.shape)==2,two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk,is_relaxed=True) - sQ = smem.allocate_tensor(element_type=self.q_dtype,layout=a_smem_s.outer,byte_alignment=128,swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype,layout=b_smem_s.outer,byte_alignment=128,swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype,layout=v_smem_s.outer,byte_alignment=128,swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype,layout=c_smem_s.outer,byte_alignment=128,swizzle=c_smem_s.inner) - gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None)) - gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None)) - gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None)) - gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - tCgV = pv_thr.partition_B(gV); tCgC = pv_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,0,None,0)).shape) - tAsQ,tAgQ = cpasync.tma_partition(tma_q,0,a_lay,cute.group_modes(sQ,0,3),cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape) - tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3)) - tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)]; tVgV = tVgV[(None,0,None,0)] - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - tStS0 = cute.make_tensor(tStS.iterator+self.tmem_s0_offset,tStS.layout) - pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - tOtO0 = cute.make_tensor(tOtO.iterator+self.tmem_o0_offset,tOtO.layout) - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None,None,None,0)] - tOrP0 = cute.make_tensor(tOrP.iterator+self.qk_acc_dtype.width//self.q_dtype.width*self.tmem_p0_offset,tOrP.layout) - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_as,self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_as,self.num_acc_stage)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # TMA LOAD - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt,unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q,tAgQ[(None,h.count)],tAsQ[(None,h.index)],tma_bar_ptr=h.barrier) - cute.copy(tma_k,tBgK[(None,h.count)],tBsK[(None,h.index)],tma_bar_ptr=h.barrier) - cute.copy(tma_v,tVgV[(None,h.count)],tVsV[(None,h.index)],tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else "FAIL"}') - if cos < 0.99: - print(f' out[0,:4]={out[0,:4].tolist()} ref[0,:4]={ref[0,:4].tolist()}') - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_pv64_kmajor_v.py b/tests/archive/test_pv64_kmajor_v.py deleted file mode 100644 index ad5392b0..00000000 --- a/tests/archive/test_pv64_kmajor_v.py +++ /dev/null @@ -1,238 +0,0 @@ -""" -Test (128,64) PV with V=(64,128) K-major (row-major contiguous). -PV MMA uses OperandMajorMode.K for V (auto-detected). -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - -HEAD_DIM = 64 - -class Pv64KMajorV: - def __init__(self): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0,1,2,3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.num_c_stage = 2 - self.num_ab_stage = 1; self.num_acc_stage = 1 - - def _setup(self, qk_mma, pv_mma): - qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (128, 128, qk_ik * 4) - pv_ik = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (128, HEAD_DIM, pv_ik * (128 // pv_ik)) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0]//cute.size(qk_mma.thr_id.shape), HEAD_DIM, self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - qk_thr = qk_mma.get_slice(0); qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - pv_thr = pv_mma.get_slice(0); pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0; self.tmem_p0_offset = 32 - self.tmem_o0_offset = find_tmem_tensor_col_offset(tOtO) - tCS = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage)) - tCO = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCS, tCO], arch="sm_100") - a_s = cute.slice_(self.a_smem_s,(None,None,None,0)); b_s = cute.slice_(self.b_smem_s,(None,None,None,0)) - v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - self.num_tma_load_bytes = (cute.size_in_bytes(self.q_dtype,a_s)+cute.size_in_bytes(self.q_dtype,b_s)+cute.size_in_bytes(self.q_dtype,v_s))*cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - # V is (64, 128) K-major — auto-detect - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, self.a_major, self.b_major, self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, self.qk_acc_dtype, self.cta_group, (128,HEAD_DIM), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - q_s = cute.slice_(self.a_smem_s,(None,None,None,0)); k_s = cute.slice_(self.b_smem_s,(None,None,None,0)) - v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id),q,q_s,self.mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id),k,k_s,self.mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape) - epi_s = cute.select(self.c_smem_s,mode=[0,1]) - tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile) - self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.a_smem_s,self.b_smem_s,self.v_smem_s,self.p_tmem_s,self.c_smem_s,self.epi_tile).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx,_,_ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage*2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage*2] - tmem_dealloc: cutlass.Int64; holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - ab_p,ab_c = pipeline.PipelineTmaUmma.create(barrier_storage=st.ab_bar.data_ptr(),num_stages=self.num_ab_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.num_tma_load_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - mma_si_prod,mma_si_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.mma_si_bar.data_ptr(),num_stages=1,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.epilogue_warp_id))).make_participants() - softmax_done_bar = pipeline.NamedBarrier(barrier_id=3, num_threads=32 + 32*len(self.epilogue_warp_id)) - acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(),num_stages=self.num_acc_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,len(self.epilogue_warp_id)),cta_layout_vmnk=cl_vmnk,defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2,num_threads=32*len((self.mma_warp_id,*self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr,barrier_for_retrieve=tmem_bar,allocator_warp_id=self.epilogue_warp_id[0],is_two_cta=cute.size(qk_mma.thr_id.shape)==2,two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk,is_relaxed=True) - sQ = smem.allocate_tensor(element_type=self.q_dtype,layout=a_smem_s.outer,byte_alignment=128,swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype,layout=b_smem_s.outer,byte_alignment=128,swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype,layout=v_smem_s.outer,byte_alignment=128,swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype,layout=c_smem_s.outer,byte_alignment=128,swizzle=c_smem_s.inner) - gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None)) - gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None)) - gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None)) - gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - tCgV = pv_thr.partition_B(gV); tCgC = pv_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,0,None,0)).shape) - tAsQ,tAgQ = cpasync.tma_partition(tma_q,0,a_lay,cute.group_modes(sQ,0,3),cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape) - tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3)) - tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)]; tVgV = tVgV[(None,0,None,0)] - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - tStS0 = cute.make_tensor(tStS.iterator+self.tmem_s0_offset,tStS.layout) - pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - tOtO0 = cute.make_tensor(tOtO.iterator+self.tmem_o0_offset,tOtO.layout) - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None,None,None,0)] - tOrP0 = cute.make_tensor(tOrP.iterator+self.qk_acc_dtype.width//self.q_dtype.width*self.tmem_p0_offset,tOrP.layout) - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_as,self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_as,self.num_acc_stage)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt,unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q,tAgQ[(None,h.count)],tAsQ[(None,h.index)],tma_bar_ptr=h.barrier) - cute.copy(tma_k,tBgK[(None,h.count)],tBsK[(None,h.index)],tma_bar_ptr=h.barrier) - cute.copy(tma_v,tVgV[(None,h.count)],tVsV[(None,h.index)],tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else "FAIL"}') - if cos < 0.99: - print(f' out[0,:4]={out[0,:4].tolist()} ref[0,:4]={ref[0,:4].tolist()}') - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_pv64_no_softmax.py b/tests/archive/test_pv64_no_softmax.py deleted file mode 100644 index 2c3baff7..00000000 --- a/tests/archive/test_pv64_no_softmax.py +++ /dev/null @@ -1,222 +0,0 @@ -""" -Test (128,64) PV WITHOUT softmax. -QK writes S to TMEM. Then PV uses S directly as P (no BF16 conversion). -If the C-fragment store path works, PV should read S and produce output. -If PV reads zeros, the P/A alias is broken for (128,64). -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - -HEAD_DIM = 64 - -class Pv64NoSoftmax: - def __init__(self): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0,1,2,3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.num_c_stage = 2 - self.num_ab_stage = 1; self.num_acc_stage = 1 - - def _setup(self, qk_mma, pv_mma): - qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (128, 128, qk_ik * 4) - pv_ik = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (128, HEAD_DIM, pv_ik * (128 // pv_ik)) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0]//cute.size(qk_mma.thr_id.shape), HEAD_DIM, self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - qk_thr = qk_mma.get_slice(0); qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - pv_thr = pv_mma.get_slice(0); pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - self.tmem_s0_offset = 0; self.tmem_p0_offset =32 - self.tmem_o0_offset = find_tmem_tensor_col_offset(tOtO) - tCS = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage)) - tCO = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCS, tCO], arch="sm_100") - a_s = cute.slice_(self.a_smem_s,(None,None,None,0)); b_s = cute.slice_(self.b_smem_s,(None,None,None,0)) - v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - self.num_tma_load_bytes = (cute.size_in_bytes(self.q_dtype,a_s)+cute.size_in_bytes(self.q_dtype,b_s)+cute.size_in_bytes(self.q_dtype,v_s))*cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, self.a_major, self.b_major, self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, self.qk_acc_dtype, self.cta_group, (128,HEAD_DIM), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - q_s = cute.slice_(self.a_smem_s,(None,None,None,0)); k_s = cute.slice_(self.b_smem_s,(None,None,None,0)) - v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id),q,q_s,self.mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id),k,k_s,self.mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape) - epi_s = cute.select(self.c_smem_s,mode=[0,1]) - tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile) - self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.a_smem_s,self.b_smem_s,self.v_smem_s,self.p_tmem_s,self.c_smem_s,self.epi_tile).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx,_,_ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage*2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage*2] - tmem_dealloc: cutlass.Int64; holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - ab_p,ab_c = pipeline.PipelineTmaUmma.create(barrier_storage=st.ab_bar.data_ptr(),num_stages=self.num_ab_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.num_tma_load_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - mma_si_prod,mma_si_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.mma_si_bar.data_ptr(),num_stages=1,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.epilogue_warp_id))).make_participants() - acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(),num_stages=self.num_acc_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,len(self.epilogue_warp_id)),cta_layout_vmnk=cl_vmnk,defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2,num_threads=32*len((self.mma_warp_id,*self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr,barrier_for_retrieve=tmem_bar,allocator_warp_id=self.epilogue_warp_id[0],is_two_cta=cute.size(qk_mma.thr_id.shape)==2,two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk,is_relaxed=True) - sQ = smem.allocate_tensor(element_type=self.q_dtype,layout=a_smem_s.outer,byte_alignment=128,swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype,layout=b_smem_s.outer,byte_alignment=128,swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype,layout=v_smem_s.outer,byte_alignment=128,swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype,layout=c_smem_s.outer,byte_alignment=128,swizzle=c_smem_s.inner) - gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None)) - gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None)) - gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None)) - gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - tCgV = pv_thr.partition_B(gV); tCgC = pv_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,0,None,0)).shape) - tAsQ,tAgQ = cpasync.tma_partition(tma_q,0,a_lay,cute.group_modes(sQ,0,3),cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape) - tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3)) - tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)]; tVgV = tVgV[(None,0,None,0)] - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - tStS0 = cute.make_tensor(tStS.iterator+self.tmem_s0_offset,tStS.layout) - pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - tOtO0 = cute.make_tensor(tOtO.iterator+self.tmem_o0_offset,tOtO.layout) - # PV reads from S offset (no separate P, no softmax) - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None,None,None,0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) # reads from S offset = 0 - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # TMA LOAD - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt,unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q,tAgQ[(None,h.count)],tAsQ[(None,h.index)],tma_bar_ptr=h.barrier) - cute.copy(tma_k,tBgK[(None,h.count)],tBsK[(None,h.index)],tma_bar_ptr=h.barrier) - cute.copy(tma_v,tVgV[(None,h.count)],tVsV[(None,h.index)],tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1 0: - print('P/A alias works for (128,64) — PV reads non-zero data from TMEM') - else: - print('P/A alias BROKEN for (128,64) — PV reads all zeros from TMEM') - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_pv64_nosoftmax_fmha_v.py b/tests/archive/test_pv64_nosoftmax_fmha_v.py deleted file mode 100644 index 62767bf6..00000000 --- a/tests/archive/test_pv64_nosoftmax_fmha_v.py +++ /dev/null @@ -1,207 +0,0 @@ -""" -Test (128,64) PV with FMHA-style V reconstruction, no softmax. -V storage is (n, hd) row-major. Reconstructed inside CuTe as (hd, n, 1) MN-major. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - -HEAD_DIM = 64 - -class Pv64NoSoftmaxFmhaV: - def __init__(self): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0,1,2,3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.num_c_stage = 2 - self.num_ab_stage = 1; self.num_acc_stage = 1 - - def _setup(self, qk_mma, pv_mma): - qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (128, 128, qk_ik * 4) - pv_ik = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (128, HEAD_DIM, pv_ik * (128 // pv_ik)) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0]//cute.size(qk_mma.thr_id.shape), HEAD_DIM, self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - qk_thr = qk_mma.get_slice(0); qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - pv_thr = pv_mma.get_slice(0); pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - self.tmem_s0_offset = 0; self.tmem_p0_offset = 0 - self.tmem_o0_offset = find_tmem_tensor_col_offset(tOtO) - tCS = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage)) - tCO = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCS, tCO], arch="sm_100") - a_s = cute.slice_(self.a_smem_s,(None,None,None,0)); b_s = cute.slice_(self.b_smem_s,(None,None,None,0)) - v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - self.num_tma_load_bytes = (cute.size_in_bytes(self.q_dtype,a_s)+cute.size_in_bytes(self.q_dtype,b_s)+cute.size_in_bytes(self.q_dtype,v_s))*cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - # FMHA-style V: logical (HEAD_DIM, s_k, 1) with stride (1, HEAD_DIM, HEAD_DIM*s_k) - v_fmha = cute.make_tensor( - v.iterator, - cute.make_layout( - (HEAD_DIM, 128, 1), - stride=(1, HEAD_DIM, HEAD_DIM * 128), - ), - ) - self.v_major = LayoutEnum.from_tensor(v_fmha).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, self.a_major, self.b_major, self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, self.qk_acc_dtype, self.cta_group, (128,HEAD_DIM), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - q_s = cute.slice_(self.a_smem_s,(None,None,None,0)); k_s = cute.slice_(self.b_smem_s,(None,None,None,0)) - v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id),q,q_s,self.mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id),k,k_s,self.mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v_fmha,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape) - epi_s = cute.select(self.c_smem_s,mode=[0,1]) - tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile) - self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.a_smem_s,self.b_smem_s,self.v_smem_s,self.p_tmem_s,self.c_smem_s,self.epi_tile).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx,_,_ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage*2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage*2] - tmem_dealloc: cutlass.Int64; holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - ab_p,ab_c = pipeline.PipelineTmaUmma.create(barrier_storage=st.ab_bar.data_ptr(),num_stages=self.num_ab_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.num_tma_load_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - mma_si_prod,mma_si_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.mma_si_bar.data_ptr(),num_stages=1,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.epilogue_warp_id))).make_participants() - acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(),num_stages=self.num_acc_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,len(self.epilogue_warp_id)),cta_layout_vmnk=cl_vmnk,defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2,num_threads=32*len((self.mma_warp_id,*self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr,barrier_for_retrieve=tmem_bar,allocator_warp_id=self.epilogue_warp_id[0],is_two_cta=cute.size(qk_mma.thr_id.shape)==2,two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk,is_relaxed=True) - sQ = smem.allocate_tensor(element_type=self.q_dtype,layout=a_smem_s.outer,byte_alignment=128,swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype,layout=b_smem_s.outer,byte_alignment=128,swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype,layout=v_smem_s.outer,byte_alignment=128,swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype,layout=c_smem_s.outer,byte_alignment=128,swizzle=c_smem_s.inner) - gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None)) - gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None)) - gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None)) - gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - tCgV = pv_thr.partition_B(gV); tCgC = pv_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,0,None,0)).shape) - tAsQ,tAgQ = cpasync.tma_partition(tma_q,0,a_lay,cute.group_modes(sQ,0,3),cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape) - tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3)) - tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)]; tVgV = tVgV[(None,0,None,0)] - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - tStS0 = cute.make_tensor(tStS.iterator+self.tmem_s0_offset,tStS.layout) - pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - tOtO0 = cute.make_tensor(tOtO.iterator+self.tmem_o0_offset,tOtO.layout) - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None,None,None,0)] - tOrP0 = cute.make_tensor(tOrP.iterator+self.qk_acc_dtype.width//self.q_dtype.width*self.tmem_s0_offset,tOrP.layout) - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt,unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q,tAgQ[(None,h.count)],tAsQ[(None,h.index)],tma_bar_ptr=h.barrier) - cute.copy(tma_k,tBgK[(None,h.count)],tBsK[(None,h.index)],tma_bar_ptr=h.barrier) - cute.copy(tma_v,tVgV[(None,h.count)],tVsV[(None,h.index)],tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_pv_mma_mn_major.py b/tests/archive/test_pv_mma_mn_major.py deleted file mode 100644 index 066b8877..00000000 --- a/tests/archive/test_pv_mma_mn_major.py +++ /dev/null @@ -1,303 +0,0 @@ -""" -Isolated test for Bug 1: PV MMA with V MN-major. - -Only tests the PV MMA (P@V) with V as MN-major B-operand. -No QK MMA, no identity softmax, no pipeline complexity. -P comes from TMEM (a_source=TMEM), V comes from SMEM (b from TMA load). - -Architecture: - - TMA load V into SMEM - - P pre-populated in TMEM (via small QK MMA or direct write) - - PV MMA: P @ V → O in TMEM - - Epilogue: TMEM → GMEM - -For simplicity, P is computed via a QK MMA first (Q@K^T → P in TMEM), -then PV MMA uses P from TMEM. No softmax — identity pass-through. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - - -class PvMmaTest: - def __init__(self, mma_tiler_mn, use_2cta_instrs=False, use_tma_store=True): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16 - self.use_2cta_instrs = use_2cta_instrs; self.use_tma_store = use_tma_store - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE - self.mma_warp_id = 0 - self.tma_warp_id = 1 - self.threads_per_cta = 64 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - self.pv_mma_tiler = (self.qk_mma_tiler[0], self.qk_mma_tiler[2], self.qk_mma_tiler[1]) - self.mma_tiler = self.qk_mma_tiler - print(f"[pv_test] qk_mma_tiler = {self.qk_mma_tiler}") - print(f"[pv_test] pv_mma_tiler = {self.pv_mma_tiler}") - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - - # Compute epilogue tile from PV output (not QK) - cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], - self.qk_mma_tiler[2], - ) - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - cta_tile_shape_mnk, self.use_2cta_instrs, self.c_layout, self.o_dtype) - - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.a_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.b_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.b_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 0 # P = S (identity softmax, same TMEM) - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.a_dtype, a_smem) + cute.size_in_bytes(self.b_dtype, b_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q: cute.Tensor, k: cute.Tensor, v: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - self.a_dtype = q.element_type; self.b_dtype = k.element_type; self.c_dtype = c.element_type - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - print(f"[pv_test] a_major (Q) = {self.a_major}") - print(f"[pv_test] b_major (K) = {self.b_major}") - print(f"[pv_test] v_major (V) = {self.v_major}") - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - # BUG 1 FIX: PV MMA uses V's MN-major mode - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.b_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32), - tx_count=self.num_tmama_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, num_threads=64) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=0, is_two_cta=False, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.a_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.b_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.b_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - # P from S TMEM — same location, MMA A-operand for PV - tP = cute.make_tensor(tStS.iterator, self.p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = tOrP # P is at same TMEM offset as S (identity softmax) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # WARP 1: TMA load - if warp_idx == self.tma_warp_id: - tmem.wait_for_alloc() - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1 O[m, n] - # This is P @ V^T in matrix notation - # So reference: Q@K^T @ V^T where V^T is (128, 64) - ref = qf @ kf.T @ vf.T # (128,128) @ (128,64) = (128,64) - - import cutlass.torch as ct - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - kernel = PvMmaTest(mma_tiler_mn=(128, 128), use_2cta_instrs=False, use_tma_store=True) - print('Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - print('Running...', flush=True) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - out = c[:,:,0].float() - cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - max_err = (out - ref).abs().max().item() - print('PV MMA test (V MN-major, no softmax):') - print(' Cosine: {:.6f}, Max error: {:.6f}'.format(cos, max_err)) - print(' {}'.format('PASS' if cos >= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_quick_rand.py b/tests/archive/test_quick_rand.py deleted file mode 100644 index e69203b5..00000000 --- a/tests/archive/test_quick_rand.py +++ /dev/null @@ -1,36 +0,0 @@ -"""Quick random test at N=32 K=32 — if cosize fix works, cosine should be ~1.0""" -import torch, sys -sys.path.insert(0, 'src') -from nvfp4_megamoe_kernel.cutlass_nvfp4_gemm.kernel import cutlass_nvfp4_blockscaled_gemm -from nvfp4_megamoe_kernel.nvfp4_mega_moe import _quantize_to_e2m1, _E2M1_MAGNITUDES - -torch.manual_seed(42) -device = "cuda" -M, N, K = 1, 32, 32 - -x_bf16 = torch.randn(M, K, dtype=torch.bfloat16, device=device) * 2.0 -w_bf16 = torch.randn(K, N, dtype=torch.bfloat16, device=device) * 0.5 - -x_fp4, x_sf = _quantize_to_e2m1(x_bf16.float()) -w_fp4, w_sf = _quantize_to_e2m1(w_bf16.T.float()) -w_fp4 = w_fp4.T; w_sf = w_sf.T - -# Dequant -x_u8 = x_fp4.view(torch.uint8) -lo = (x_u8 & 0x0F).long(); hi = ((x_u8 >> 4) & 0x0F).long() -x_nib = torch.stack([lo, hi], dim=-1).reshape(M, -1) -x_deq = ((x_nib >> 3).float() * -2 + 1) * _E2M1_MAGNITUDES.to(device)[(x_nib & 0x07)] -x_recon = (x_deq * x_sf.to(torch.float32).repeat_interleave(16, dim=-1)).to(torch.bfloat16) - -w_u8 = w_fp4.view(torch.uint8) -wlo = (w_u8 & 0x0F).long(); whi = ((w_u8 >> 4) & 0x0F).long() -w_nib = torch.stack([wlo, whi], dim=-1).reshape(w_u8.shape[0]*2, w_u8.shape[1]) -w_deq = ((w_nib >> 3).float() * -2 + 1) * _E2M1_MAGNITUDES.to(device)[(w_nib & 0x07)] -w_recon = (w_deq * w_sf.to(torch.float32).repeat_interleave(16, dim=0)).to(torch.bfloat16) - -quant_ref = torch.nn.functional.linear(x_recon, w_recon.T) -nvfp4_out = cutlass_nvfp4_blockscaled_gemm(x_fp4, x_sf, w_fp4, w_sf, M, N, K, alpha=1.0) -cos = torch.nn.functional.cosine_similarity(nvfp4_out.float(), quant_ref.float(), dim=-1).mean().item() -print(f"M={M} N={N} K={K} cosine={cos:.6f}") -print(f"NVFP4 first 8: {nvfp4_out[0,:8].tolist()}") -print(f"REF first 8: {quant_ref[0,:8].tolist()}") diff --git a/tests/archive/test_recast_minimal.py b/tests/archive/test_recast_minimal.py deleted file mode 100644 index 3a20d711..00000000 --- a/tests/archive/test_recast_minimal.py +++ /dev/null @@ -1,237 +0,0 @@ -"""Absolute minimal: ld FP32 from S0, st FP32 to S1, epi reads S1. -No recast, no BF16, no packing. Pure FP32 copy between TMEM regions.""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - -class RecastMinimal: - def __init__(self, mma_tiler_mn): - self.qk_acc_dtype = Float32; self.q_dtype = BFloat16; self.o_dtype = BFloat16 - self.c_dtype = BFloat16; self.acc_dtype = Float32 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2; self.use_2cta_instrs = False - self.epilog_sync_bar_id = 1 - - def _setup(self, qk_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - self.mma_tiler = self.qk_mma_tiler - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - c_layout = LayoutEnum.ROW_MAJOR; self.c_layout = c_layout - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, False, c_layout, self.o_dtype) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, c_layout, self.epi_tile, 2) - self.num_ab_stage = 1; self.num_acc_stage = 1 - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - self.s_cols = find_tmem_tensor_col_offset(tStS) - self.tmem_alloc_cols = self.s_cols * 2 - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, - LayoutEnum.from_tensor(a).mma_major_mode(), - LayoutEnum.from_tensor(b).mma_major_mode(), - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, - tcgen05.OperandSource.SMEM) - self._setup(qk_mma) - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - tma_a, tma_ta = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - a, a_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_b, tma_tb = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - b, b_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - self._kernel(qk_mma, tma_a, tma_ta, tma_b, tma_tb, tma_c, tma_tc, - self.cluster_layout_vmnk, self.a_smem_s, self.b_smem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, tma_a, mA, tma_b, mB, tma_c, mC, cl_vmnk, - a_smem_s, b_smem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_a); cpasync.prefetch_descriptor(tma_b); cpasync.prefetch_descriptor(tma_c) - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, len(self.epilogue_warp_id)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=False, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - sA = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sB = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - gA = cute.local_tile(mA, cute.slice_(self.mma_tiler, (None,0,None)), (None,None,None)) - gB = cute.local_tile(mB, cute.slice_(self.mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gA, mode=[3]) - qk_thr = qk_mma.get_slice(0) - tCgA = qk_thr.partition_A(gA); tCgB = qk_thr.partition_B(gB); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsA, tAgA = cpasync.tma_partition(tma_a, 0, a_lay, cute.group_modes(sA,0,3), cute.group_modes(tCgA,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsB, tBgB = cpasync.tma_partition(tma_b, 0, b_lay, cute.group_modes(sB,0,3), cute.group_modes(tCgB,0,3)) - tAgA = tAgA[(None,0,None,0)]; tBgB = tBgB[(None,0,None,0)] - tCrA = qk_mma.make_fragment_A(sA); tCrB = qk_mma.make_fragment_B(sB) - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator, tStS.layout) - tStS1 = cute.make_tensor(tStS.iterator + self.s_cols, tStS.layout) - - # LD and ST on same layout - tmem_ld = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tmem_st = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_ld = tcgen05.make_tmem_copy(tmem_ld, tStS0) - tiled_st = tcgen05.make_tmem_copy(tmem_st, tStS1) - sfw = tidx % (32 * len(self.epilogue_warp_id)) - thr_ld = tiled_ld.get_slice(sfw) - thr_st = tiled_st.get_slice(sfw) - tLdS = thr_ld.partition_S(tStS0) - tStS = thr_st.partition_D(tStS1) - cS_id = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS_id) - tLdcS = thr_ld.partition_D(tScS) - tStcS = thr_st.partition_S(tScS) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, 1)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_a, tAgA[(None,h.count)], tAsA[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_b, tBgB[(None,h.count)], tBsB[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_ref_minimal.py b/tests/archive/test_ref_minimal.py deleted file mode 100644 index 937eb2ee..00000000 --- a/tests/archive/test_ref_minimal.py +++ /dev/null @@ -1,41 +0,0 @@ -"""Minimal test: CUTLASS reference FMHA, n=256 only.""" -import sys -sys.path.insert(0, '/root/cutlass/examples/python/CuTeDSL') - -import torch, math, cutlass, cutlass.cute as cute, cuda.bindings.driver as cuda -from cute.blackwell.kernel.attention.fmha.fmha import BlackwellFusedMultiHeadAttentionForward, FMHA_OperandMajorMode - -HEAD_DIM = 64 -n = 256 -torch.manual_seed(42) -m = 128; batch = 1 -q = torch.randn(batch, 1, m, HEAD_DIM, dtype=torch.bfloat16, device='cuda') -k = torch.randn(batch, 1, n, HEAD_DIM, dtype=torch.bfloat16, device='cuda') -v = torch.randn(batch, 1, n, HEAD_DIM, dtype=torch.bfloat16, device='cuda') -c = torch.zeros(batch, 1, m, HEAD_DIM, dtype=torch.bfloat16, device='cuda') - -qf = q[0,0].float(); kf = k[0,0].float(); vf = v[0,0].float() -scale = 1.0/math.sqrt(HEAD_DIM) -ref = torch.softmax(qf @ kf.T * scale, dim=-1) @ vf - -stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) -kernel = BlackwellFusedMultiHeadAttentionForward( - q_major_mode=FMHA_OperandMajorMode.K, k_major_mode=FMHA_OperandMajorMode.K, - v_major_mode=FMHA_OperandMajorMode.MN, o_major_mode=FMHA_OperandMajorMode.K, - q_head_dim=HEAD_DIM, kv_head_dim=HEAD_DIM, num_q_heads=1, num_kv_heads=1, - q_dtype=cutlass.BFloat16, k_dtype=cutlass.BFloat16, v_dtype=cutlass.BFloat16, - o_dtype=cutlass.BFloat16, acc_dtype=cutlass.Float32, epilogue_dtype=cutlass.Float32, - use_2cta_instrs=False, -) -print(f'n={n}: Compiling reference FMHA...', flush=True) -try: - kernel.run(q, k, v, c, stream) - torch.cuda.synchronize() - out = c[0,0].float() - cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - print(f'Reference FMHA n={n} (2 tiles): cos {cos:.6f} {"PASS" if cos >= 0.99 else "FAIL"}') - if cos < 0.99: - print(f' out[0,:4]={out[0,:4].tolist()}') - print(f' ref[0,:4]={ref[0,:4].tolist()}') -except Exception as e: - import traceback; traceback.print_exc() diff --git a/tests/archive/test_reference_fmha.py b/tests/archive/test_reference_fmha.py deleted file mode 100644 index c4f3da3c..00000000 --- a/tests/archive/test_reference_fmha.py +++ /dev/null @@ -1,80 +0,0 @@ -"""Test the CUTLASS reference Blackwell FMHA on the B200. -Does it actually work multi-tile?""" -import sys -sys.path.insert(0, '/root/cutlass/examples/python/CuTeDSL') - -import torch -import math -import cutlass -import cutlass.cute as cute -import cutlass.torch as ct -import cuda.bindings.driver as cuda - -from cute.blackwell.kernel.attention.fmha.fmha import ( - BlackwellFusedMultiHeadAttentionForward, - FusedMask, FusedMaskScale, FMHA_OperandMajorMode -) -from cutlass.utils import LayoutEnum - -def test_reference(): - HEAD_DIM = 64 - for n in [128, 256, 512]: - torch.manual_seed(42) - m = 128 - batch = 1 - - q = torch.randn(batch, 1, m, HEAD_DIM, dtype=torch.bfloat16, device='cuda') - k = torch.randn(batch, 1, n, HEAD_DIM, dtype=torch.bfloat16, device='cuda') - v = torch.randn(batch, 1, n, HEAD_DIM, dtype=torch.bfloat16, device='cuda') - c = torch.zeros(batch, 1, m, HEAD_DIM, dtype=torch.bfloat16, device='cuda') - - # Reference: PyTorch softmax attention - qf = q[0, 0].float() - kf = k[0, 0].float() - vf = v[0, 0].float() - scale = 1.0 / math.sqrt(HEAD_DIM) - attn = qf @ kf.T * scale - attn = torch.softmax(attn, dim=-1) - ref = attn @ vf - - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - kernel = BlackwellFusedMultiHeadAttentionForward( - q_major_mode=FMHA_OperandMajorMode.K, - k_major_mode=FMHA_OperandMajorMode.K, - v_major_mode=FMHA_OperandMajorMode.MN, - o_major_mode=FMHA_OperandMajorMode.K, - q_head_dim=HEAD_DIM, - kv_head_dim=HEAD_DIM, - num_q_heads=1, - num_kv_heads=1, - q_dtype=cutlass.BFloat16, - k_dtype=cutlass.BFloat16, - v_dtype=cutlass.BFloat16, - o_dtype=cutlass.BFloat16, - acc_dtype=cutlass.Float32, - epilogue_dtype=cutlass.Float32, - use_2cta_instrs=False, - ) - - print(f'n={n}: Compiling reference FMHA...', flush=True) - try: - result = kernel.run(q, k, v, c, stream) - torch.cuda.synchronize() - - out = c[0, 0].float() - cos = torch.nn.functional.cosine_similarity( - out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0) - ).item() - n_tiles = n // 128 - print(f'Reference FMHA n={n} ({n_tiles} tiles): cos {cos:.6f} {"PASS" if cos >= 0.99 else "FAIL"}') - if cos < 0.99: - print(f' out[0,:4]={out[0,:4].tolist()}') - print(f' ref[0,:4]={ref[0,:4].tolist()}') - except Exception as e: - import traceback - print(f'Reference FMHA n={n}: FAILED') - traceback.print_exc() - -if __name__ == '__main__': - test_reference() diff --git a/tests/archive/test_rope_kv_b200.py b/tests/archive/test_rope_kv_b200.py deleted file mode 100644 index ceb53710..00000000 --- a/tests/archive/test_rope_kv_b200.py +++ /dev/null @@ -1,152 +0,0 @@ -#!/usr/bin/env python3 -""" -Quick test: verify that applying RoPE to KV fixes the NaN issue. -Test the full attention pipeline with RoPE on both Q and KV. -""" -import sys, os, json, torch, torch.nn.functional as F, math -from safetensors import safe_open - -REPO = "/root/nvfp4-megamoe-kernel" -sys.path.insert(0, REPO) -MODEL = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4" -DEV = "cuda:0" - -H = 7168; NH = 128; HD = 512; NOPE = 448; ROPE = 64 -QL = 1536; OL = 1024; OG = 16; HPG = NH // OG -EPS = 1e-6; SCALE = HD ** -0.5 - -E2M1 = torch.tensor([0,.5,1.,1.5,2.,3.,4.,6.,-0,-.5,-1.,-1.5,-2.,-3.,-4.,-6.], dtype=torch.float32) - -_cache = {} -def P(k, wm, md): - if k in _cache: return _cache[k] - with safe_open(os.path.join(md, wm[k]), framework="pt") as f: - t = f.get_tensor(k) - _cache[k] = t - return t - -def dequant(w, sf, gs): - d = w.device; lut = E2M1.to(d) - lo = lut[(w & 0xF).long()]; hi = lut[((w >> 4) & 0xF).long()] - O, I2 = w.shape; I = I2*2 - u = torch.empty(O, I, dtype=torch.float32, device=d) - u[:,0::2] = lo; u[:,1::2] = hi - bs = sf.float().repeat_interleave(16, dim=1)[:O,:I] - return (u * bs * gs).to(torch.bfloat16) - -def rms(x, w, eps=1e-6): - v = x.float().pow(2).mean(-1, keepdim=True) - return (w.float() * (x * torch.rsqrt(v+eps)).float()).to(x.dtype) - -def make_runner(w, sf, gs_t, inf, outf): - from dsv4.layers.linear import Nvfp4Linear - fp4 = w.view(torch.float4_e2m1fn_x2).permute(1,0).contiguous() - s = sf.to(torch.float8_e4m3fn) if sf.dtype != torch.float8_e4m3fn else sf - s = s.permute(1,0).contiguous() - gs = gs_t.max().item() if gs_t.numel() > 1 else gs_t.item() - r = Nvfp4Linear(in_features=inf, out_features=outf, max_num_tokens=8192, device=str(w.device)) - r.fp4 = [fp4]; r.sf = [s]; r.gs = [gs] - r.finalize_weights(); r._ensure_initialized() - return r - -def build_cos_sin(max_pos=8192, rope_dim=ROPE): - half = rope_dim // 2 - inv_freq = 1.0 / (10000.0 ** (torch.arange(0, half, dtype=torch.float32) / half)) - freqs = torch.outer(torch.arange(max_pos, dtype=torch.float32), inv_freq) - return torch.cat([freqs.cos(), freqs.sin()], dim=-1) - -def apply_gptj_rope(x, positions, cos_sin, nope, rope): - if rope == 0 or x.numel() == 0: return x - half = rope // 2 - cos = cos_sin[positions, :half].to(x.dtype) - sin = cos_sin[positions, half:2*half].to(x.dtype) - if x.dim() == 3: cos = cos.unsqueeze(1); sin = sin.unsqueeze(1) - x_rope = x[..., nope:].clone() - even = x_rope[..., 0::2]; odd = x_rope[..., 1::2] - out = x.clone() - out[..., nope:][..., 0::2] = even * cos - odd * sin - out[..., nope:][..., 1::2] = even * sin + odd * cos - return out - -def full_sdpa_attention(q, kv, scale): - T, NH, HD = q.shape - q_2d = q.reshape(T * NH, HD) - kv_exp = kv.unsqueeze(1).expand(-1, NH, -1).contiguous() - k_2d = kv_exp.permute(1, 0, 2).unsqueeze(1).expand(NH, T, T, -1).contiguous().reshape(T * NH, T, HD) - v_2d = k_2d.clone() - scores = torch.matmul(q_2d.unsqueeze(1), k_2d.transpose(-1, -2)) * scale - qpos = torch.arange(T, device=q.device).unsqueeze(1).repeat(1, NH).reshape(T * NH) - kpos = torch.arange(T, device=q.device).unsqueeze(0) - causal = kpos <= qpos.unsqueeze(1) - scores = scores.squeeze(1).masked_fill(~causal, float('-inf')) - weights = F.softmax(scores.float(), dim=-1).to(q.dtype) - out = torch.matmul(weights.unsqueeze(1), v_2d).squeeze(1) - return out.reshape(T, NH, HD) - -torch.cuda.set_device(0) -torch.manual_seed(42) - -with open(os.path.join(MODEL, "model.safetensors.index.json")) as f: - wm = json.load(f)["weight_map"] -G = lambda k: P(k, wm, MODEL).to(DEV) - -p = "model.layers.0"; a = f"{p}.self_attn" -emb = G("model.embed_tokens.weight") -anorm = G(f"{p}.input_layernorm.weight") -qn = G(f"{a}.q_a_norm.weight"); kvn = G(f"{a}.kv_norm.weight") -woa = G(f"{a}.o_a_proj.weight") -qa_w = G(f"{a}.q_a_proj.weight"); qa_sf = G(f"{a}.q_a_proj.weight_scale"); qa_gs = G(f"{a}.q_a_proj.weight_scale_2") -qb_w = G(f"{a}.q_b_proj.weight"); qb_sf = G(f"{a}.q_b_proj.weight_scale"); qb_gs = G(f"{a}.q_b_proj.weight_scale_2") -kv_w = G(f"{a}.kv_proj.weight"); kv_sf = G(f"{a}.kv_proj.weight_scale"); kv_gs = G(f"{a}.kv_proj.weight_scale_2") -wob_w = G(f"{a}.o_b_proj.weight"); wob_sf = G(f"{a}.o_b_proj.weight_scale"); wob_gs = G(f"{a}.o_b_proj.weight_scale_2") - -r_qa = make_runner(qa_w, qa_sf, qa_gs, H, qa_w.shape[0]) -r_qb = make_runner(qb_w, qb_sf, qb_gs, QL, qb_w.shape[0]) -r_kv = make_runner(kv_w, kv_sf, kv_gs, H, kv_w.shape[0]) -r_wob = make_runner(wob_w, wob_sf, wob_gs, OG*OL, wob_w.shape[0]) - -cos_sin = build_cos_sin().to(DEV) -NT = 6 -token_ids = torch.tensor([1, 450, 8403, 315, 5413, 374], dtype=torch.long, device=DEV) -positions = torch.arange(NT, dtype=torch.int64, device=DEV) - -with torch.no_grad(): - hidden = emb[token_ids] - normed = rms(hidden, anorm, EPS) - - # Projections - qa = r_qa.run(normed); kv = r_kv.run(normed) - qa_n = rms(qa, qn, EPS); kv_n = rms(kv, kvn, EPS) - q = r_qb.run(qa_n).view(NT, NH, HD) - q_rope = apply_gptj_rope(q, positions, cos_sin, NOPE, ROPE) - - # Test 1: NO RoPE on KV (the bug) - print("--- Test 1: No RoPE on KV (BUG) ---") - o_no_rope = full_sdpa_attention(q_rope, kv_n, SCALE) - print(f" Output: amax={o_no_rope.amax():.4f} NaN={torch.isnan(o_no_rope).any()}") - - # Test 2: RoPE on KV (the fix) - print("--- Test 2: RoPE on KV (FIX) ---") - kv_rope = apply_gptj_rope(kv_n.unsqueeze(1), positions, cos_sin, NOPE, ROPE).squeeze(1) - o_with_rope = full_sdpa_attention(q_rope, kv_rope, SCALE) - print(f" Output: amax={o_with_rope.amax():.4f} NaN={torch.isnan(o_with_rope).any()}") - - # Test 3: Full pipeline - from dsv4.reference.csa_attention import apply_inv_gptj_rope - o_inv = apply_inv_gptj_rope(o_with_rope, positions, cos_sin, NOPE, ROPE) - o_grouped = o_inv.view(NT, OG, HPG * HD).permute(1, 0, 2) - woa_3d = woa.view(OG, OL, HPG * HD) - z = torch.bmm(o_grouped, woa_3d.transpose(1, 2)).permute(1, 0, 2).reshape(NT, OG * OL) - attn_out = r_wob.run(z) - - # LM head - fnorm_w = G("model.norm.weight"); lm_head = G("lm_head.weight") - x = hidden + attn_out - x_n = rms(x, fnorm_w, EPS) - logits = x_n @ lm_head.T - log_std = logits[-1].float().std().item() - top5 = torch.topk(logits[-1], 5) - print(f"\n--- Logits ---") - print(f" std={log_std:.4f} {'✅' if 0.5 < log_std < 50 else '❌'}") - print(f" top5 tokens: {top5.indices.tolist()}") - print(f" NaN in logits: {torch.isnan(logits).any()}") diff --git a/tests/archive/test_router.py b/tests/archive/test_router.py deleted file mode 100644 index 30f1423f..00000000 --- a/tests/archive/test_router.py +++ /dev/null @@ -1,217 +0,0 @@ -"""Unit tests for DSV4 Router — dense and hash modes. - -Test strategy: - Each kernel has a closed-form mathematical spec. The unit test computes - the spec in one expression in FP32 (PyTorch) and compares against the - kernel output. This is not "a PyTorch reference implementation" — it's - the math. Compare against that. No "ref/" file, no second implementation - drift, no two debug streams. - - The oracle is the same five lines of math as the kernel spec, written - declaratively. Compare against that. - -DO NOT RUN THESE TESTS — Carmine is actively testing Stage C. -Write the tests, commit them, they'll be run later. - -Tie-breaking: When two scores are exactly equal, torch.topk and the kernel -may pick different indices. Use the same tie-break rule: lower index wins -on ties. If the test fails on tie-breaking, fix the kernel, not the test. -""" - -import torch -import math - - -def test_fused_activation_topk(N=64, E=256, k=6, seed=42): - """Test the fused activation + top-k kernel against the math spec. - - Oracle: - logits = X @ W (FP32) - act = sqrt(softplus(logits)) - score = act + bias - ids = argtopk(score, k) with lower-index tie-break - raw_w = gather(act, ids) - topk_w = raw_w / sum(raw_w) * scaling - """ - torch.manual_seed(seed) - scaling = 2.5 - - logits = torch.randn(N, E, dtype=torch.float32, device='cuda') - e_bias = torch.randn(E, dtype=torch.float32, device='cuda') * 0.01 - - # Oracle — the math, one expression at a time - act = torch.sqrt(torch.nn.functional.softplus(logits)) - score = act + e_bias - # torch.topk tie-breaking: picks lower index on ties (matches our kernel) - topk_result = score.topk(k, dim=-1) - ids = topk_result.indices - raw_w = act.gather(-1, ids) - w = raw_w / raw_w.sum(-1, keepdim=True) * scaling - - # Kernel under test: - from dsv4.kernels.router._activation_topk import run_fused_activation_topk - out_w = torch.empty(N, k, dtype=torch.float32, device='cuda') - out_ids = torch.empty(N, k, dtype=torch.int32, device='cuda') - run_fused_activation_topk(logits, e_bias, scaling, k, out_w, out_ids) - - # Verify - assert (out_ids == ids).all(), f"top-k indices mismatch" - torch.testing.assert_close(out_w, w, atol=1e-4, rtol=1e-3) - - -def test_fused_activation_topk_decode_shapes(): - """Test the activation+topk kernel at decode-relevant N values.""" - for N in [1, 4, 16, 64]: - test_fused_activation_topk(N=N, E=256, k=6, seed=N) - - -def test_fused_activation_topk_pro_experts(): - """Test with 384 experts (Pro model).""" - test_fused_activation_topk(N=64, E=384, k=6, seed=123) - - -def test_hash_router(N=128, vocab_size=128000, k=6, num_experts=256, seed=42): - """Test the hash router against the math spec. - - Oracle: - topk_ids[n, h] = hash_lut[token_ids[n], h] - topk_w[n, h] = 1.0 / k - """ - torch.manual_seed(seed) - - # Build a random LUT - hash_lut = torch.randint(0, num_experts, (vocab_size, k), dtype=torch.int32, device='cuda') - token_ids = torch.randint(0, vocab_size, (N,), dtype=torch.int32, device='cuda') - - # Oracle — literally just indexing - expected_ids = hash_lut[token_ids] # [N, k] - expected_w = torch.full((N, k), 1.0 / k, dtype=torch.float32, device='cuda') - - # Kernel under test: - from dsv4.kernels.router import hash_router_dispatch - out_w = torch.empty(N, k, dtype=torch.float32, device='cuda') - out_ids = torch.empty(N, k, dtype=torch.int32, device='cuda') - hash_router_dispatch(token_ids, hash_lut, k, out_w, out_ids) - - assert (out_ids == expected_ids).all(), f"hash router IDs mismatch" - torch.testing.assert_close(out_w, expected_w, atol=1e-7, rtol=1e-7) - - -def test_hash_router_edge_cases(): - """Test hash router with N=1 and N=max_num_tokens.""" - test_hash_router(N=1, vocab_size=128000, k=6) - test_hash_router(N=8192, vocab_size=128000, k=6) - - -def test_topk_select(N=64, E=256, k=6, seed=42): - """Test standalone top-k selection against torch.topk. - - Oracle: - (values, indices) = score.topk(k, dim=-1) - Lower index wins on ties (torch.topk default). - """ - torch.manual_seed(seed) - scores = torch.randn(N, E, dtype=torch.float32, device='cuda') - - # Oracle - expected = scores.topk(k, dim=-1) - expected_ids = expected.indices - expected_values = expected.values - - # Kernel under test: - from dsv4.ops.topk import topk_select - out_values, out_ids = topk_select(scores, k) - - assert (out_ids == expected_ids).all(), f"top-k IDs mismatch" - torch.testing.assert_close(out_values, expected_values, atol=1e-6, rtol=1e-6) - - -def test_dense_router_decode(N=64, H=4096, E=256, k=6, seed=42): - """Test the full dense router (GEMM + activation + topk) against the spec. - - Oracle: - logits = (X.float() @ W.float()) - act = sqrt(softplus(logits)) - score = act + bias - ids = score.topk(k).indices - w = act.gather(-1, ids) - w = w / w.sum(-1, keepdim=True) * scaling - """ - torch.manual_seed(seed) - scaling = 2.5 - - X = torch.randn(N, H, dtype=torch.bfloat16, device='cuda') - W = torch.randn(H, E, dtype=torch.bfloat16, device='cuda') - bias = torch.randn(E, dtype=torch.float32, device='cuda') * 0.01 - - # Oracle — the math, in one expression, in FP32 - logits = (X.float() @ W.float()) - act = torch.sqrt(torch.nn.functional.softplus(logits)) - score = act + bias - ids = score.topk(k, dim=-1).indices - w = act.gather(-1, ids) - w = w / w.sum(-1, keepdim=True) * scaling - - # Kernel under test: - from dsv4.layers.router import Router - router = Router(H, E, k, scaling, mode='dense', max_num_tokens=N) - router.load_weights(W_gate=W, e_bias=bias) - router.finalize_weights() - out_w, out_ids = router(X) - - assert (out_ids == ids).all(), f"router IDs mismatch" - torch.testing.assert_close(out_w, w, atol=1e-3, rtol=1e-3) - - -def test_dense_router_decode_shapes(): - """Test dense router at decode-relevant N values.""" - for N in [1, 4, 16, 64]: - test_dense_router_decode(N=N, H=4096, E=256, k=6, seed=N) - - -def test_hash_router_via_router_class(): - """Test the Router class in hash mode.""" - vocab_size = 128000 - k = 6 - num_experts = 256 - N = 64 - - hash_lut = torch.randint(0, num_experts, (vocab_size, k), dtype=torch.int32, device='cuda') - token_ids = torch.randint(0, vocab_size, (N,), dtype=torch.int32, device='cuda') - - # Oracle - expected_ids = hash_lut[token_ids] - expected_w = torch.full((N, k), 1.0 / k, dtype=torch.float32, device='cuda') - - # Router class - from dsv4.layers.router import Router - router = Router( - hidden_size=4096, # not used in hash mode - num_experts=num_experts, - top_k=k, - mode='hash', - vocab_size=vocab_size, - max_num_tokens=N, - ) - router.load_weights(hash_lut=hash_lut) - router.finalize_weights() - out_w, out_ids = router(hidden_states=None, token_ids=token_ids) - - assert (out_ids == expected_ids).all(), f"hash router class IDs mismatch" - torch.testing.assert_close(out_w, expected_w, atol=1e-7, rtol=1e-7) - - -def test_softplus_numerical_stability(): - """Verify the numerically stable softplus matches the spec. - - For x = -100: softplus(x) ≈ 0, sqrt(softplus(x)) ≈ 0 - For x = 0: softplus(x) = log(2) ≈ 0.693, sqrt ≈ 0.832 - For x = 100: softplus(x) ≈ 100, sqrt(softplus(x)) ≈ 10 - """ - # This tests the Python math, not the kernel. It's a sanity check - # that the formula max(x,0) + log1p(exp(-|x|)) works correctly. - x = torch.tensor([-100.0, 0.0, 100.0], dtype=torch.float32) - sp = torch.nn.functional.softplus(x) - act = torch.sqrt(sp) - expected = torch.tensor([0.0, math.sqrt(math.log(2.0)), 10.0], dtype=torch.float32) - torch.testing.assert_close(act, expected, atol=1e-3, rtol=1e-3) diff --git a/tests/archive/test_runner_vs_pipeline.py b/tests/archive/test_runner_vs_pipeline.py deleted file mode 100644 index 26f6d25a..00000000 --- a/tests/archive/test_runner_vs_pipeline.py +++ /dev/null @@ -1,210 +0,0 @@ -#!/usr/bin/env python3 -""" -Test A: Compare moe_pipeline output vs Nvfp4MoE output. - -Uses the same weights and inputs. If they differ, the runner is broken. -Runs on the B200 host (not inside Docker): - source /root/nvfp4-megamoe-kernel/tests/.venv/bin/activate - python3 tests/test_runner_vs_pipeline.py -""" -import os, sys, json, torch -from safetensors import safe_open - -REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) -sys.path.insert(0, REPO_ROOT) - -from dsv4.reference.moe_pipeline import run_nvfp4_moe -from vllm.nvfp4_cutedsl import Nvfp4MoE -from dsv4.ops.quantize import ( - quantize_to_nvfp4, - quantize_weight_to_nvfp4, -) -from dsv4.ops.layouts import ( - make_b_k_major, - assemble_scales_3d_side, - compute_expert_offsets, -) - -MODEL_DIR = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4" -DEVICE = "cuda" -LAYER_IDX = 0 -E2M1_LUT = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, - -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0], dtype=torch.float32) - - -def find_shards(model_dir): - index_path = os.path.join(model_dir, "model.safetensors.index.json") - key_to_shard = {} - if os.path.exists(index_path): - with open(index_path) as f: - index = json.load(f) - for key, shard in index["weight_map"].items(): - key_to_shard[key] = os.path.join(model_dir, shard) - return key_to_shard - - -def load_layer_tensors(model_dir, layer_idx): - key_to_shard = find_shards(model_dir) - layer_prefix = f"layers.{layer_idx}." - tensors = {} - for key, shard in key_to_shard.items(): - norm_key = key.removeprefix("model.") - if not norm_key.startswith(layer_prefix): - continue - with safe_open(shard, framework="pt") as f: - if key in f.keys(): - tensors[norm_key] = f.get_tensor(key) - return tensors - - -def dequantize_nvfp4_weight(packed_uint8, scale_e4m3, global_scale): - device = packed_uint8.device - lut = E2M1_LUT.to(device) - lower = lut[(packed_uint8 & 0x0F).long()] - upper = lut[((packed_uint8 >> 4) & 0x0F).long()] - out_features = packed_uint8.shape[0] - in_features = packed_uint8.shape[1] * 2 - unpacked = torch.empty(out_features, in_features, dtype=torch.float32, device=device) - unpacked[:, 0::2] = lower - unpacked[:, 1::2] = upper - block_scale = scale_e4m3.float() - block_expanded = block_scale.repeat_interleave(16, dim=1)[:, :in_features] - return (unpacked * block_expanded * global_scale).to(torch.bfloat16) - - -def prepare_direct_weights(nvfp4_tensors, layer_idx, expert_indices, intermediate_size): - """Direct view-cast path (same as layertest).""" - l1_fp4, l1_sf, l1_gs = [], [], [] - l2_fp4, l2_sf, l2_gs = [], [], [] - for e in expert_indices: - gate_w = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.gate_proj.weight"].to(DEVICE) - up_w = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.up_proj.weight"].to(DEVICE) - gate_sf = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.gate_proj.weight_scale"].to(DEVICE) - up_sf = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.up_proj.weight_scale"].to(DEVICE) - gate_gs = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.gate_proj.weight_scale_2"].item() - up_gs = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.up_proj.weight_scale_2"].item() - fused_w = torch.cat([gate_w, up_w], dim=0) - fused_w_fp4 = fused_w.view(torch.float4_e2m1fn_x2).permute(1, 0).contiguous() - fused_sf = torch.cat([gate_sf, up_sf], dim=0).permute(1, 0).contiguous() - max_gs = max(gate_gs, up_gs) - if gate_gs != up_gs: - f32 = fused_sf.float() - f32[:, :intermediate_size] *= (gate_gs / max_gs) - f32[:, intermediate_size:] *= (up_gs / max_gs) - fused_sf = f32.to(torch.float8_e4m3fn) - l1_fp4.append(fused_w_fp4) - l1_sf.append(fused_sf) - l1_gs.append(max_gs) - down_w = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.down_proj.weight"].to(DEVICE) - down_sf = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.down_proj.weight_scale"].to(DEVICE) - down_gs = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.down_proj.weight_scale_2"].item() - l2_fp4.append(down_w.view(torch.float4_e2m1fn_x2).permute(1, 0).contiguous()) - l2_sf.append(down_sf.permute(1, 0).contiguous()) - l2_gs.append(down_gs) - return {'l1_fp4': l1_fp4, 'l1_sf': l1_sf, 'l1_gs': l1_gs, - 'l2_fp4': l2_fp4, 'l2_sf': l2_sf, 'l2_gs': l2_gs} - - -def main(): - torch.manual_seed(42) - expert_indices = [0, 1, 2] - num_experts = len(expert_indices) - hidden_size = 7168 - intermediate_size = 3072 - top_k = 2 - num_tokens = 4 - - print("=" * 70) - print(" Loading checkpoint") - print("=" * 70) - nvfp4_tensors = load_layer_tensors(MODEL_DIR, LAYER_IDX) - print(f" {len(nvfp4_tensors)} tensors loaded") - - weights = prepare_direct_weights(nvfp4_tensors, LAYER_IDX, expert_indices, intermediate_size) - - hidden_states = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device=DEVICE) * 2.0 - expert_ids = torch.tensor([[0, 1]] * num_tokens, dtype=torch.int32, device=DEVICE) - expert_weights = torch.tensor([[0.6, 0.4]] * num_tokens, dtype=torch.float32, device=DEVICE) - - # ── Path 1: moe_pipeline (reference, uses quantize_to_nvfp4) ── - print("\n Running moe_pipeline (dynamic gs)...") - pipeline_out = run_nvfp4_moe( - hidden_states.clone(), expert_ids.clone(), expert_weights.clone(), - weights, expert_indices, - ) - print(f" Pipeline: amax={pipeline_out.abs().max():.4f}, mean={pipeline_out.float().mean():.6f}") - - # ── Path 2: Nvfp4MoE with checkpoint input_scale (what vLLM uses) ── - print("\n Running Nvfp4MoE (checkpoint gs)...") - runner = Nvfp4MoE(num_experts, hidden_size, intermediate_size, device=DEVICE) - runner.prepare_weights_direct( - [w.clone() for w in weights['l1_fp4']], - [w.clone() for w in weights['l1_sf']], - list(weights['l1_gs']), - [w.clone() for w in weights['l2_fp4']], - [w.clone() for w in weights['l2_sf']], - list(weights['l2_gs']), - ) - # Set checkpoint input_scale (what vLLM does in finalize_weights) - igs = nvfp4_tensors[f"layers.{LAYER_IDX}.mlp.experts.0.gate_proj.input_scale"].item() - runner._l1_activation_global_scale = igs - runner._l2_activation_global_scale = igs - print(f" Checkpoint input_scale: {igs:.10f}") - - # Build topk_weights and topk_ids in the format the runner expects - # runner.run expects topk_ids as expert indices (0-based within our expert set) - topk_weights = expert_weights - topk_ids = expert_ids - - runner_out = runner.run(hidden_states.clone(), topk_weights, topk_ids) - print(f" Runner (ckpt gs): amax={runner_out.abs().max():.4f}, mean={runner_out.float().mean():.6f}") - - cos_ckpt = torch.nn.functional.cosine_similarity( - runner_out.flatten().unsqueeze(0).float(), - pipeline_out.flatten().unsqueeze(0).float(), - ).item() - print(f" Cosine vs pipeline: {cos_ckpt:.6f}") - - # ── Path 3: Nvfp4MoE with dynamic gs ── - print("\n Running Nvfp4MoE (dynamic gs)...") - # We can't use quantize_to_nvfp4 in the runner (cudagraph), but we can - # compute the gs from the input and set it before calling run - x_igs = (hidden_states.abs().max().item()) / (6.0 * 448.0) - runner2 = Nvfp4MoE(num_experts, hidden_size, intermediate_size, device=DEVICE) - runner2.prepare_weights_direct( - [w.clone() for w in weights['l1_fp4']], - [w.clone() for w in weights['l1_sf']], - list(weights['l1_gs']), - [w.clone() for w in weights['l2_fp4']], - [w.clone() for w in weights['l2_sf']], - list(weights['l2_gs']), - ) - runner2._l1_activation_global_scale = x_igs - runner2._l2_activation_global_scale = x_igs - print(f" Dynamic gs (from input amax): {x_igs:.10f}") - - runner2_out = runner2.run(hidden_states.clone(), topk_weights, topk_ids) - print(f" Runner (dynamic gs): amax={runner2_out.abs().max():.4f}, mean={runner2_out.float().mean():.6f}") - - cos_dyn = torch.nn.functional.cosine_similarity( - runner2_out.flatten().unsqueeze(0).float(), - pipeline_out.flatten().unsqueeze(0).float(), - ).item() - print(f" Cosine vs pipeline: {cos_dyn:.6f}") - - # ── Summary ── - print(f"\n{'=' * 70}") - print(f" RESULTS") - print(f"{'=' * 70}") - print(f" Runner with checkpoint gs vs pipeline: {cos_ckpt:.6f}") - print(f" Runner with dynamic gs vs pipeline: {cos_dyn:.6f}") - if cos_dyn > 0.95: - print(f" ✅ Dynamic gs fixes the problem — gs is the only bug") - elif cos_dyn < 0.5 and cos_ckpt < 0.5: - print(f" ❌ Both runner paths are broken — scale assembly is also wrong") - else: - print(f" ⚠️ Partial match — multiple issues") - - -if __name__ == "__main__": - main() diff --git a/tests/archive/test_scale_assembly.py b/tests/archive/test_scale_assembly.py deleted file mode 100644 index 7be030a7..00000000 --- a/tests/archive/test_scale_assembly.py +++ /dev/null @@ -1,116 +0,0 @@ -#!/usr/bin/env python3 -""" -Test B: Compare _assemble_scales_cudagraph_safe vs assemble_scales_2d_side. - -Both should produce identical output given the same x_sf and expert_offsets. -If they differ, the cudagraph-safe path has a bug. - -Runs on the B200 host: - source /root/nvfp4-megamoe-kernel/tests/.venv/bin/activate - python3 tests/test_scale_assembly.py -""" -import os, sys, torch - -REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) -sys.path.insert(0, REPO_ROOT) - -from dsv4.ops.quantize import ( - quantize_to_nvfp4, -) -from dsv4.ops.layouts import ( - assemble_scales_2d_side, -) -from dsv4.kernels.gemm.grouped import pad_and_swizzle_single, ceil_div -from vllm.nvfp4_cutedsl import Nvfp4MoE - - -def test_scale_assembly(): - """Compare the two scale assembly methods with realistic data.""" - DEVICE = "cuda" - num_experts = 3 - hidden_size = 7168 - intermediate_size = 3072 - - # Create a runner just to use its _assemble_scales_cudagraph_safe - runner = Nvfp4MoE(num_experts, hidden_size, intermediate_size, device=DEVICE) - # Trigger _ensure_stacked and buffer allocation with dummy weights - def rand_fp4(*shape): - return torch.randint(0, 256, shape, dtype=torch.uint8, device=DEVICE).view(torch.float4_e2m1fn_x2) - def rand_sf(*shape): - return torch.rand(shape, dtype=torch.float16, device=DEVICE).to(torch.float8_e4m3fn) - runner.prepare_weights_direct( - [rand_fp4(3584, intermediate_size * 2) for _ in range(num_experts)], - [rand_sf(3584 // 16, intermediate_size * 2) for _ in range(num_experts)], - [0.1] * num_experts, - [rand_fp4(1536, hidden_size) for _ in range(num_experts)], - [rand_sf(1536 // 16, hidden_size) for _ in range(num_experts)], - [0.1] * num_experts, - ) - runner._ensure_stacked() - - # Test with different token distributions - test_cases = [ - ("4 tokens, expert 0 gets 2, expert 1 gets 2, expert 2 gets 0", [2, 2, 0]), - ("8 tokens, expert 0 gets 4, expert 1 gets 3, expert 2 gets 1", [4, 3, 1]), - ("4 tokens, expert 0 gets 4, expert 1 gets 0, expert 2 gets 0", [4, 0, 0]), - ("3 tokens, expert 0 gets 1, expert 1 gets 1, expert 2 gets 1", [1, 1, 1]), - ] - - all_pass = True - for desc, tokens_per_expert in test_cases: - total_tokens = sum(tokens_per_expert) - - # Create input and quantize - x = torch.randn(total_tokens, hidden_size, dtype=torch.bfloat16, device=DEVICE) * 2.0 - x_fp4, x_sf, x_igs = quantize_to_nvfp4(x) - - # Path 1: assemble_scales_2d_side (per-expert split) - x_sf_parts = [] - offset = 0 - for tpe in tokens_per_expert: - x_sf_parts.append(x_sf[offset:offset + tpe]) - offset += tpe - scale_a_ref = assemble_scales_2d_side(x_sf_parts) - - # Path 2: _assemble_scales_cudagraph_safe (GPU-only) - expert_offsets = torch.zeros(num_experts + 1, dtype=torch.int32, device=DEVICE) - expert_offsets[1:] = torch.tensor(tokens_per_expert, dtype=torch.int32).cumsum(0) - scale_a_cudagraph = runner._assemble_scales_cudagraph_safe( - x_sf, expert_offsets, - runner._padded_x_sf_buf_l1, runner._per_expert_scale_bufs_l1 - ) - - # Compare - # Note: shapes may differ due to padding, but the data in the - # padded rows should match (up to the total number of rows used by the kernel) - if scale_a_ref.shape != scale_a_cudagraph.shape: - print(f" {desc}") - print(f" Shape mismatch: ref={scale_a_ref.shape}, cg={scale_a_cudagraph.shape}") - all_pass = False - continue - - match = torch.equal(scale_a_ref, scale_a_cudagraph) - if not match: - # Check how many bytes differ - diff = (scale_a_ref.view(torch.uint8) != scale_a_cudagraph.view(torch.uint8)).sum().item() - total = scale_a_ref.numel() - pct = diff / total * 100 - print(f" {desc}") - print(f" MISMATCH: {diff}/{total} bytes differ ({pct:.1f}%)") - print(f" ref range: [{scale_a_ref.view(torch.uint8).min()}, {scale_a_ref.view(torch.uint8).max()}]") - print(f" cg range: [{scale_a_cudagraph.view(torch.uint8).min()}, {scale_a_cudagraph.view(torch.uint8).max()}]") - all_pass = False - else: - print(f" {desc}: ✅ MATCH") - - print(f"\n{'=' * 70}") - if all_pass: - print(" ALL SCALE ASSEMBLY TESTS PASSED ✅") - else: - print(" SCALE ASSEMBLY TESTS FAILED ❌") - print(f"{'=' * 70}") - return all_pass - - -if __name__ == "__main__": - test_scale_assembly() diff --git a/tests/archive/test_scale_debug.py b/tests/archive/test_scale_debug.py deleted file mode 100644 index 794f3285..00000000 --- a/tests/archive/test_scale_debug.py +++ /dev/null @@ -1,77 +0,0 @@ -#!/usr/bin/env python3 -"""Deep-dive: compare scale assembly byte-by-byte for expert 0.""" -import os, sys, torch -REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) -sys.path.insert(0, REPO_ROOT) -from dsv4.ops.quantize import ( - quantize_to_nvfp4, -) -from dsv4.ops.layouts import ( - assemble_scales_2d_side, -) -from dsv4.kernels.gemm.grouped import pad_and_swizzle_single, ceil_div -from vllm.nvfp4_cutedsl import Nvfp4MoE - -DEVICE = "cuda" -num_experts = 3 -hidden_size = 7168 - -runner = Nvfp4MoE(num_experts, hidden_size, 3072, device=DEVICE) -def rand_fp4(*shape): - return torch.randint(0, 256, shape, dtype=torch.uint8, device=DEVICE).view(torch.float4_e2m1fn_x2) -def rand_sf(*shape): - return torch.rand(shape, dtype=torch.float16, device=DEVICE).to(torch.float8_e4m3fn) -runner.prepare_weights_direct( - [rand_fp4(3584, 3072*2) for _ in range(num_experts)], - [rand_sf(3584//16, 3072*2) for _ in range(num_experts)], - [0.1]*num_experts, - [rand_fp4(1536, hidden_size) for _ in range(num_experts)], - [rand_sf(1536//16, hidden_size) for _ in range(num_experts)], - [0.1]*num_experts, -) -runner._ensure_stacked() - -# 8 tokens, expert 0 gets 4, expert 1 gets 3, expert 2 gets 1 -tokens_per_expert = [4, 3, 1] -total = sum(tokens_per_expert) -x = torch.randn(total, hidden_size, dtype=torch.bfloat16, device=DEVICE) * 2.0 -x_fp4, x_sf, x_igs = quantize_to_nvfp4(x) - -# Reference: assemble_scales_2d_side -x_sf_parts = [x_sf[0:4], x_sf[4:7], x_sf[7:8]] -ref = assemble_scales_2d_side(x_sf_parts) - -# Cudagraph-safe -expert_offsets = torch.zeros(num_experts + 1, dtype=torch.int32, device=DEVICE) -expert_offsets[1:] = torch.tensor(tokens_per_expert, dtype=torch.int32).cumsum(0) -cg = runner._assemble_scales_cudagraph_safe(x_sf, expert_offsets) - -print(f"ref shape: {ref.shape}, cg shape: {cg.shape}") - -# Compare expert 0's block (first 128 rows) -ref_e0 = ref[:128].view(torch.uint8) -cg_e0 = cg[:128].view(torch.uint8) -diff_e0 = (ref_e0 != cg_e0).sum().item() -print(f"Expert 0: {diff_e0}/{ref_e0.numel()} bytes differ") - -# Where do they differ? -if diff_e0 > 0: - diff_idx = torch.where(ref_e0.flatten() != cg_e0.flatten())[0] - for i in diff_idx[:20]: - print(f" byte {i}: ref={ref_e0.flatten()[i].item()}, cg={cg_e0.flatten()[i].item()}") - -# Also test: does pad_and_swizzle_single on the SAME input give the same output? -buf = torch.zeros(128, 448, dtype=torch.float16, device=DEVICE).to(torch.float8_e4m3fn) -buf[:4, :x_sf.shape[1]] = x_sf[0:4] -swizzled_buf = pad_and_swizzle_single(buf) - -# Now compare swizzled_buf vs the reference path's expert 0 swizzled block -# Reference path swizzles x_sf[0:4] padded to 128 rows -buf2 = torch.zeros(4, 448, dtype=torch.float16, device=DEVICE).to(torch.float8_e4m3fn) -buf2[:4, :x_sf.shape[1]] = x_sf[0:4] -# Wait, assemble_scales_2d_side pads to 128 rows first -buf3 = torch.zeros(128, 448, dtype=torch.float16, device=DEVICE).to(torch.float8_e4m3fn) -buf3[:4, :x_sf.shape[1]] = x_sf[0:4] -swizzled_ref = pad_and_swizzle_single(buf3) - -print(f"\nSame-input swizzle comparison: {torch.equal(swizzled_buf.view(torch.uint8), swizzled_ref.view(torch.uint8))}") diff --git a/tests/archive/test_shared_expert.py b/tests/archive/test_shared_expert.py deleted file mode 100644 index bf6856ab..00000000 --- a/tests/archive/test_shared_expert.py +++ /dev/null @@ -1,163 +0,0 @@ -"""Standalone test: Shared expert using CuTeDSL dedicated runner. - -Tests the Nvfp4SharedExpert for the shared expert path. -Compares against BF16 dequantized reference. - -Usage: python3 test_shared_expert.py -""" - -import torch -import torch.nn.functional as F -import sys, os, json -from safetensors import safe_open - -MODEL_PATH = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4" -DEVICE = "cuda:0" -LAYER_IDX = 0 -HIDDEN_SIZE = 7168 # shared expert input dim (from checkpoint weight shapes) -INTERMEDIATE_SIZE = 3072 -SWIGLU_LIMIT = 10.0 -NUM_TOKENS = 4 - -E2M1_LUT = torch.tensor([0., 0.5, 1., 1.5, 2., 3., 4., 6., -0., -0.5, -1., -1.5, -2., -3., -4., -6.], - dtype=torch.float32) - -_cache = {} - -def load_tensor(key, wm, model_dir): - if key in _cache: - return _cache[key] - shard_path = os.path.join(model_dir, wm[key]) - with safe_open(shard_path, framework="pt") as f: - t = f.get_tensor(key) - _cache[key] = t - return t - - -def dequant_nvfp4(packed_uint8, scale_e4m3, global_scale): - """Dequantize NVFP4 weight to BF16 for reference.""" - device = packed_uint8.device - lut = E2M1_LUT.to(device) - lower = lut[(packed_uint8 & 0x0F).long()] - upper = lut[((packed_uint8 >> 4) & 0x0F).long()] - out_features = packed_uint8.shape[0] - in_features = packed_uint8.shape[1] * 2 - unpacked = torch.empty(out_features, in_features, dtype=torch.float32, device=device) - unpacked[:, 0::2] = lower - unpacked[:, 1::2] = upper - block_scale = scale_e4m3.float() - block_expanded = block_scale.repeat_interleave(16, dim=1)[:out_features, :in_features] - return (unpacked * block_expanded * global_scale).to(torch.bfloat16) - - -def main(): - torch.cuda.set_device(0) - torch.manual_seed(42) - - sys.path.insert(0, "/root/nvfp4-megamoe-kernel") - from dsv4.layers.shared_expert import Nvfp4SharedExpert - - with open(os.path.join(MODEL_PATH, "model.safetensors.index.json")) as f: - wm = json.load(f)["weight_map"] - P = lambda key: load_tensor(key, wm, MODEL_PATH).to(DEVICE) - - print("=== Shared Expert Test (CuTeDSL SharedExpertRunner) ===\n") - - # Load shared expert weights - prefix = f"model.layers.{LAYER_IDX}.mlp.shared_experts" - - gate_w = P(f"{prefix}.gate_proj.weight") - gate_sf = P(f"{prefix}.gate_proj.weight_scale") - gate_gs = P(f"{prefix}.gate_proj.weight_scale_2").item() - up_w = P(f"{prefix}.up_proj.weight") - up_sf = P(f"{prefix}.up_proj.weight_scale") - up_gs = P(f"{prefix}.up_proj.weight_scale_2").item() - down_w = P(f"{prefix}.down_proj.weight") - down_sf = P(f"{prefix}.down_proj.weight_scale") - down_gs = P(f"{prefix}.down_proj.weight_scale_2").item() - - print(f"gate_proj: shape={gate_w.shape} gs={gate_gs:.8f} sf_shape={gate_sf.shape}") - print(f"up_proj: shape={up_w.shape} gs={up_gs:.8f} sf_shape={up_sf.shape}") - print(f"down_proj: shape={down_w.shape} gs={down_gs:.8f} sf_shape={down_sf.shape}") - - # Stack gate + up into gate_up_proj (same format as MoE L1) - # gate/up weights are (intermediate, hidden) uint8 packed - gate_up_w = torch.cat([gate_w, up_w], dim=0) - gate_up_sf = torch.cat([gate_sf, up_sf], dim=0) - mgs = max(gate_gs, up_gs) - if gate_gs != up_gs: - sf32 = gate_up_sf.float() - sf32[:INTERMEDIATE_SIZE] *= (gate_gs / mgs) - sf32[INTERMEDIATE_SIZE:] *= (up_gs / mgs) - gate_up_sf = sf32.to(torch.float8_e4m3fn) - - # Convert to CuTeDSL format: - # Checkpoint weights are (out_features, in_features) uint8 packed - # We need float4_e2m1fn_x2 with (out_features, in_features // 2) after view - # Then permute to (in_features // 2, out_features) for K-major (K=in_features) - l1_fp4 = [gate_up_w.view(torch.float4_e2m1fn_x2).permute(1, 0).contiguous()] - l1_sf = [gate_up_sf.permute(1, 0).contiguous()] - l2_fp4 = [down_w.view(torch.float4_e2m1fn_x2).permute(1, 0).contiguous()] - l2_sf = [down_sf.permute(1, 0).contiguous()] - - # Create runner - runner = Nvfp4SharedExpert( - hidden_size=HIDDEN_SIZE, - intermediate_size=INTERMEDIATE_SIZE, - max_num_tokens=8192, - device=DEVICE, - swiglu_limit=SWIGLU_LIMIT, - ) - runner.l1_fp4 = l1_fp4 - runner.l1_sf = l1_sf - runner.l1_gs = [mgs] - runner.l2_fp4 = l2_fp4 - runner.l2_sf = l2_sf - runner.l2_gs = [down_gs] - runner.finalize_weights() - - # Warmup to compute activation global scales - dummy = torch.randn(NUM_TOKENS, HIDDEN_SIZE, dtype=torch.bfloat16, device=DEVICE) * 2.0 - runner._ensure_initialized() - runner.compute_activation_global_scales(dummy) - print(f"Warmup gs: L1={runner._l1_activation_global_scale:.6f} " - f"L2={runner._l2_activation_global_scale:.6f}") - - # Run CuTeDSL - print("\n--- CuTeDSL Forward ---") - hidden = torch.randn(NUM_TOKENS, HIDDEN_SIZE, dtype=torch.bfloat16, device=DEVICE) * 2.0 - - with torch.no_grad(): - output = runner.run(hidden) - print(f"CuTeDSL output: shape={output.shape} amax={output.amax():.4f} " - f"NaN={torch.isnan(output).any()}") - - # BF16 reference - print("\n--- BF16 Reference ---") - gate_bf16 = dequant_nvfp4(gate_w, gate_sf, gate_gs) - up_bf16 = dequant_nvfp4(up_w, up_sf, up_gs) - down_bf16 = dequant_nvfp4(down_w, down_sf, down_gs) - - with torch.no_grad(): - gate = hidden @ gate_bf16.T - up = hidden @ up_bf16.T - gate_silu = F.silu(gate).clamp(max=SWIGLU_LIMIT) - up = up.clamp(min=-SWIGLU_LIMIT, max=SWIGLU_LIMIT) - intermediate = gate_silu * up - ref_output = intermediate @ down_bf16.T - - print(f"BF16 ref: shape={ref_output.shape} amax={ref_output.amax():.4f}") - - # Compare - cos = F.cosine_similarity(ref_output.flatten().unsqueeze(0), - output.flatten().unsqueeze(0)).item() - mse = (ref_output - output).pow(2).mean().item() - print(f"\n=== RESULT: cosine={cos:.6f} MSE={mse:.6e} ===") - if cos >= 0.98: - print("✅ PASS") - else: - print("❌ FAIL") - - -if __name__ == "__main__": - main() diff --git a/tests/archive/test_silu_step1.py b/tests/archive/test_silu_step1.py deleted file mode 100644 index 8b11b35c..00000000 --- a/tests/archive/test_silu_step1.py +++ /dev/null @@ -1,97 +0,0 @@ -"""Test: Validate that cute.exp works on register tensors in the fused epilogue. - -Step 1 of the fused SwiGLU validation. We test with fused_swiglu=True but -with the full SiLU applied (not gate/up pairing yet). This confirms that: -1. cute.exp works on register tensors -2. The element-wise SiLU (x / (1+exp(-x))) produces correct values -3. The register tensor can be converted to BF16 and stored to C - -The test compares the fused kernel output (SiLU applied in registers) -against the PyTorch equivalent (SiLU applied to the BF16 L1 output). -""" -import torch -import sys -sys.path.insert(0, '/root/dsv4-nvfp4-workspace/kernel') - -from dsv4.ops.quantize import ( - quantize_weight_to_nvfp4, - quantize_activation_nvfp4, -) -from dsv4.ops.layouts import ( - make_b_k_major, - assemble_scales_2d_side, - assemble_scales_3d_side, -) -from dsv4.ops.gemm_runner import ( - run_nvfp4_grouped_gemm, - warmup_compilation, -) - - -def test_silu_in_registers(): - """Compare SiLU applied in registers vs SiLU applied in PyTorch.""" - device = "cuda" - num_experts = 4 - hidden = 512 - intermediate = 256 - num_tokens = 32 - - torch.manual_seed(42) - x = torch.randn(num_tokens, hidden, dtype=torch.bfloat16, device=device) - - # Create and quantize L1 weights (gate+up fused) - l1_w = torch.randn(num_experts, 2 * intermediate, hidden, dtype=torch.bfloat16, device=device) - l1_fp4_list, l1_sf_list, l1_gs_list = [], [], [] - for e in range(num_experts): - w_fp4, w_sf, w_gs = quantize_weight_to_nvfp4(l1_w[e].T) - l1_fp4_list.append(w_fp4) - l1_sf_list.append(w_sf) - l1_gs_list.append(w_gs) - - l1_mat_b = make_b_k_major(torch.stack(l1_fp4_list)) - l1_scale_b = assemble_scales_3d_side(l1_sf_list) - l1_gs = torch.tensor(l1_gs_list, dtype=torch.float32, device=device) - - gs_val = x.abs().max().item() / (6.0 * 448.0) - x_fp4, x_sf = quantize_activation_nvfp4(x, gs_val) - - tokens_per_expert = [num_tokens // num_experts] * num_experts - scale_a = assemble_scales_2d_side([x_sf[i*tpe:(i+1)*tpe] for i, tpe in enumerate(tokens_per_expert)]) - expert_offsets = torch.tensor( - [sum(tokens_per_expert[:e+1]) for e in range(num_experts)], - dtype=torch.int32, device=device, - ) - global_scale_a = torch.full((num_experts,), gs_val, dtype=torch.float32, device=device) - - # Warmup standard GEMM - warmup_compilation(num_experts, hidden // 2, (2 * intermediate) // 2, device) - - # Run standard L1 GEMM (no SiLU) - out_bf16 = run_nvfp4_grouped_gemm( - mat_a=x_fp4, mat_b=l1_mat_b, - scale_a=scale_a, scale_b=l1_scale_b, - expert_offsets=expert_offsets, - global_scale_a=global_scale_a, global_scale_b=l1_gs, - ) - - # Apply SiLU in PyTorch (reference) - silu_ref = torch.nn.functional.silu(out_bf16) - - print(f"L1 BF16 output shape: {out_bf16.shape}") - print(f"SiLU reference shape: {silu_ref.shape}") - print(f"L1 output amax: {out_bf16.abs().amax().item():.4f}") - print(f"SiLU reference amax: {silu_ref.abs().amax().item():.4f}") - print() - print("Step 1 validation: SiLU in PyTorch on BF16 GEMM output") - print("Next step: Run fused kernel with SiLU in registers and compare") - print() - print("NOTE: The fused kernel with SiLU on the full acc_vec should produce") - print("the same result as torch.nn.functional.silu on the BF16 output,") - print("within NVFP4 quantization tolerance (~5e-2).") - print() - print("This test validates the SiLU math. The gate/up pairing (Step 2)") - print("will change which values get SiLU applied (gate only, not up).") - - -if __name__ == "__main__": - test_silu_in_registers() diff --git a/tests/archive/test_softmax_only.py b/tests/archive/test_softmax_only.py deleted file mode 100644 index 3de3f670..00000000 --- a/tests/archive/test_softmax_only.py +++ /dev/null @@ -1,288 +0,0 @@ -""" -Test: QK + softmax packing only (no PV). -Output is the softmax-packed P (BF16) stored to GMEM via epilogue. -This tests whether the FP32→BF16 packing in TMEM works correctly. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - - -class SoftmaxOnlyKernel: - def __init__(self, mma_tiler_mn, use_2cta_instrs=False, use_tma_store=True): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = use_2cta_instrs; self.use_tma_store = use_tma_store - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.epilog_sync_bar_id = 1; self.tmem_alloc_sync_bar_id = 2; self.tmem_dealloc_sync_bar_id = 3 - self.num_c_stage = 2 - - def _setup(self, tiled_mma): - mma_inst_k = cute.size(tiled_mma.shape_mnk, mode=[2]) - self.mma_tiler = (*self.mma_tiler_mn, mma_inst_k * 4) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (tiled_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), - self.mma_tiler[1], self.mma_tiler[2]) - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, self.use_2cta_instrs, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(tiled_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(tiled_mma, self.mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) - tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(tCtAcc_fake, arch="sm_100") - self.tilePlikeFP32 = self.mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 # BF16 P location in TMEM - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) - ) * cute.size(tiled_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - self.q_dtype = a.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(a).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(b).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - tiled_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - self._setup(tiled_mma) - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - tma_a, tma_ta = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma.thr_id), - a, a_smem, self.mma_tiler, tiled_mma, self.cluster_layout_vmnk.shape) - tma_b, tma_tb = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, tiled_mma.thr_id), - b, b_smem, self.mma_tiler, tiled_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(tiled_mma, tma_a, tma_ta, tma_b, tma_tb, tma_c, tma_tc, - self.cluster_layout_vmnk, self.a_smem_s, self.b_smem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, tiled_mma, tma_a, mA, tma_b, mB, tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(tiled_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_a); cpasync.prefetch_descriptor(tma_b); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=self.tmem_alloc_sync_bar_id, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sA = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sB = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gA = cute.local_tile(mA, cute.slice_(self.mma_tiler, (None,0,None)), (None,None,None)) - gB = cute.local_tile(mB, cute.slice_(self.mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gA, mode=[3]) - - thr_mma = tiled_mma.get_slice(0) - tCgA = thr_mma.partition_A(gA); tCgB = thr_mma.partition_B(gB); tCgC = thr_mma.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsA, tAgA = cpasync.tma_partition(tma_a, 0, a_lay, cute.group_modes(sA,0,3), cute.group_modes(tCgA,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsB, tBgB = cpasync.tma_partition(tma_b, 0, b_lay, cute.group_modes(sB,0,3), cute.group_modes(tCgB,0,3)) - tAgA = tAgA[(None,0,None,0)]; tBgB = tBgB[(None,0,None,0)] - - tCrA = tiled_mma.make_fragment_A(sA); tCrB = tiled_mma.make_fragment_B(sB) - acc_shape = thr_mma.partition_shape_C(self.mma_tiler[:2]) - tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_acc_stage)) - - # S in TMEM - tStS = thr_mma.make_fragment_C(acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # TMA WARP - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_a, tAgA[(None,h.count)], tAsA[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_b, tBgB[(None,h.count)], tBsB[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_softmax_store_debug.py b/tests/archive/test_softmax_store_debug.py deleted file mode 100644 index 2bc3716b..00000000 --- a/tests/archive/test_softmax_store_debug.py +++ /dev/null @@ -1,247 +0,0 @@ -""" -Debug: Verify softmax store actually writes non-zero P to TMEM. -Test: QK only + identity softmax + read back S region (should be partially overwritten by P). -If S region is modified after softmax, the store is working. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - -HEAD_DIM = 64 - -class SoftmaxStoreDebug: - def __init__(self): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0,1,2,3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.num_c_stage = 2 - self.kv_stage = 2; self.q_stage = 1; self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (128, 128, qk_ik * 4) - pv_ik = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (128, HEAD_DIM, pv_ik * (128 // pv_ik)) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0]//cute.size(qk_mma.thr_id.shape), HEAD_DIM, self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - self.q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, self.q_dtype, self.q_stage) - self.k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, self.q_dtype, self.kv_stage) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, self.kv_stage) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - qk_thr = qk_mma.get_slice(0); qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - pv_thr = pv_mma.get_slice(0); pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - self.tmem_s0_offset = 0; self.tmem_p0_offset = 32 - self.tmem_o0_offset = find_tmem_tensor_col_offset(tOtO) - tCS = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage)) - tCO = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCS, tCO], arch="sm_100") - cta = cute.size(qk_mma.thr_id.shape) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)) - self.q_tx_bytes = cute.size_in_bytes(self.q_dtype, q_s) * cta - self.kv_tx_bytes = cute.size_in_bytes(self.q_dtype, k_s) * cta - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type - a_major = LayoutEnum.from_tensor(q).mma_major_mode() - b_major = LayoutEnum.from_tensor(k).mma_major_mode() - v_major = LayoutEnum.from_tensor(v).mma_major_mode() - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, a_major, b_major, self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, v_major, self.qk_acc_dtype, self.cta_group, (128,HEAD_DIM), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)); v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id),q,q_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id),k,k_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape) - epi_s = cute.select(self.c_smem_s,mode=[0,1]) - tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile) - self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.v_smem_s,self.p_tmem_s,self.c_smem_s,self.epi_tile).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, q_smem_s, k_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx,_,_ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k); cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - q_bar: cute.struct.MemRange[cutlass.Int64, self.q_stage*2] - kv_bar: cute.struct.MemRange[cutlass.Int64, self.kv_stage*2] - s_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage*2] - tmem_dealloc: cutlass.Int64; holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - qp,qc = pipeline.PipelineTmaUmma.create(barrier_storage=st.q_bar.data_ptr(),num_stages=self.q_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.q_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - kvp,kvc = pipeline.PipelineTmaUmma.create(barrier_storage=st.kv_bar.data_ptr(),num_stages=self.kv_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.kv_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - s_prod,s_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.s_bar.data_ptr(),num_stages=1,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.epilogue_warp_id))).make_participants() - softmax_done_bar = pipeline.NamedBarrier(barrier_id=3, num_threads=32 + 32*len(self.epilogue_warp_id)) - acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(),num_stages=self.num_acc_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,len(self.epilogue_warp_id)),cta_layout_vmnk=cl_vmnk,defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2,num_threads=32*len((self.mma_warp_id,*self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr,barrier_for_retrieve=tmem_bar,allocator_warp_id=self.epilogue_warp_id[0],is_two_cta=cute.size(qk_mma.thr_id.shape)==2,two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk,is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype,layout=q_smem_s.outer,byte_alignment=128,swizzle=q_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype,layout=k_smem_s.outer,byte_alignment=128,swizzle=k_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype,layout=v_smem_s.outer,byte_alignment=128,swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype,layout=c_smem_s.outer,byte_alignment=128,swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None)) - gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None)) - gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None)) - gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None)) - n_kv_tiles = cute.size(gK, mode=[3]) - - qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - tCgV = pv_thr.partition_B(gV); tCgC = pv_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,0,None,0)).shape) - tAsQ,tAgQ = cpasync.tma_partition(tma_q,0,a_lay,cute.group_modes(sQ,0,3),cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape) - tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3)) - tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)]; tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - # PV read view (MMA only) - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None,None,None,0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # TMA LOAD - if warp_idx == self.tma_warp_id: - qp.reset(); qh = qp.acquire_and_advance() - cute.copy(tma_q,tAgQ[(None,qh.count)],tAsQ[(None,qh.index)],tma_bar_ptr=qh.barrier) - qp.tail() - kvp.reset(); pk = kvp.try_acquire() - for kt in cutlass.range(n_kv_tiles,unroll=1): - kh = kvp.acquire_and_advance(pk) - cute.copy(tma_k,tBgK[(None,kh.count)],tBsK[(None,kh.index)],tma_bar_ptr=kh.barrier) - pk = cutlass.Boolean(1) - vh = kvp.acquire_and_advance(pk) - cute.copy(tma_v,tVgV[(None,vh.count)],tVsV[(None,vh.index)],tma_bar_ptr=vh.barrier) - pk = cutlass.Boolean(1) - kvp.tail() - - # MMA — same as v3 - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - qc.reset(); qh = qc.wait_and_advance(); qh.release() - kvc.reset(); pk = kvc.try_wait() - acc_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_stage) - acc_pipe.producer_acquire(acc_st) - for kt in range(n_kv_tiles): - kh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - sh = s_prod.acquire_and_advance() - qk_mma.set(tcgen05.Field.ACCUMULATE, False) - for kb in cutlass.range(cute.size(tCrQ,mode=[2]), unroll_full=True): - cute.gemm(qk_mma, tStS0, tCrQ[(None,None,kb,0)], tCrK[(None,None,kb,kh.index)], tStS0) - qk_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - sh.commit(); kh.release() - softmax_done_bar.arrive_and_wait() - vh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - pv_mma.set(tcgen05.Field.ACCUMULATE, kt != 0) - for kb in cutlass.range(cute.size(tOrP0,mode=[2]), unroll_full=True): - cute.gemm(pv_mma, tOtO0, tOrP0[(None,None,kb)], tCrV[(None,None,kb,vh.index)], tOtO0) - pv_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_async_view_tmem_store() - vh.release() - acc_pipe.producer_commit(acc_st); acc_st.advance() - acc_pipe.producer_tail(acc_st) - - # EPILOGUE — debug: dump S+P region to GMEM instead of doing PV - if warp_idx < self.mma_warp_id: - tmem.allocate(self.num_tmem_alloc_cols) - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) - sfw_idx = tidx % (32 * len(self.epilogue_warp_id)) - - # S load + P store (same as v3) - tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS0) - thr_load = tiled_tmem_load.get_slice(sfw_idx) - tTMEM_LOADtS = thr_load.partition_S(tStS0) - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tTMEM_LOADcS = thr_load.partition_D(tScS) - - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - tStP_layout = cute.composition(tStS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tStP0 = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStP_layout) - tmem_store_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStP0) - thr_store = tiled_tmem_store.get_slice(sfw_idx) - tTMEM_STOREtP = thr_store.partition_D(tStP0) - tScP_layout = cute.composition(tScS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tScP = cute.make_tensor(tScS.iterator, tScP_layout) - tTMEM_STOREcP = thr_store.partition_S(tScP) - - for kt in range(n_kv_tiles): - si_handle = s_cons.wait_and_advance() - tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS) - rP_words = cute.make_rmem_tensor(tTMEM_STOREcP.shape, self.qk_acc_dtype) - rP_bf16 = cute.make_tensor(cute.recast_ptr(rP_words.iterator, dtype=self.q_dtype), tTMEM_LOADrS.layout) - frg_cnt = 4; frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt - tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile)) - rP_bf16_frg = cute.logical_divide(rP_bf16, cute.make_layout(frg_tile)) - for j in range(frg_cnt): - s_vec = tTMEM_LOADrS_frg[None, j].load() - rP_bf16_frg[None, j].store(s_vec.to(self.q_dtype)) - cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP) - cute.arch.fence_view_async_tmem_store() - si_handle.release() - softmax_done_bar.arrive() - - # Now read back the P region directly - # Read from tStP0 (the same region we wrote to) - tmem_read_P_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_read_P = tcgen05.make_tmem_copy(tmem_read_P_atom, tStP0) - thr_read_P = tiled_tmem_read_P.get_slice(sfw_idx) - tTMEM_READtP = thr_read_P.partition_S(tStP0) - rP_read = cute.make_rmem_tensor(thr_read_P.partition_D(tScP).shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_read_P, tTMEM_READtP, rP_read) - # Print the first few values - if sfw_idx == 0: - print(f"DEBUG P readback: size={cute.size(rP_read)} val0={float(rP_read.load()[0])}") - - # Epilogue as before - tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout) - acc_cons_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage) - c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) - c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp) - acc_cons_st = utils.gemm.sm100.epilogue_tma_store(self, tidx, warp_idx, tma_c, tCtO_base, sC, tCgC, epi_tile, 0, const_expr(lambda x: x), (0,0,0), acc_cons_st, acc_pipe, c_pipe) - c_pipe.producer_tail() - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) diff --git a/tests/archive/test_sparse_attn_b200.py b/tests/archive/test_sparse_attn_b200.py deleted file mode 100644 index da7adc42..00000000 --- a/tests/archive/test_sparse_attn_b200.py +++ /dev/null @@ -1,364 +0,0 @@ -#!/usr/bin/env python3 -""" -DeepSeek-V4 CSA/HCA Sparse Attention Kernel - -NOT MLA. CSA = Compressed Sparse Attention. HCA = Heavily Compressed Attention. - -The sparse attention works as follows: -1. KV latent is stored in a compressed cache (cr=4 for CSA, cr=128 for HCA) -2. The indexer finds the top-k most relevant positions in the compressed cache -3. Sparse attention: Q attends only to KV at those top-k positions -4. SWA attention: Q attends to the local sliding window -5. Merge: combine sparse + SWA outputs using attention sink weights - -This kernel implements step 3 (sparse attention with paged FP8 KV cache). - -Usage (on B200): - cd /root/nvfp4-megamoe-kernel - PYTHONPATH=/root/nvfp4-megamoe-kernel tests/venv/bin/python tests/test_sparse_attn_b200.py -""" - -import sys, os, json, torch, torch.nn.functional as F, math -from safetensors import safe_open - -REPO = "/root/nvfp4-megamoe-kernel" -sys.path.insert(0, REPO) -MODEL = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4" -DEV = "cuda:0" - -H = 7168; NH = 128; HD = 512; NOPE = 448; ROPE = 64 -QL = 1536; OL = 1024; OG = 16; HPG = NH // OG -EPS = 1e-6; WINDOW = 128; SCALE = HD ** -0.5 - -E2M1 = torch.tensor([0,.5,1.,1.5,2.,3.,4.,6.,-0,-.5,-1.,-1.5,-2.,-3.,-4.,-6.], dtype=torch.float32) - -_cache = {} -def P(k, wm, md): - if k in _cache: return _cache[k] - with safe_open(os.path.join(md, wm[k]), framework="pt") as f: - t = f.get_tensor(k) - _cache[k] = t - return t - -def dequant(w, sf, gs): - d = w.device; lut = E2M1.to(d) - lo = lut[(w & 0xF).long()]; hi = lut[((w >> 4) & 0xF).long()] - O, I2 = w.shape; I = I2*2 - u = torch.empty(O, I, dtype=torch.float32, device=d) - u[:,0::2] = lo; u[:,1::2] = hi - bs = sf.float().repeat_interleave(16, dim=1)[:O,:I] - return (u * bs * gs).to(torch.bfloat16) - -def rms(x, w, eps=1e-6): - v = x.float().pow(2).mean(-1, keepdim=True) - return (w.float() * (x * torch.rsqrt(v+eps)).float()).to(x.dtype) - -def make_runner(w, sf, gs_t, inf, outf, fused=False, lw=None): - from dsv4.layers.linear import Nvfp4Linear - fp4 = w.view(torch.float4_e2m1fn_x2).permute(1,0).contiguous() - s = sf.to(torch.float8_e4m3fn) if sf.dtype != torch.float8_e4m3fn else sf - s = s.permute(1,0).contiguous() - if fused and gs_t.numel() == 2: - g1,g2 = gs_t[0].item(), gs_t[1].item(); gs = max(g1,g2) - if g1 != g2: - s32 = s.float(); sp = lw[0] if lw else outf//2 - s32[:sp] *= g1/gs; s32[sp:] *= g2/gs; s = s32.to(torch.float8_e4m3fn) - else: - gs = gs_t.max().item() if gs_t.numel() > 1 else gs_t.item() - r = Nvfp4Linear(in_features=inf, out_features=outf, max_num_tokens=8192, device=str(w.device)) - r.fp4 = [fp4]; r.sf = [s]; r.gs = [gs] - r.finalize_weights(); r._ensure_initialized() - return r - -def build_cos_sin(max_pos=4096, rope_dim=ROPE): - half = rope_dim // 2 - inv_freq = 1.0 / (10000.0 ** (torch.arange(0, half, dtype=torch.float32) / half)) - freqs = torch.outer(torch.arange(max_pos, dtype=torch.float32), inv_freq) - return torch.cat([freqs.cos(), freqs.sin()], dim=-1) - -def apply_gptj_rope(x, positions, cos_sin, nope, rope): - if rope == 0 or x.numel() == 0: return x - half = rope // 2 - cos = cos_sin[positions, :half].to(x.dtype) - sin = cos_sin[positions, half:2*half].to(x.dtype) - if x.dim() == 3: cos = cos.unsqueeze(1); sin = sin.unsqueeze(1) - x_rope = x[..., nope:].clone() - even = x_rope[..., 0::2]; odd = x_rope[..., 1::2] - out = x.clone() - out[..., nope:][..., 0::2] = even * cos - odd * sin - out[..., nope:][..., 1::2] = even * sin + odd * cos - return out - - -# ── KV Cache Kernels ──────────────────────────────────────────────── - -def kv_quantize_fp8(kv_bf16): - amax = kv_bf16.float().abs().amax(dim=-1, keepdim=True).clamp(min=1e-12) - scale = 448.0 / amax - kv_fp8 = (kv_bf16.float() * scale).to(torch.float8_e4m3fn) - inv_scale = (amax / 448.0).to(torch.bfloat16) - return kv_fp8, inv_scale - -def kv_dequantize_fp8(kv_fp8, inv_scale): - return (kv_fp8.to(torch.bfloat16) * inv_scale).to(torch.bfloat16) - - -def sparse_attention(q, kv_cache_bf16, topk_indices, topk_lens, scale, - cos_sin_cache, positions, nope_dim=NOPE, rope_dim=ROPE, - attn_sink=None): - """CSA/HCA sparse attention. - - Args: - q: (T, NH, HD) with RoPE applied - kv_cache_bf16: (cache_len, HD) BF16 KV latent (already dequantized from fp8) - topk_indices: (T, num_topk) global position indices in the KV cache - topk_lens: (T,) valid length per token (how many topk positions are valid) - scale: 1/sqrt(HD) - cos_sin_cache: (max_pos, 2*half) for RoPE on gathered KV - positions: (T,) query position IDs - nope_dim: 448 - rope_dim: 64 - attn_sink: (NH,) sink bias weights - - Returns: (T, NH, HD) attention output - """ - T, NH, HD = q.shape - device = q.device - num_topk = topk_indices.shape[-1] - - # Clamp indices to valid range - safe_indices = topk_indices.clamp(min=0, max=kv_cache_bf16.shape[0] - 1) - - # Gather KV from cache: (T, num_topk, HD) - # For each query token, gather its top-k KV vectors - idx_expanded = safe_indices.unsqueeze(-1).expand(-1, -1, HD) - # kv_cache_bf16 is (cache_len, HD) → expand to (T, cache_len, HD) for gather - kv_expanded = kv_cache_bf16.unsqueeze(0).expand(T, -1, -1) - k_gathered = torch.gather(kv_expanded, 1, idx_expanded) # (T, num_topk, HD) - - # Apply RoPE to gathered KV at their original positions - if rope_dim > 0 and cos_sin_cache is not None: - kv_positions = safe_indices # (T, num_topk) - half = rope_dim // 2 - cos_kv = cos_sin_cache[kv_positions, :half].to(k_gathered.dtype) # (T, num_topk, half) - sin_kv = cos_sin_cache[kv_positions, half:2*half].to(k_gathered.dtype) - - k_rope = k_gathered[:, :, nope_dim:].clone() - k_even = k_rope[:, :, 0::2] - k_odd = k_rope[:, :, 1::2] - k_gathered[:, :, nope_dim:][:, :, 0::2] = k_even * cos_kv - k_odd * sin_kv - k_gathered[:, :, nope_dim:][:, :, 1::2] = k_even * sin_kv + k_odd * cos_kv - - # V = K in this attention (KV latent, shared K/V) - v_gathered = k_gathered.clone() - - # Expand for multi-head: (T, num_topk, HD) → (T, NH, num_topk, HD) - k_heads = k_gathered.unsqueeze(1).expand(-1, NH, -1, -1) - v_heads = v_gathered.unsqueeze(1).expand(-1, NH, -1, -1) - - # Q: (T, NH, HD) → (T, NH, 1, HD) - q_4d = q.unsqueeze(2) - - # Attention scores: (T, NH, 1, num_topk) - attn_weights = torch.matmul(q_4d, k_heads.transpose(-1, -2)) * scale - - # Apply attention sink bias to first position - if attn_sink is not None: - # attn_sink: (NH,) → (1, NH, 1, 1) - sink_bias = attn_sink.view(1, NH, 1, 1) - attn_weights[:, :, :, 0] += sink_bias.squeeze(-1) - - # Mask invalid positions - valid_mask = torch.arange(num_topk, device=device).unsqueeze(0) < topk_lens.unsqueeze(1) - attn_weights = attn_weights.masked_fill(~valid_mask.unsqueeze(1).unsqueeze(2), float('-inf')) - - # Softmax - attn_weights = F.softmax(attn_weights.float(), dim=-1).to(q.dtype) - - # Weighted sum: (T, NH, 1, HD) - output = torch.matmul(attn_weights, v_heads) - return output.squeeze(2) # (T, NH, HD) - - -def swa_attention(q, kv_cache_bf16, positions, scale, window_size=WINDOW): - """Sliding window attention: attend to last window_size tokens. - - For testing with small T, this is just causal attention. - """ - T, NH, HD = q.shape - device = q.device - - # Full causal attention (for T <= window_size) - q_2d = q.reshape(T * NH, HD) - kv_exp = kv_cache_bf16.unsqueeze(1).expand(-1, NH, -1).contiguous() - k_2d = kv_exp.permute(1, 0, 2).unsqueeze(1).expand(NH, T, T, -1).contiguous().reshape(T * NH, T, HD) - v_2d = k_2d.clone() - scores = torch.matmul(q_2d.unsqueeze(1), k_2d.transpose(-1, -2)) * scale - qpos = torch.arange(T, device=device).unsqueeze(1).repeat(1, NH).reshape(T * NH) - kpos = torch.arange(T, device=device).unsqueeze(0) - causal = kpos <= qpos.unsqueeze(1) - scores = scores.squeeze(1).masked_fill(~causal, float('-inf')) - weights = F.softmax(scores.float(), dim=-1).to(q.dtype) - out = torch.matmul(weights.unsqueeze(1), v_2d).squeeze(1) - return out.reshape(T, NH, HD) - - -def csa_hca_merged_attention(q, kv_cache_bf16, topk_indices, topk_lens, - positions, scale, cos_sin_cache, - compress_ratio, attn_sink=None): - """Full CSA/HCA + SWA merged attention. - - For compress_ratio <= 1: SWA only - For compress_ratio > 1: sparse + SWA, merged with sink weights - """ - if compress_ratio <= 1: - return swa_attention(q, kv_cache_bf16, positions, scale) - - # Sparse attention on compressed cache - sparse_out = sparse_attention( - q, kv_cache_bf16, topk_indices, topk_lens, scale, - cos_sin_cache, positions, attn_sink=attn_sink, - ) - - # SWA attention - swa_out = swa_attention(q, kv_cache_bf16, positions, scale) - - # Merge: sigmoid(sink) weights sparse vs SWA - if attn_sink is not None: - sink_weight = torch.sigmoid(attn_sink).view(1, NH, 1) - return sparse_out * (1 - sink_weight) + swa_out * sink_weight - else: - return sparse_out + swa_out - - -def main(): - torch.cuda.set_device(0) - torch.manual_seed(42) - - print("=" * 70) - print(" DeepSeek-V4 CSA/HCA Sparse Attention Kernel Test") - print(" Compressed Sparse Attention (NOT MLA)") - print("=" * 70) - - # Load model weights - with open(os.path.join(MODEL, "model.safetensors.index.json")) as f: - wm = json.load(f)["weight_map"] - G = lambda k: P(k, wm, MODEL).to(DEV) - - p = "model.layers.0"; a = f"{p}.self_attn" - emb = G("model.embed_tokens.weight") - anorm = G(f"{p}.input_layernorm.weight") - qn = G(f"{a}.q_a_norm.weight"); kvn = G(f"{a}.kv_norm.weight") - sinks = G(f"{a}.sinks") - - qa_w = G(f"{a}.q_a_proj.weight"); qa_sf = G(f"{a}.q_a_proj.weight_scale"); qa_gs = G(f"{a}.q_a_proj.weight_scale_2") - qb_w = G(f"{a}.q_b_proj.weight"); qb_sf = G(f"{a}.q_b_proj.weight_scale"); qb_gs = G(f"{a}.q_b_proj.weight_scale_2") - kv_w = G(f"{a}.kv_proj.weight"); kv_sf = G(f"{a}.kv_proj.weight_scale"); kv_gs = G(f"{a}.kv_proj.weight_scale_2") - - r_qa = make_runner(qa_w, qa_sf, qa_gs, H, qa_w.shape[0]) - r_qb = make_runner(qb_w, qb_sf, qb_gs, QL, qb_w.shape[0]) - r_kv = make_runner(kv_w, kv_sf, kv_gs, H, kv_w.shape[0]) - - cos_sin = build_cos_sin(max_pos=8192).to(DEV) - - NT = 32 # More tokens for sparse attention - token_ids = torch.randint(0, 129280, (NT,), dtype=torch.long, device=DEV) - positions = torch.arange(NT, dtype=torch.int64, device=DEV) - - with torch.no_grad(): - hidden = emb[token_ids] - normed = rms(hidden, anorm, EPS) - - # Projections - qa_cute = r_qa.run(normed) - kv_cute = r_kv.run(normed) - qa_n = rms(qa_cute, qn, EPS) - kv_n = rms(kv_cute, kvn, EPS) - q_cute = r_qb.run(qa_n).view(NT, NH, HD) - q_rope = apply_gptj_rope(q_cute, positions, cos_sin, NOPE, ROPE) - - # FP8 KV cache - kv_fp8, inv_scale = kv_quantize_fp8(kv_n) - kv_from_cache = kv_dequantize_fp8(kv_fp8, inv_scale) - kv_from_cache_rope = apply_gptj_rope(kv_from_cache.unsqueeze(1), positions, cos_sin, NOPE, ROPE).squeeze(1) - - # ── Test 1: SWA attention (no compression) ─────────────────── - print("\n--- Test 1: SWA attention (cr=1, layer 60) ---") - swa_out = swa_attention(q_rope, kv_from_cache_rope, positions, SCALE) - print(f" SWA attention output: amax={swa_out.amax():.4f} NaN={torch.isnan(swa_out).any()}") - - # Compare with full causal attention - full_out = swa_attention(q_rope, kv_from_cache_rope, positions, SCALE) # same for T<=WINDOW - c = F.cosine_similarity(swa_out.flatten().unsqueeze(0).float(), full_out.flatten().unsqueeze(0).float()).item() - print(f" SWA vs full attention cosine: {c:.6f} {'✅' if c>=0.99 else '❌'}") - - # ── Test 2: CSA sparse attention (cr=4) ────────────────────── - print("\n--- Test 2: CSA sparse attention (cr=4) ---") - # Simulate indexer: select top-8 positions (simplified — pick evenly spaced) - num_topk = 8 - # For a real indexer, this would be the output of the scoring + topk - # Here, simulate: every 4th position + some random - topk_indices = torch.zeros(NT, num_topk, dtype=torch.long, device=DEV) - topk_lens = torch.full((NT,), num_topk, dtype=torch.long, device=DEV) - for t in range(NT): - # Pick 8 evenly spaced positions from 0..t - if t + 1 <= num_topk: - topk_indices[t, :t+1] = torch.arange(t+1, device=DEV) - topk_lens[t] = t + 1 - else: - step = (t + 1) / num_topk - for k in range(num_topk): - topk_indices[t, k] = int(k * step) - - csa_out = sparse_attention( - q_rope, kv_from_cache_rope, topk_indices, topk_lens, SCALE, - cos_sin, positions, attn_sink=sinks[:NH], - ) - print(f" CSA sparse attention output: amax={csa_out.amax():.4f} NaN={torch.isnan(csa_out).any()}") - - # ── Test 3: HCA sparse attention (cr=128) ──────────────────── - print("\n--- Test 3: HCA sparse attention (cr=128) ---") - num_topk_128 = 4 # Fewer positions in HCA cache - topk_indices_128 = torch.zeros(NT, num_topk_128, dtype=torch.long, device=DEV) - topk_lens_128 = torch.full((NT,), num_topk_128, dtype=torch.long, device=DEV) - for t in range(NT): - # Pick 4 evenly spaced positions - if t + 1 <= num_topk_128: - topk_indices_128[t, :t+1] = torch.arange(t+1, device=DEV) - topk_lens_128[t] = t + 1 - else: - step = (t + 1) / num_topk_128 - for k in range(num_topk_128): - topk_indices_128[t, k] = int(k * step) - - hca_out = sparse_attention( - q_rope, kv_from_cache_rope, topk_indices_128, topk_lens_128, SCALE, - cos_sin, positions, attn_sink=sinks[:NH], - ) - print(f" HCA sparse attention output: amax={hca_out.amax():.4f} NaN={torch.isnan(hca_out).any()}") - - # ── Test 4: Merged CSA + SWA ──────────────────────────────── - print("\n--- Test 4: Merged CSA + SWA attention (cr=4) ---") - merged_out = csa_hca_merged_attention( - q_rope, kv_from_cache_rope, topk_indices, topk_lens, - positions, SCALE, cos_sin, compress_ratio=4, attn_sink=sinks[:NH], - ) - print(f" Merged attention output: amax={merged_out.amax():.4f} NaN={torch.isnan(merged_out).any()}") - - # ── Test 5: Full pipeline with real sink weights ───────────── - print("\n--- Test 5: Sink weights analysis ---") - print(f" Sink weights: min={sinks.min():.4f} max={sinks.max():.4f} mean={sinks.mean():.4f}") - print(f" Sigmoid(sink) range: {torch.sigmoid(sinks).min():.4f} to {torch.sigmoid(sinks).max():.4f}") - print(f" → Near 0 = mostly sparse, Near 1 = mostly SWA") - - print(f"\n{'='*70}") - print(f" DONE — All attention kernels tested") - print(f" SWA: ✅") - print(f" CSA sparse: {'✅' if not torch.isnan(csa_out).any() else '❌'}") - print(f" HCA sparse: {'✅' if not torch.isnan(hca_out).any() else '❌'}") - print(f" Merged CSA+SWA: {'✅' if not torch.isnan(merged_out).any() else '❌'}") - print(f"{'='*70}") - - -if __name__ == "__main__": - main() diff --git a/tests/archive/test_sparse_decode.py b/tests/archive/test_sparse_decode.py deleted file mode 100644 index 31964e4e..00000000 --- a/tests/archive/test_sparse_decode.py +++ /dev/null @@ -1,71 +0,0 @@ -import sys, torch, torch.nn.functional as F -sys.path.insert(0, "/root/dsv4-nvfp4-workspace/kernel") -from dsv4.ops.decode_sparse import native_sparse_decode_attention - -torch.manual_seed(42) -torch.cuda.set_device(0) -NH, HD, BS, WIN, TOPK = 128, 512, 256, 128, 16 - -for nt, swa_l, topk_l in [(2,32,8), (2,64,16), (4,32,16), (4,64,8)]: - q = torch.randn(nt, NH, HD, dtype=torch.bfloat16, device="cuda:0") * 0.1 - nb = 4 - # SWA cache - kv_bf = torch.randn(nb*BS, HD, dtype=torch.bfloat16, device="cuda:0") * 0.5 - am = kv_bf.float().abs().amax(-1, keepdim=True).clamp(min=1e-12) - f8m = torch.tensor(448.0, dtype=torch.float32, device="cuda:0") - swa_cache = (kv_bf.float() * f8m / am).to(torch.float8_e4m3fn)[:nb*BS].reshape(nb,BS,HD).view(torch.uint8) - inv_sc = (am / f8m).to(torch.bfloat16) - # Compressed cache - comp_bf = torch.randn(nb*BS, HD, dtype=torch.bfloat16, device="cuda:0") * 0.3 - am2 = comp_bf.float().abs().amax(-1, keepdim=True).clamp(min=1e-12) - comp_cache = (comp_bf.float() * f8m / am2).to(torch.float8_e4m3fn)[:nb*BS].reshape(nb,BS,HD).view(torch.uint8) - inv_sc2 = (am2 / f8m).to(torch.bfloat16) - - si = torch.zeros(nt, WIN, dtype=torch.int64, device="cuda:0") - sl = torch.zeros(nt, dtype=torch.int64, device="cuda:0") - ti = torch.zeros(nt, TOPK, dtype=torch.int64, device="cuda:0") - tl = torch.zeros(nt, dtype=torch.int64, device="cuda:0") - for t in range(nt): - sl[t] = swa_l - for i in range(swa_l): si[t,i] = i - for i in range(swa_l, WIN): si[t,i] = -1 - tl[t] = topk_l - for i in range(topk_l): ti[t,i] = 1000+i - for i in range(topk_l, TOPK): ti[t,i] = -1 - - sink = torch.full((NH,), float("-inf"), dtype=torch.float32, device="cuda:0") - ascale = HD ** -0.5 - - # Reference: combined SDPA - safe_swa = si.clamp(min=0) - swa_raw = swa_cache[safe_swa//BS, safe_swa%BS].view(torch.float8_e4m3fn) - swa_kv = (swa_raw.to(torch.bfloat16)*inv_sc[safe_swa]).to(torch.bfloat16) - comp_bs = comp_cache.shape[1] - safe_topk = ti.clamp(min=0) - comp_raw = comp_cache[safe_topk//comp_bs, safe_topk%comp_bs].view(torch.float8_e4m3fn) - comp_kv = (comp_raw.to(torch.bfloat16)*inv_sc2[safe_topk]).to(torch.bfloat16) - kv_comb = torch.cat([swa_kv, comp_kv], dim=1) - total = WIN + TOPK - cl = sl + tl - - # Build mask - pos = torch.arange(total, device="cuda:0").unsqueeze(0) - lm = pos >= cl.unsqueeze(1) - inv_s = si < 0 - inv_t = ti < 0 - inv = torch.cat([inv_s, inv_t], dim=1) - mask = lm | inv - fm = torch.zeros(mask.shape, dtype=torch.bfloat16, device="cuda:0") - fm[mask] = float("-inf") - - qt = q.permute(1,0,2).reshape(NH*nt,1,HD) - kve = kv_comb.unsqueeze(0).expand(NH,nt,total,HD).reshape(NH*nt,total,HD) - mb = fm.unsqueeze(0).unsqueeze(2).expand(NH,nt,1,total).reshape(NH*nt,1,total) - ref = F.scaled_dot_product_attention(qt, kve, kve, attn_mask=mb, is_causal=False, scale=ascale).reshape(NH,nt,HD).permute(1,0,2) - - try: - nat = native_sparse_decode_attention(q, swa_cache, inv_sc, si, sl, comp_cache, inv_sc2, ti, tl, sink, BS, ascale, WIN, compress_ratio=4) - c = F.cosine_similarity(ref.flatten().unsqueeze(0).float(), nat.flatten().unsqueeze(0).float()).item() - print(f"tokens={nt} swa={swa_l} topk={topk_l} cosine={c:.6f} {'OK' if c>=0.99 else 'LOW'}") - except Exception as e: - print(f"tokens={nt} swa={swa_l} topk={topk_l} FAILED: {e}") diff --git a/tests/archive/test_stage_a_copy.py b/tests/archive/test_stage_a_copy.py deleted file mode 100644 index 07d61b09..00000000 --- a/tests/archive/test_stage_a_copy.py +++ /dev/null @@ -1,372 +0,0 @@ -""" -Stage A: Bare Q@K^T via tcgen05.mma → TMEM → GMEM -Follows the CUTLASS dense_gemm_persistent.py pattern EXACTLY. -BF16 inputs, FP32 accumulator, TMA load/store, warp specialization. -Single tile (no persistent scheduler), cluster (1,1). -""" -import torch -import cutlass -import cutlass.cute as cute -import cutlass.utils as utils -import cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.cute.runtime import make_ptr -import cuda.bindings.driver as cuda - - -class StageAQKTKernel: - def __init__(self, mma_tiler_mn, use_2cta_instrs=False, use_tma_store=True): - self.acc_dtype = Float32 - self.use_2cta_instrs = use_2cta_instrs - self.mma_tiler_mn = mma_tiler_mn - self.mma_tiler = (*mma_tiler_mn, 1) - self.use_tma_store = use_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4 - self.tma_warp_id = 5 - self.threads_per_cta = 32 * 6 # 192 - self.epilog_sync_bar_id = 1 - self.tmem_alloc_sync_bar_id = 2 - self.tmem_dealloc_sync_bar_id = 3 - - def _create_tiled_mma(self): - return utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.a_major_mode, self.b_major_mode, - self.acc_dtype, self.cta_group, self.mma_tiler_mn, - ) - - def _setup_attributes(self): - tiled_mma = self._create_tiled_mma() - mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2]) - mma_inst_tile_k = 4 - self.mma_tiler = (self.mma_tiler[0], self.mma_tiler[1], mma_inst_shape_k * mma_inst_tile_k) - self.cta_tile_shape_mnk = ( - self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), - self.mma_tiler[1], - self.mma_tiler[2], - ) - self.cluster_layout_vmnk = cute.tiled_divide( - cute.make_layout((1, 1, 1)), (tiled_mma.thr_id.shape,)) - self.num_mcast_ctas_a = 1 - self.num_mcast_ctas_b = 1 - self.is_a_mcast = False - self.is_b_mcast = False - - # Epilogue tile - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, self.use_2cta_instrs, self.c_layout, self.c_dtype) - - # Stage counts: 1 AB stage (single tile, no double-buffer), 1 acc stage, 2 C stages - self.num_ab_stage = 1 - self.num_acc_stage = 1 - self.num_c_stage = 2 - - # SMEM layouts - self.a_smem_layout_staged = utils.sm100.make_smem_layout_a( - tiled_mma, self.mma_tiler, self.a_dtype, self.num_ab_stage) - self.b_smem_layout_staged = utils.sm100.make_smem_layout_b( - tiled_mma, self.mma_tiler, self.b_dtype, self.num_ab_stage) - self.c_smem_layout_staged = utils.sm100.make_smem_layout_epi( - self.c_dtype, self.c_layout, self.epi_tile, self.num_c_stage) - - # TMEM alloc cols - acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) - tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(tCtAcc_fake, arch="sm_100") - - # TMA load bytes - a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) - b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.a_dtype, a_smem_layout) + - cute.size_in_bytes(self.b_dtype, b_smem_layout) - ) * cute.size(tiled_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, - stream: cuda.CUstream): - self.a_dtype = a.element_type - self.b_dtype = b.element_type - self.c_dtype = c.element_type - self.a_major_mode = LayoutEnum.from_tensor(a).mma_major_mode() - self.b_major_mode = LayoutEnum.from_tensor(b).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - tiled_mma = self._create_tiled_mma() - self._setup_attributes() - - # TMA load A - a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) - tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma.thr_id), - a, a_smem_layout, self.mma_tiler, tiled_mma, - self.cluster_layout_vmnk.shape, - ) - - # TMA load B - b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) - tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, tiled_mma.thr_id), - b, b_smem_layout, self.mma_tiler, tiled_mma, - self.cluster_layout_vmnk.shape, - ) - - # TMA store C - epi_smem_layout = cute.select(self.c_smem_layout_staged, mode=[0, 1]) - tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( - cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem_layout, self.epi_tile) - - self._kernel( - tiled_mma, tma_atom_a, tma_tensor_a, tma_atom_b, tma_tensor_b, - tma_atom_c, tma_tensor_c, self.cluster_layout_vmnk, - self.a_smem_layout_staged, self.b_smem_layout_staged, - self.c_smem_layout_staged, self.epi_tile, - ).launch(grid=(1, 1, 1), block=[self.threads_per_cta, 1, 1], stream=stream) - - @cute.kernel - def _kernel(self, tiled_mma, tma_atom_a, mA_mkl, tma_atom_b, mB_nkl, - tma_atom_c, mC_mnl, cluster_layout_vmnk, - a_smem_layout_staged, b_smem_layout_staged, c_smem_layout_staged, epi_tile): - warp_idx = cute.arch.warp_idx() - warp_idx = cute.arch.make_warp_uniform(warp_idx) - tidx, _, _ = cute.arch.thread_idx() - use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2 - is_leader_cta = True # single CTA, always leader - - # Prefetch TMA descriptors - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_atom_a) - cpasync.prefetch_descriptor(tma_atom_b) - cpasync.prefetch_descriptor(tma_atom_c) - - # ── Shared storage ─────────────────────────────────── - @cute.struct - class SharedStorage: - ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc_mbar: cutlass.Int64 - tmem_holding_buf: cutlass.Int32 - - smem = utils.SmemAllocator() - storage = smem.allocate(SharedStorage) - - # AB pipeline - ab_producer, ab_consumer = pipeline.PipelineTmaUmma.create( - barrier_storage=storage.ab_full_mbar_ptr.data_ptr(), - num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, - cta_layout_vmnk=cluster_layout_vmnk, - defer_sync=True, - ).make_participants() - - # ACC pipeline - acc_pipeline = pipeline.PipelineUmmaAsync.create( - barrier_storage=storage.acc_full_mbar_ptr.data_ptr(), - num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta_instrs else 1)), - cta_layout_vmnk=cluster_layout_vmnk, - defer_sync=True, - ) - - # TMEM allocator - tmem_alloc_barrier = pipeline.NamedBarrier( - barrier_id=self.tmem_alloc_sync_bar_id, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id)), - ) - tmem = utils.TmemAllocator( - storage.tmem_holding_buf.ptr, - barrier_for_retrieve=tmem_alloc_barrier, - allocator_warp_id=self.epilogue_warp_id[0], - is_two_cta=use_2cta_instrs, - two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr, - ) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cluster_layout_vmnk, is_relaxed=True) - - # SMEM tensors - sA = smem.allocate_tensor( - element_type=self.a_dtype, layout=a_smem_layout_staged.outer, - byte_alignment=128, swizzle=a_smem_layout_staged.inner) - sB = smem.allocate_tensor( - element_type=self.b_dtype, layout=b_smem_layout_staged.outer, - byte_alignment=128, swizzle=b_smem_layout_staged.inner) - sC = smem.allocate_tensor( - element_type=self.c_dtype, layout=c_smem_layout_staged.outer, - byte_alignment=128, swizzle=c_smem_layout_staged.inner) - - # Partition global tensors - gA_mkl = cute.local_tile(mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)) - gB_nkl = cute.local_tile(mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None)) - gC_mnl = cute.local_tile(mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None)) - k_tile_cnt = cute.size(gA_mkl, mode=[3]) - - # Partition for TiledMMA - thr_mma = tiled_mma.get_slice(0) # leader CTA - tCgA = thr_mma.partition_A(gA_mkl) - tCgB = thr_mma.partition_B(gB_nkl) - tCgC = thr_mma.partition_C(gC_mnl) - - # TMA partition A/B - a_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape) - tAsA, tAgA = cpasync.tma_partition( - tma_atom_a, 0, a_cta_layout, - cute.group_modes(sA, 0, 3), cute.group_modes(tCgA, 0, 3)) - b_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape) - tBsB, tBgB = cpasync.tma_partition( - tma_atom_b, 0, b_cta_layout, - cute.group_modes(sB, 0, 3), cute.group_modes(tCgB, 0, 3)) - - # Slice to tile coord (0, 0, 0) - tAgA_slice = tAgA[(None, 0, None, 0)] - tBgB_slice = tBgB[(None, 0, None, 0)] - - # MMA fragments - tCrA = tiled_mma.make_fragment_A(sA) - tCrB = tiled_mma.make_fragment_B(sB) - acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) - tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cluster_layout_vmnk) - - # ══════════════════════════════════════════════════════════ - # TMA LOAD WARP (warp 5) - # ══════════════════════════════════════════════════════════ - if warp_idx == self.tma_warp_id: - ab_producer.reset() - peek_ab_empty_status = ab_producer.try_acquire() - - for k_tile in cutlass.range(k_tile_cnt, unroll=1): - handle = ab_producer.acquire_and_advance(peek_ab_empty_status) - cute.copy(tma_atom_a, tAgA_slice[(None, handle.count)], tAsA[(None, handle.index)], - tma_bar_ptr=handle.barrier) - cute.copy(tma_atom_b, tBgB_slice[(None, handle.count)], tBsB[(None, handle.index)], - tma_bar_ptr=handle.barrier) - peek_ab_empty_status = cutlass.Boolean(1) - if handle.count + 1 < k_tile_cnt: - peek_ab_empty_status = ab_producer.try_acquire() - - ab_producer.tail() - - # ══════════════════════════════════════════════════════════ - # MMA WARP (warp 4) - # ══════════════════════════════════════════════════════════ - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) - tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) - tCtAcc = tCtAcc_base[(None, None, None, 0)] - - ab_consumer.reset() - peek_ab_full_status = cutlass.Boolean(1) - if is_leader_cta: - peek_ab_full_status = ab_consumer.try_wait() - - acc_producer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Producer, self.num_acc_stage) - if is_leader_cta: - acc_pipeline.producer_acquire(acc_producer_state) - tiled_mma.set(tcgen05.Field.ACCUMULATE, False) - - for k_tile in range(k_tile_cnt): - if is_leader_cta: - handle = ab_consumer.wait_and_advance(peek_ab_full_status) - num_kblocks = cute.size(tCrA, mode=[2]) - for kblk_idx in cutlass.range(num_kblocks, unroll_full=True): - kblk_crd = (None, None, kblk_idx, handle.index) - cute.gemm(tiled_mma, tCtAcc, tCrA[kblk_crd], tCrB[kblk_crd], tCtAcc) - tiled_mma.set(tcgen05.Field.ACCUMULATE, True) - handle.release() - peek_ab_full_status = cutlass.Boolean(1) - if handle.count + 1 < k_tile_cnt: - peek_ab_full_status = ab_consumer.try_wait() - - if is_leader_cta: - acc_pipeline.producer_commit(acc_producer_state) - acc_producer_state.advance() - acc_pipeline.producer_tail(acc_producer_state) - - # ══════════════════════════════════════════════════════════ - # EPILOGUE WARPS (0..3) - # ══════════════════════════════════════════════════════════ - if warp_idx < self.mma_warp_id: - tmem.allocate(self.num_tmem_alloc_cols) - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) - tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) - - acc_consumer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Consumer, self.num_acc_stage) - - c_producer_group = pipeline.CooperativeGroup( - pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) - c_pipeline = pipeline.PipelineTmaStore.create( - num_stages=self.num_c_stage, producer_group=c_producer_group) - - # Use the reference epilogue implementation - mma_tile_coord_mnl = (0, 0, 0) - epilogue_op = const_expr(lambda x: x) - num_tiles_executed = 0 - - acc_consumer_state = utils.gemm.sm100.epilogue_tma_store( - self, tidx, warp_idx, tma_atom_c, tCtAcc_base, sC, tCgC, - epi_tile, num_tiles_executed, epilogue_op, - mma_tile_coord_mnl, acc_consumer_state, acc_pipeline, c_pipeline) - - c_pipeline.producer_tail() - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def test_stage_a(): - """Test Stage A: Q @ K^T → TMEM → GMEM""" - device = torch.device("cuda") - torch.manual_seed(42) - - m, n, k = 128, 128, 512 - - # Tensors must be 3D (M, K, L) for the CUTLASS pattern - a = torch.randn(m, k, 1, dtype=torch.bfloat16, device="cuda") - b = torch.randn(n, k, 1, dtype=torch.bfloat16, device="cuda") - c = torch.zeros(m, n, 1, dtype=torch.bfloat16, device="cuda") - - ref = a[:, :, 0].float() @ b[:, :, 0].float().T - - # Create cute tensors - import cutlass.torch as cutlass_torch - mA = cutlass_torch.from_dlpack(a).mark_layout_dynamic( - leading_dim=cutlass_torch.get_leading_dim(a)) - mB = cutlass_torch.from_dlpack(b).mark_layout_dynamic( - leading_dim=cutlass_torch.get_leading_dim(b)) - mC = cutlass_torch.from_dlpack(c).mark_layout_dynamic( - leading_dim=cutlass_torch.get_leading_dim(c)) - - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - kernel = StageAQKTKernel(mma_tiler_mn=(128, 128), use_2cta_instrs=False, use_tma_store=True) - compiled = cute.compile(kernel, mA, mB, mC, stream) - - # Run with the same tensors - compiled(mA, mB, mC, stream) - torch.cuda.synchronize() - - output = c[:, :, 0].float() - cos = torch.nn.functional.cosine_similarity( - output.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - max_err = (output - ref).abs().max().item() - - print("Stage A: Q({},{}) @ K^T({}, {}) -> S({}, {})".format(m, k, k, n, m, n)) - print(" Cosine: {:.6f}, Max error: {:.6f}".format(cos, max_err)) - print(" {}".format("PASS" if cos >= 0.99 else "FAIL")) - return cos - - -if __name__ == "__main__": - test_stage_a() diff --git a/tests/archive/test_stage_a_minimal.py b/tests/archive/test_stage_a_minimal.py deleted file mode 100644 index 7ce3256d..00000000 --- a/tests/archive/test_stage_a_minimal.py +++ /dev/null @@ -1,395 +0,0 @@ -""" -Stage A Minimal: Just tcgen05.mma with TMEM accumulator, epilogue writes to GMEM. -No TMA loads (use regular SMEM loads). No pipeline. Just validate the MMA path. -""" - -import torch -import cutlass -import cutlass.cute as cute -import cutlass.utils as utils -import cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import tcgen05 -from cutlass import BFloat16, Float32 -import cuda.bindings.driver as cuda - -# Q: (128, 512) BF16, K^T: (64, 512) BF16 -> S: (128, 64) FP32 -# These are the minimum valid tile sizes for tcgen05.mma -M = 128 -N = 64 -K = 512 - - -class MinimalMMAKernel: - def __init__(self): - self.cta_group = tcgen05.CtaGroup.ONE - self.mma_tiler_mn = (M, N) - self.threads_per_cta = 192 # 6 warps - self.epilog_warp_ids = [0, 1, 2, 3] - self.mma_warp_id = 4 - self.tma_warp_id = 5 - - @cute.jit - def __call__(self, a_ptr, b_ptr, c_ptr, problem_m, problem_n, problem_k, stream): - a_dtype = a_ptr.value_type - b_dtype = b_ptr.value_type - c_dtype = c_ptr.value_type - acc_dtype = Float32 - - m, n, k = problem_m, problem_n, problem_k - - # TiledMMA - tiled_mma = sm100_utils.make_trivial_tiled_mma( - a_dtype, b_dtype, - tcgen05.OperandMajorMode.K, - tcgen05.OperandMajorMode.K, - acc_dtype, - self.cta_group, - self.mma_tiler_mn, - ) - - atom_thr_size = cute.size(tiled_mma.thr_id.shape) - self.atom_thr_size = atom_thr_size - mma_tiler = (self.mma_tiler_mn[0], self.mma_tiler_mn[1], cute.size(tiled_mma.shape_mnk, mode=[2]) * 4) - self.mma_tiler = mma_tiler - cta_tile_shape_mnk = (mma_tiler[0] // atom_thr_size, mma_tiler[1], mma_tiler[2]) - self.cta_tile_shape_mnk = cta_tile_shape_mnk - - # SMEM layouts - num_ab_stages = 1 - a_smem_layout = sm100_utils.make_smem_layout_a(tiled_mma, mma_tiler, a_dtype, num_ab_stages) - b_smem_layout = sm100_utils.make_smem_layout_b(tiled_mma, mma_tiler, b_dtype, num_ab_stages) - - # Epilogue tile - c_layout_enum = utils.LayoutEnum.ROW_MAJOR - epi_tile = sm100_utils.compute_epilogue_tile_shape( - cta_tile_shape_mnk, False, c_layout_enum, c_dtype) - self.epi_tile = epi_tile - - c_smem_layout = sm100_utils.make_smem_layout_epi(c_dtype, c_layout_enum, epi_tile, 2) - self.c_smem_layout = c_smem_layout - - # TMEM columns - self.num_accumulator_tmem_cols = cta_tile_shape_mnk[1] - - # GMEM tensors - a_gmem_layout = cute.make_ordered_layout((m, k), order=(1, 0)) - b_gmem_layout = cute.make_ordered_layout((n, k), order=(1, 0)) - c_gmem_layout = cute.make_ordered_layout((m, n), order=(1, 0)) - gA = cute.make_tensor(a_ptr, a_gmem_layout) - gB = cute.make_tensor(b_ptr, b_gmem_layout) - gC = cute.make_tensor(c_ptr, c_gmem_layout) - - # TMA descriptors - a_smem_layout_one = cute.slice_(a_smem_layout, (None, None, None, 0)) - b_smem_layout_one = cute.slice_(b_smem_layout, (None, None, None, 0)) - c_smem_layout_one = cute.slice_(c_smem_layout, (None, None, 0)) - - tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( - sm100_utils.cluster_shape_to_tma_atom_A((1, 1), tiled_mma.thr_id), - gA, a_smem_layout_one, mma_tiler, tiled_mma, - (1, 1, 1), - ) - tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( - sm100_utils.cluster_shape_to_tma_atom_B((1, 1), tiled_mma.thr_id), - gB, b_smem_layout_one, mma_tiler, tiled_mma, - (1, 1, 1), - ) - tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( - cpasync.CopyBulkTensorTileS2GOp(), - gC, c_smem_layout_one, epi_tile, - ) - - # Pipeline barriers - a_copy_size = cute.size_in_bytes(a_dtype, a_smem_layout_one) - b_copy_size = cute.size_in_bytes(b_dtype, b_smem_layout_one) - tma_load_bytes = (a_copy_size + b_copy_size) * atom_thr_size - - self.tma_load_bytes = tma_load_bytes - - # Named barriers - self.epilog_sync_barrier = pipeline.NamedBarrier( - barrier_id=1, - num_threads=32 * len(self.epilog_warp_ids), - ) - self.tmem_alloc_barrier = pipeline.NamedBarrier( - barrier_id=2, - num_threads=32 * (1 + len(self.epilog_warp_ids)), - ) - - @cute.struct - class SharedStorage: - ab_full_mbar: cute.struct.MemRange[cutlass.Int64, num_ab_stages] - ab_empty_mbar: cute.struct.MemRange[cutlass.Int64, num_ab_stages] - acc_full_mbar: cute.struct.MemRange[cutlass.Int64, 1] - acc_empty_mbar: cute.struct.MemRange[cutlass.Int64, 1] - tmem_dealloc_mbar: cutlass.Int64 - tmem_holding_buf: cutlass.Int32 - sA: cute.struct.Align[cute.struct.MemRange[a_dtype, cute.cosize(a_smem_layout.outer)], 1024] - sB: cute.struct.Align[cute.struct.MemRange[b_dtype, cute.cosize(b_smem_layout.outer)], 1024] - sC: cute.struct.Align[cute.struct.MemRange[c_dtype, cute.cosize(c_smem_layout.outer)], 1024] - - self.shared_storage = SharedStorage - - # Cluster - self.cluster_shape_mn = (1, 1) - cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1, 1, 1)), (tiled_mma.thr_id.shape,)) - self.cluster_layout_vmnk = cluster_layout_vmnk - - self._kernel( - tiled_mma, tma_atom_a, tma_tensor_a, tma_atom_b, tma_tensor_b, - tma_atom_c, tma_tensor_c, cluster_layout_vmnk, - a_smem_layout, b_smem_layout, c_smem_layout, epi_tile, - gA, gB, gC, mma_tiler, - ).launch(grid=(1, 1, 1), block=[self.threads_per_cta, 1, 1], stream=stream) - - @cute.kernel - def _kernel(self, tiled_mma, tma_atom_a, mA, tma_atom_b, mB, - tma_atom_c, mC, cluster_layout_vmnk, - a_smem_layout, b_smem_layout, c_smem_layout, epi_tile, - gA, gB, gC, mma_tiler): - warp_idx = cute.arch.warp_idx() - warp_idx = cute.arch.make_warp_uniform(warp_idx) - tidx, _, _ = cute.arch.thread_idx() - - use_2cta = cute.size(tiled_mma.thr_id.shape) == 2 - - smem = utils.SmemAllocator() - storage = smem.allocate(self.shared_storage) - - # AB pipeline - ab_pipeline = pipeline.PipelineTmaUmma.create( - barrier_storage=storage.ab_full_mbar.data_ptr(), - num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.tma_load_bytes, - cta_layout_vmnk=cluster_layout_vmnk, - defer_sync=True, - ) - - # Accumulator pipeline - acc_pipeline = pipeline.PipelineUmmaAsync.create( - barrier_storage=storage.acc_full_mbar.data_ptr(), - num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, - 32 * len(self.epilog_warp_ids) * (2 if use_2cta else 1)), - cta_layout_vmnk=cluster_layout_vmnk, - defer_sync=True, - ) - - # TMEM allocator - tmem = utils.TmemAllocator( - storage.tmem_holding_buf.ptr, - barrier_for_retrieve=self.tmem_alloc_barrier, - allocator_warp_id=self.epilog_warp_ids[0], - is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr, - ) - - pipeline.pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) - - # SMEM tensors - sA = storage.sA.get_tensor(a_smem_layout.outer, swizzle=a_smem_layout.inner) - sB = storage.sB.get_tensor(b_smem_layout.outer, swizzle=b_smem_layout.inner) - sC = storage.sC.get_tensor(c_smem_layout.outer, swizzle=c_smem_layout.inner) - - # gC tiled for epilogue partition - gC_tiled = cute.local_tile(gC, (mma_tiler[0], mma_tiler[1]), (0, 0)) - - # Partition for TMA — use TMA tensors directly (they have static shape) - thr_mma = tiled_mma.get_slice(0) - tCgC = thr_mma.partition_C(gC_tiled) - - tAsA, tAgA = cpasync.tma_partition( - tma_atom_a, 0, cute.make_layout(1), - cute.group_modes(sA, 0, 3), - cute.group_modes(mA, 0, 3), - ) - tBsB, tBgB = cpasync.tma_partition( - tma_atom_b, 0, cute.make_layout(1), - cute.group_modes(sB, 0, 3), - cute.group_modes(mB, 0, 3), - ) - - # MMA fragments - tCrA = tiled_mma.make_fragment_A(sA) - tCrB = tiled_mma.make_fragment_B(sB) - - # TMEM accumulator - acc_shape = tiled_mma.partition_shape_C(mma_tiler[:2]) - tCtAcc_all = tiled_mma.make_fragment_C(cute.append(acc_shape, 1)) - - k_tile_cnt = cute.size(mA, mode=[3]) - - # gC tiled for epilogue partition - gC_tiled = cute.local_tile(gC, (mma_tiler[0], mma_tiler[1]), (0, 0)) - - pipeline.pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) - - # ══════════════════════════════════════════════════════════ - # TMA LOAD WARP (warp 5) - # ══════════════════════════════════════════════════════════ - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_atom_a) - cpasync.prefetch_descriptor(tma_atom_b) - cpasync.prefetch_descriptor(tma_atom_c) - - ab_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, 1) - for k_tile in cutlass.range(k_tile_cnt, unroll=1): - ab_pipeline.producer_acquire(ab_state) - cute.copy(tma_atom_a, tAgA[(None, ab_state.count)], tAsA[(None, ab_state.index)], - tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_state)) - cute.copy(tma_atom_b, tBgB[(None, ab_state.count)], tBsB[(None, ab_state.index)], - tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_state)) - ab_state.advance() - ab_pipeline.producer_tail(ab_state) - - # ══════════════════════════════════════════════════════════ - # MMA WARP (warp 4) - # ══════════════════════════════════════════════════════════ - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - acc_ptr = tmem.retrieve_ptr(Float32) - tCtAcc_base = cute.make_tensor(acc_ptr, tCtAcc_all.layout) - tCtAcc = tCtAcc_base[(None, None, None, 0)] - - ab_cstate = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, 1) - acc_pstate = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, 1) - - acc_pipeline.producer_acquire(acc_pstate) - tiled_mma.set(tcgen05.Field.ACCUMULATE, False) - - for k_tile in range(k_tile_cnt): - ab_pipeline.consumer_wait(ab_cstate, cutlass.Boolean(1)) - for kblock in cutlass.range(cute.size(tCrA, mode=[2]), unroll_full=True): - coord = (None, None, kblock, ab_cstate.index) - cute.gemm(tiled_mma, tCtAcc, tCrA[coord], tCrB[coord], tCtAcc) - tiled_mma.set(tcgen05.Field.ACCUMULATE, True) - ab_pipeline.consumer_release(ab_cstate) - ab_cstate.advance() - - acc_pipeline.producer_commit(acc_pstate) - acc_pstate.advance() - acc_pipeline.producer_tail(acc_pstate) - - # ══════════════════════════════════════════════════════════ - # EPILOGUE WARPS (0..3) - # ══════════════════════════════════════════════════════════ - if warp_idx < self.mma_warp_id: - tmem.allocate(self.num_accumulator_tmem_cols) - tmem.wait_for_alloc() - - acc_ptr = tmem.retrieve_ptr(Float32) - tCtAcc_base = cute.make_tensor(acc_ptr, tCtAcc_all.layout) - - c_layout_enum = utils.LayoutEnum.ROW_MAJOR - c_dtype = BFloat16 - - # TMEM→reg - copy_atom_t2r = sm100_utils.get_tmem_load_op( - self.cta_tile_shape_mnk, c_layout_enum, c_dtype, Float32, epi_tile, False) - tAcc_epi = cute.flat_divide(tCtAcc_base[((None, None), 0, 0, None)], epi_tile) - tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tAcc_epi[(None, None, 0, 0, 0)]) - thr_t2r = tiled_copy_t2r.get_slice(tidx) - tTR_tAcc = thr_t2r.partition_S(tAcc_epi) - tTR_rAcc = cute.make_rmem_tensor( - thr_t2r.partition_D( - cute.flat_divide(tCgC[((None, None), 0, 0, None, None)], epi_tile) - )[(None, None, None, 0, 0, 0, 0)].shape, Float32) - tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, c_dtype) - - # reg→SMEM - copy_atom_r2s = sm100_utils.get_smem_store_op(c_layout_enum, c_dtype, Float32, tiled_copy_t2r) - tiled_copy_r2s = cute.make_tiled_copy_D(copy_atom_r2s, tiled_copy_t2r) - thr_r2s = tiled_copy_r2s.get_slice(tidx) - tRS_sC = thr_r2s.partition_D(sC) - tRS_rC = tiled_copy_r2s.retile(tTR_rC) - - # SMEM→GMEM (TMA) - gC_epi = cute.flat_divide(tCgC[((None, None), 0, 0, None, None)], epi_tile) - bSG_sC, bSG_gC = cpasync.tma_partition( - tma_atom_c, 0, cute.make_layout(1), - cute.group_modes(sC, 0, 2), - cute.group_modes(gC_epi, 0, 2)) - - acc_cstate = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, 1) - c_pipeline = pipeline.PipelineTmaStore.create( - num_stages=2, - producer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, 32 * len(self.epilog_warp_ids))) - - acc_pipeline.consumer_wait(acc_cstate) - - tTR_tAcc_g = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) - bSG_gC_g = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC)) - - for subtile in cutlass.range(cute.size(tTR_tAcc_g.shape, mode=[3])): - cute.copy(tiled_copy_t2r, tTR_tAcc_g[(None, None, None, subtile)], tTR_rAcc) - acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load() - tRS_rC.store(acc_vec.to(c_dtype)) - - c_buf = subtile % 2 - cute.copy(tiled_copy_r2s, tRS_rC, tRS_sC[(None, None, None, c_buf)]) - cute.arch.fence_proxy("async.shared", space="cta") - self.epilog_sync_barrier.arrive_and_wait() - - if warp_idx == self.epilog_warp_ids[0]: - cute.copy(tma_atom_c, bSG_sC[(None, c_buf)], bSG_gC_g[(None, subtile)]) - c_pipeline.producer_commit() - c_pipeline.producer_acquire() - self.epilog_sync_barrier.arrive_and_wait() - - acc_pipeline.consumer_release(acc_cstate) - tmem.relinquish_alloc_permit() - self.epilog_sync_barrier.arrive_and_wait() - tmem.free(acc_ptr) - c_pipeline.producer_tail() - - -from cutlass.cute.runtime import make_ptr -from cutlass.cute.nvgpu import cpasync -import cutlass.utils.blackwell_helpers as sm100_utils - -def test_minimal_mma(): - device = torch.device("cuda") - torch.manual_seed(42) - - prob_m, prob_n, prob_k = M, N, K - - tA = torch.randn(prob_m, prob_k, dtype=torch.bfloat16, device=device) - tB = torch.randn(prob_n, prob_k, dtype=torch.bfloat16, device=device) - ref = torch.matmul(tA.to(torch.float32), tB.to(torch.float32).T) - tC = torch.zeros(prob_m, prob_n, dtype=torch.bfloat16, device=device) - - a_ptr = make_ptr(BFloat16, 0, cute.AddressSpace.gmem, assumed_align=16) - b_ptr = make_ptr(BFloat16, 0, cute.AddressSpace.gmem, assumed_align=16) - c_ptr = make_ptr(BFloat16, 0, cute.AddressSpace.gmem, assumed_align=16) - - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - kernel = MinimalMMAKernel() - compiled = cute.compile(kernel, a_ptr, b_ptr, c_ptr, - cutlass.Int32(prob_m), cutlass.Int32(prob_n), cutlass.Int32(prob_k), stream) - - a_ptr_r = make_ptr(BFloat16, tA.data_ptr(), cute.AddressSpace.gmem, assumed_align=16) - b_ptr_r = make_ptr(BFloat16, tB.data_ptr(), cute.AddressSpace.gmem, assumed_align=16) - c_ptr_r = make_ptr(BFloat16, tC.data_ptr(), cute.AddressSpace.gmem, assumed_align=16) - - compiled(a_ptr_r, b_ptr_r, c_ptr_r, prob_m, prob_n, prob_k, stream) - torch.cuda.synchronize() - - output = tC.to(torch.float32) - cos = torch.nn.functional.cosine_similarity( - output.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - max_err = (output - ref).abs().max().item() - - print(f"Minimal MMA: A({prob_m},{prob_k}) @ B^T({prob_k},{prob_n}) → C({prob_m},{prob_n})") - print(f" Cosine: {cos:.6f}, Max error: {max_err:.6f}") - print(f" {'✅ PASS' if cos >= 0.99 else '❌ FAIL'}") - return cos - - -if __name__ == "__main__": - test_minimal_mma() diff --git a/tests/archive/test_stage_a_pv_created.py b/tests/archive/test_stage_a_pv_created.py deleted file mode 100644 index aa5fe4d1..00000000 --- a/tests/archive/test_stage_a_pv_created.py +++ /dev/null @@ -1,376 +0,0 @@ -""" -Stage A: Bare Q@K^T via tcgen05.mma → TMEM → GMEM -Follows the CUTLASS dense_gemm_persistent.py pattern EXACTLY. -BF16 inputs, FP32 accumulator, TMA load/store, warp specialization. -Single tile (no persistent scheduler), cluster (1,1). -""" -import torch -import cutlass -import cutlass.cute as cute -import cutlass.utils as utils -import cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.cute.runtime import make_ptr -import cuda.bindings.driver as cuda - - -class StageAQKTKernel: - def __init__(self, mma_tiler_mn, use_2cta_instrs=False, use_tma_store=True): - self.acc_dtype = Float32 - self.use_2cta_instrs = use_2cta_instrs - self.mma_tiler_mn = mma_tiler_mn - self.mma_tiler = (*mma_tiler_mn, 1) - self.use_tma_store = use_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4 - self.tma_warp_id = 5 - self.threads_per_cta = 32 * 6 # 192 - self.epilog_sync_bar_id = 1 - self.tmem_alloc_sync_bar_id = 2 - self.tmem_dealloc_sync_bar_id = 3 - - def _create_tiled_mma(self): - return utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.a_major_mode, self.b_major_mode, - self.acc_dtype, self.cta_group, self.mma_tiler_mn, - ) - - def _setup_attributes(self): - # Create pv_mma but DO NOT use it - pv_mma = utils.sm100.make_trivial_tiled_mma(self.a_dtype, self.b_dtype, cute.nvgpu.OperandMajorMode.K, self.b_major_mode, self.acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - tiled_mma = self._create_tiled_mma() - mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2]) - mma_inst_tile_k = 4 - self.mma_tiler = (self.mma_tiler[0], self.mma_tiler[1], mma_inst_shape_k * mma_inst_tile_k) - self.cta_tile_shape_mnk = ( - self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), - self.mma_tiler[1], - self.mma_tiler[2], - ) - self.cluster_layout_vmnk = cute.tiled_divide( - cute.make_layout((1, 1, 1)), (tiled_mma.thr_id.shape,)) - self.num_mcast_ctas_a = 1 - self.num_mcast_ctas_b = 1 - self.is_a_mcast = False - self.is_b_mcast = False - - # Epilogue tile - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, self.use_2cta_instrs, self.c_layout, self.c_dtype) - - # Stage counts: 1 AB stage (single tile, no double-buffer), 1 acc stage, 2 C stages - self.num_ab_stage = 1 - self.num_acc_stage = 1 - self.num_c_stage = 2 - - # SMEM layouts - self.a_smem_layout_staged = utils.sm100.make_smem_layout_a( - tiled_mma, self.mma_tiler, self.a_dtype, self.num_ab_stage) - self.b_smem_layout_staged = utils.sm100.make_smem_layout_b( - tiled_mma, self.mma_tiler, self.b_dtype, self.num_ab_stage) - self.c_smem_layout_staged = utils.sm100.make_smem_layout_epi( - self.c_dtype, self.c_layout, self.epi_tile, self.num_c_stage) - - # TMEM alloc cols - acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) - tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(tCtAcc_fake, arch="sm_100") - - # TMA load bytes - a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) - b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.a_dtype, a_smem_layout) + - cute.size_in_bytes(self.b_dtype, b_smem_layout) - ) * cute.size(tiled_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, - stream: cuda.CUstream): - self.a_dtype = a.element_type - self.b_dtype = b.element_type - self.c_dtype = c.element_type - self.a_major_mode = LayoutEnum.from_tensor(a).mma_major_mode() - self.b_major_mode = LayoutEnum.from_tensor(b).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - # Create pv_mma but DO NOT use it - pv_mma = utils.sm100.make_trivial_tiled_mma(self.a_dtype, self.b_dtype, cute.nvgpu.OperandMajorMode.K, self.b_major_mode, self.acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - tiled_mma = self._create_tiled_mma() - self._setup_attributes() - - # TMA load A - a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) - tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma.thr_id), - a, a_smem_layout, self.mma_tiler, tiled_mma, - self.cluster_layout_vmnk.shape, - ) - - # TMA load B - b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) - tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, tiled_mma.thr_id), - b, b_smem_layout, self.mma_tiler, tiled_mma, - self.cluster_layout_vmnk.shape, - ) - - # TMA store C - epi_smem_layout = cute.select(self.c_smem_layout_staged, mode=[0, 1]) - tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( - cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem_layout, self.epi_tile) - - self._kernel( - tiled_mma, tma_atom_a, tma_tensor_a, tma_atom_b, tma_tensor_b, - tma_atom_c, tma_tensor_c, self.cluster_layout_vmnk, - self.a_smem_layout_staged, self.b_smem_layout_staged, - self.c_smem_layout_staged, self.epi_tile, - ).launch(grid=(1, 1, 1), block=[self.threads_per_cta, 1, 1], stream=stream) - - @cute.kernel - def _kernel(self, tiled_mma, tma_atom_a, mA_mkl, tma_atom_b, mB_nkl, - tma_atom_c, mC_mnl, cluster_layout_vmnk, - a_smem_layout_staged, b_smem_layout_staged, c_smem_layout_staged, epi_tile): - warp_idx = cute.arch.warp_idx() - warp_idx = cute.arch.make_warp_uniform(warp_idx) - tidx, _, _ = cute.arch.thread_idx() - use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2 - is_leader_cta = True # single CTA, always leader - - # Prefetch TMA descriptors - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_atom_a) - cpasync.prefetch_descriptor(tma_atom_b) - cpasync.prefetch_descriptor(tma_atom_c) - - # ── Shared storage ─────────────────────────────────── - @cute.struct - class SharedStorage: - ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc_mbar: cutlass.Int64 - tmem_holding_buf: cutlass.Int32 - - smem = utils.SmemAllocator() - storage = smem.allocate(SharedStorage) - - # AB pipeline - ab_producer, ab_consumer = pipeline.PipelineTmaUmma.create( - barrier_storage=storage.ab_full_mbar_ptr.data_ptr(), - num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, - cta_layout_vmnk=cluster_layout_vmnk, - defer_sync=True, - ).make_participants() - - # ACC pipeline - acc_pipeline = pipeline.PipelineUmmaAsync.create( - barrier_storage=storage.acc_full_mbar_ptr.data_ptr(), - num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta_instrs else 1)), - cta_layout_vmnk=cluster_layout_vmnk, - defer_sync=True, - ) - - # TMEM allocator - tmem_alloc_barrier = pipeline.NamedBarrier( - barrier_id=self.tmem_alloc_sync_bar_id, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id)), - ) - tmem = utils.TmemAllocator( - storage.tmem_holding_buf.ptr, - barrier_for_retrieve=tmem_alloc_barrier, - allocator_warp_id=self.epilogue_warp_id[0], - is_two_cta=use_2cta_instrs, - two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr, - ) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cluster_layout_vmnk, is_relaxed=True) - - # SMEM tensors - sA = smem.allocate_tensor( - element_type=self.a_dtype, layout=a_smem_layout_staged.outer, - byte_alignment=128, swizzle=a_smem_layout_staged.inner) - sB = smem.allocate_tensor( - element_type=self.b_dtype, layout=b_smem_layout_staged.outer, - byte_alignment=128, swizzle=b_smem_layout_staged.inner) - sC = smem.allocate_tensor( - element_type=self.c_dtype, layout=c_smem_layout_staged.outer, - byte_alignment=128, swizzle=c_smem_layout_staged.inner) - - # Partition global tensors - gA_mkl = cute.local_tile(mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)) - gB_nkl = cute.local_tile(mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None)) - gC_mnl = cute.local_tile(mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None)) - k_tile_cnt = cute.size(gA_mkl, mode=[3]) - - # Partition for TiledMMA - thr_mma = tiled_mma.get_slice(0) # leader CTA - tCgA = thr_mma.partition_A(gA_mkl) - tCgB = thr_mma.partition_B(gB_nkl) - tCgC = thr_mma.partition_C(gC_mnl) - - # TMA partition A/B - a_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape) - tAsA, tAgA = cpasync.tma_partition( - tma_atom_a, 0, a_cta_layout, - cute.group_modes(sA, 0, 3), cute.group_modes(tCgA, 0, 3)) - b_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape) - tBsB, tBgB = cpasync.tma_partition( - tma_atom_b, 0, b_cta_layout, - cute.group_modes(sB, 0, 3), cute.group_modes(tCgB, 0, 3)) - - # Slice to tile coord (0, 0, 0) - tAgA_slice = tAgA[(None, 0, None, 0)] - tBgB_slice = tBgB[(None, 0, None, 0)] - - # MMA fragments - tCrA = tiled_mma.make_fragment_A(sA) - tCrB = tiled_mma.make_fragment_B(sB) - acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) - tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cluster_layout_vmnk) - - # ══════════════════════════════════════════════════════════ - # TMA LOAD WARP (warp 5) - # ══════════════════════════════════════════════════════════ - if warp_idx == self.tma_warp_id: - ab_producer.reset() - peek_ab_empty_status = ab_producer.try_acquire() - - for k_tile in cutlass.range(k_tile_cnt, unroll=1): - handle = ab_producer.acquire_and_advance(peek_ab_empty_status) - cute.copy(tma_atom_a, tAgA_slice[(None, handle.count)], tAsA[(None, handle.index)], - tma_bar_ptr=handle.barrier) - cute.copy(tma_atom_b, tBgB_slice[(None, handle.count)], tBsB[(None, handle.index)], - tma_bar_ptr=handle.barrier) - peek_ab_empty_status = cutlass.Boolean(1) - if handle.count + 1 < k_tile_cnt: - peek_ab_empty_status = ab_producer.try_acquire() - - ab_producer.tail() - - # ══════════════════════════════════════════════════════════ - # MMA WARP (warp 4) - # ══════════════════════════════════════════════════════════ - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) - tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) - tCtAcc = tCtAcc_base[(None, None, None, 0)] - - ab_consumer.reset() - peek_ab_full_status = cutlass.Boolean(1) - if is_leader_cta: - peek_ab_full_status = ab_consumer.try_wait() - - acc_producer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Producer, self.num_acc_stage) - if is_leader_cta: - acc_pipeline.producer_acquire(acc_producer_state) - tiled_mma.set(tcgen05.Field.ACCUMULATE, False) - - for k_tile in range(k_tile_cnt): - if is_leader_cta: - handle = ab_consumer.wait_and_advance(peek_ab_full_status) - num_kblocks = cute.size(tCrA, mode=[2]) - for kblk_idx in cutlass.range(num_kblocks, unroll_full=True): - kblk_crd = (None, None, kblk_idx, handle.index) - cute.gemm(tiled_mma, tCtAcc, tCrA[kblk_crd], tCrB[kblk_crd], tCtAcc) - tiled_mma.set(tcgen05.Field.ACCUMULATE, True) - handle.release() - peek_ab_full_status = cutlass.Boolean(1) - if handle.count + 1 < k_tile_cnt: - peek_ab_full_status = ab_consumer.try_wait() - - if is_leader_cta: - acc_pipeline.producer_commit(acc_producer_state) - acc_producer_state.advance() - acc_pipeline.producer_tail(acc_producer_state) - - # ══════════════════════════════════════════════════════════ - # EPILOGUE WARPS (0..3) - # ══════════════════════════════════════════════════════════ - if warp_idx < self.mma_warp_id: - tmem.allocate(self.num_tmem_alloc_cols) - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) - tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) - - acc_consumer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Consumer, self.num_acc_stage) - - c_producer_group = pipeline.CooperativeGroup( - pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) - c_pipeline = pipeline.PipelineTmaStore.create( - num_stages=self.num_c_stage, producer_group=c_producer_group) - - # Use the reference epilogue implementation - mma_tile_coord_mnl = (0, 0, 0) - epilogue_op = const_expr(lambda x: x) - num_tiles_executed = 0 - - acc_consumer_state = utils.gemm.sm100.epilogue_tma_store( - self, tidx, warp_idx, tma_atom_c, tCtAcc_base, sC, tCgC, - epi_tile, num_tiles_executed, epilogue_op, - mma_tile_coord_mnl, acc_consumer_state, acc_pipeline, c_pipeline) - - c_pipeline.producer_tail() - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def test_stage_a(): - """Test Stage A: Q @ K^T → TMEM → GMEM""" - device = torch.device("cuda") - torch.manual_seed(42) - - m, n, k = 128, 128, 512 - - # Tensors must be 3D (M, K, L) for the CUTLASS pattern - a = torch.randn(m, k, 1, dtype=torch.bfloat16, device="cuda") - b = torch.randn(n, k, 1, dtype=torch.bfloat16, device="cuda") - c = torch.zeros(m, n, 1, dtype=torch.bfloat16, device="cuda") - - ref = a[:, :, 0].float() @ b[:, :, 0].float().T - - # Create cute tensors - import cutlass.torch as cutlass_torch - mA = cutlass_torch.from_dlpack(a).mark_layout_dynamic( - leading_dim=cutlass_torch.get_leading_dim(a)) - mB = cutlass_torch.from_dlpack(b).mark_layout_dynamic( - leading_dim=cutlass_torch.get_leading_dim(b)) - mC = cutlass_torch.from_dlpack(c).mark_layout_dynamic( - leading_dim=cutlass_torch.get_leading_dim(c)) - - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - kernel = StageAQKTKernel(mma_tiler_mn=(128, 128), use_2cta_instrs=False, use_tma_store=True) - compiled = cute.compile(kernel, mA, mB, mC, stream) - - # Run with the same tensors - compiled(mA, mB, mC, stream) - torch.cuda.synchronize() - - output = c[:, :, 0].float() - cos = torch.nn.functional.cosine_similarity( - output.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - max_err = (output - ref).abs().max().item() - - print("Stage A: Q({},{}) @ K^T({}, {}) -> S({}, {})".format(m, k, k, n, m, n)) - print(" Cosine: {:.6f}, Max error: {:.6f}".format(cos, max_err)) - print(" {}".format("PASS" if cos >= 0.99 else "FAIL")) - return cos - - -if __name__ == "__main__": - test_stage_a() diff --git a/tests/archive/test_stage_a_pv_param.py b/tests/archive/test_stage_a_pv_param.py deleted file mode 100644 index a170225c..00000000 --- a/tests/archive/test_stage_a_pv_param.py +++ /dev/null @@ -1,374 +0,0 @@ -""" -Stage A: Bare Q@K^T via tcgen05.mma → TMEM → GMEM -Follows the CUTLASS dense_gemm_persistent.py pattern EXACTLY. -BF16 inputs, FP32 accumulator, TMA load/store, warp specialization. -Single tile (no persistent scheduler), cluster (1,1). -""" -import torch -import cutlass -import cutlass.cute as cute -import cutlass.utils as utils -import cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.cute.runtime import make_ptr -import cuda.bindings.driver as cuda - - -class StageAWithPVParam: - def __init__(self, mma_tiler_mn, use_2cta_instrs=False, use_tma_store=True): - self.acc_dtype = Float32 - self.use_2cta_instrs = use_2cta_instrs - self.mma_tiler_mn = mma_tiler_mn - self.mma_tiler = (*mma_tiler_mn, 1) - self.use_tma_store = use_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4 - self.tma_warp_id = 5 - self.threads_per_cta = 32 * 6 # 192 - self.epilog_sync_bar_id = 1 - self.tmem_alloc_sync_bar_id = 2 - self.tmem_dealloc_sync_bar_id = 3 - - def _create_tiled_mma(self): - return utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.a_major_mode, self.b_major_mode, - self.acc_dtype, self.cta_group, self.mma_tiler_mn, - ) - - def _setup_attributes(self): - tiled_mma = self._create_tiled_mma() - pv_mma = utils.sm100.make_trivial_tiled_mma(self.a_dtype, self.b_dtype, cute.nvgpu.OperandMajorMode.K, self.b_major_mode, self.acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2]) - mma_inst_tile_k = 4 - self.mma_tiler = (self.mma_tiler[0], self.mma_tiler[1], mma_inst_shape_k * mma_inst_tile_k) - self.cta_tile_shape_mnk = ( - self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), - self.mma_tiler[1], - self.mma_tiler[2], - ) - self.cluster_layout_vmnk = cute.tiled_divide( - cute.make_layout((1, 1, 1)), (tiled_mma.thr_id.shape,)) - self.num_mcast_ctas_a = 1 - self.num_mcast_ctas_b = 1 - self.is_a_mcast = False - self.is_b_mcast = False - - # Epilogue tile - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, self.use_2cta_instrs, self.c_layout, self.c_dtype) - - # Stage counts: 1 AB stage (single tile, no double-buffer), 1 acc stage, 2 C stages - self.num_ab_stage = 1 - self.num_acc_stage = 1 - self.num_c_stage = 2 - - # SMEM layouts - self.a_smem_layout_staged = utils.sm100.make_smem_layout_a( - tiled_mma, self.mma_tiler, self.a_dtype, self.num_ab_stage) - self.b_smem_layout_staged = utils.sm100.make_smem_layout_b( - tiled_mma, self.mma_tiler, self.b_dtype, self.num_ab_stage) - self.c_smem_layout_staged = utils.sm100.make_smem_layout_epi( - self.c_dtype, self.c_layout, self.epi_tile, self.num_c_stage) - - # TMEM alloc cols - acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) - tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(tCtAcc_fake, arch="sm_100") - - # TMA load bytes - a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) - b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.a_dtype, a_smem_layout) + - cute.size_in_bytes(self.b_dtype, b_smem_layout) - ) * cute.size(tiled_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, - stream: cuda.CUstream): - self.a_dtype = a.element_type - self.b_dtype = b.element_type - self.c_dtype = c.element_type - self.a_major_mode = LayoutEnum.from_tensor(a).mma_major_mode() - self.b_major_mode = LayoutEnum.from_tensor(b).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - tiled_mma = self._create_tiled_mma() - pv_mma = utils.sm100.make_trivial_tiled_mma(self.a_dtype, self.b_dtype, cute.nvgpu.OperandMajorMode.K, self.b_major_mode, self.acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup_attributes() - - # TMA load A - a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) - tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma.thr_id), - a, a_smem_layout, self.mma_tiler, tiled_mma, - self.cluster_layout_vmnk.shape, - ) - - # TMA load B - b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) - tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, tiled_mma.thr_id), - b, b_smem_layout, self.mma_tiler, tiled_mma, - self.cluster_layout_vmnk.shape, - ) - - # TMA store C - epi_smem_layout = cute.select(self.c_smem_layout_staged, mode=[0, 1]) - tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( - cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem_layout, self.epi_tile) - - self._kernel(pv_mma, - tiled_mma, pv_mma, tma_atom_a, tma_tensor_a, tma_atom_b, tma_tensor_b, - tma_atom_c, tma_tensor_c, self.cluster_layout_vmnk, - self.a_smem_layout_staged, self.b_smem_layout_staged, - self.c_smem_layout_staged, self.epi_tile, - ).launch(grid=(1, 1, 1), block=[self.threads_per_cta, 1, 1], stream=stream) - - @cute.kernel - def _kernel(self, tiled_mma, pv_mma, tma_atom_a, mA_mkl, tma_atom_b, mB_nkl, - tma_atom_c, mC_mnl, cluster_layout_vmnk, - a_smem_layout_staged, b_smem_layout_staged, c_smem_layout_staged, epi_tile): - warp_idx = cute.arch.warp_idx() - warp_idx = cute.arch.make_warp_uniform(warp_idx) - tidx, _, _ = cute.arch.thread_idx() - use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2 - is_leader_cta = True # single CTA, always leader - - # Prefetch TMA descriptors - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_atom_a) - cpasync.prefetch_descriptor(tma_atom_b) - cpasync.prefetch_descriptor(tma_atom_c) - - # ── Shared storage ─────────────────────────────────── - @cute.struct - class SharedStorage: - ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc_mbar: cutlass.Int64 - tmem_holding_buf: cutlass.Int32 - - smem = utils.SmemAllocator() - storage = smem.allocate(SharedStorage) - - # AB pipeline - ab_producer, ab_consumer = pipeline.PipelineTmaUmma.create( - barrier_storage=storage.ab_full_mbar_ptr.data_ptr(), - num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, - cta_layout_vmnk=cluster_layout_vmnk, - defer_sync=True, - ).make_participants() - - # ACC pipeline - acc_pipeline = pipeline.PipelineUmmaAsync.create( - barrier_storage=storage.acc_full_mbar_ptr.data_ptr(), - num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta_instrs else 1)), - cta_layout_vmnk=cluster_layout_vmnk, - defer_sync=True, - ) - - # TMEM allocator - tmem_alloc_barrier = pipeline.NamedBarrier( - barrier_id=self.tmem_alloc_sync_bar_id, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id)), - ) - tmem = utils.TmemAllocator( - storage.tmem_holding_buf.ptr, - barrier_for_retrieve=tmem_alloc_barrier, - allocator_warp_id=self.epilogue_warp_id[0], - is_two_cta=use_2cta_instrs, - two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr, - ) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cluster_layout_vmnk, is_relaxed=True) - - # SMEM tensors - sA = smem.allocate_tensor( - element_type=self.a_dtype, layout=a_smem_layout_staged.outer, - byte_alignment=128, swizzle=a_smem_layout_staged.inner) - sB = smem.allocate_tensor( - element_type=self.b_dtype, layout=b_smem_layout_staged.outer, - byte_alignment=128, swizzle=b_smem_layout_staged.inner) - sC = smem.allocate_tensor( - element_type=self.c_dtype, layout=c_smem_layout_staged.outer, - byte_alignment=128, swizzle=c_smem_layout_staged.inner) - - # Partition global tensors - gA_mkl = cute.local_tile(mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)) - gB_nkl = cute.local_tile(mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None)) - gC_mnl = cute.local_tile(mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None)) - k_tile_cnt = cute.size(gA_mkl, mode=[3]) - - # Partition for TiledMMA - thr_mma = tiled_mma.get_slice(0) # leader CTA - tCgA = thr_mma.partition_A(gA_mkl) - tCgB = thr_mma.partition_B(gB_nkl) - tCgC = thr_mma.partition_C(gC_mnl) - - # TMA partition A/B - a_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape) - tAsA, tAgA = cpasync.tma_partition( - tma_atom_a, 0, a_cta_layout, - cute.group_modes(sA, 0, 3), cute.group_modes(tCgA, 0, 3)) - b_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape) - tBsB, tBgB = cpasync.tma_partition( - tma_atom_b, 0, b_cta_layout, - cute.group_modes(sB, 0, 3), cute.group_modes(tCgB, 0, 3)) - - # Slice to tile coord (0, 0, 0) - tAgA_slice = tAgA[(None, 0, None, 0)] - tBgB_slice = tBgB[(None, 0, None, 0)] - - # MMA fragments - tCrA = tiled_mma.make_fragment_A(sA) - tCrB = tiled_mma.make_fragment_B(sB) - acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) - tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cluster_layout_vmnk) - - # ══════════════════════════════════════════════════════════ - # TMA LOAD WARP (warp 5) - # ══════════════════════════════════════════════════════════ - if warp_idx == self.tma_warp_id: - ab_producer.reset() - peek_ab_empty_status = ab_producer.try_acquire() - - for k_tile in cutlass.range(k_tile_cnt, unroll=1): - handle = ab_producer.acquire_and_advance(peek_ab_empty_status) - cute.copy(tma_atom_a, tAgA_slice[(None, handle.count)], tAsA[(None, handle.index)], - tma_bar_ptr=handle.barrier) - cute.copy(tma_atom_b, tBgB_slice[(None, handle.count)], tBsB[(None, handle.index)], - tma_bar_ptr=handle.barrier) - peek_ab_empty_status = cutlass.Boolean(1) - if handle.count + 1 < k_tile_cnt: - peek_ab_empty_status = ab_producer.try_acquire() - - ab_producer.tail() - - # ══════════════════════════════════════════════════════════ - # MMA WARP (warp 4) - # ══════════════════════════════════════════════════════════ - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) - tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) - tCtAcc = tCtAcc_base[(None, None, None, 0)] - - ab_consumer.reset() - peek_ab_full_status = cutlass.Boolean(1) - if is_leader_cta: - peek_ab_full_status = ab_consumer.try_wait() - - acc_producer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Producer, self.num_acc_stage) - if is_leader_cta: - acc_pipeline.producer_acquire(acc_producer_state) - tiled_mma.set(tcgen05.Field.ACCUMULATE, False) - - for k_tile in range(k_tile_cnt): - if is_leader_cta: - handle = ab_consumer.wait_and_advance(peek_ab_full_status) - num_kblocks = cute.size(tCrA, mode=[2]) - for kblk_idx in cutlass.range(num_kblocks, unroll_full=True): - kblk_crd = (None, None, kblk_idx, handle.index) - cute.gemm(tiled_mma, tCtAcc, tCrA[kblk_crd], tCrB[kblk_crd], tCtAcc) - tiled_mma.set(tcgen05.Field.ACCUMULATE, True) - handle.release() - peek_ab_full_status = cutlass.Boolean(1) - if handle.count + 1 < k_tile_cnt: - peek_ab_full_status = ab_consumer.try_wait() - - if is_leader_cta: - acc_pipeline.producer_commit(acc_producer_state) - acc_producer_state.advance() - acc_pipeline.producer_tail(acc_producer_state) - - # ══════════════════════════════════════════════════════════ - # EPILOGUE WARPS (0..3) - # ══════════════════════════════════════════════════════════ - if warp_idx < self.mma_warp_id: - tmem.allocate(self.num_tmem_alloc_cols) - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) - tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) - - acc_consumer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Consumer, self.num_acc_stage) - - c_producer_group = pipeline.CooperativeGroup( - pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) - c_pipeline = pipeline.PipelineTmaStore.create( - num_stages=self.num_c_stage, producer_group=c_producer_group) - - # Use the reference epilogue implementation - mma_tile_coord_mnl = (0, 0, 0) - epilogue_op = const_expr(lambda x: x) - num_tiles_executed = 0 - - acc_consumer_state = utils.gemm.sm100.epilogue_tma_store( - self, tidx, warp_idx, tma_atom_c, tCtAcc_base, sC, tCgC, - epi_tile, num_tiles_executed, epilogue_op, - mma_tile_coord_mnl, acc_consumer_state, acc_pipeline, c_pipeline) - - c_pipeline.producer_tail() - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def test_stage_a_with_pv_param(): - """Test Stage A: Q @ K^T → TMEM → GMEM""" - device = torch.device("cuda") - torch.manual_seed(42) - - m, n, k = 128, 128, 512 - - # Tensors must be 3D (M, K, L) for the CUTLASS pattern - a = torch.randn(m, k, 1, dtype=torch.bfloat16, device="cuda") - b = torch.randn(n, k, 1, dtype=torch.bfloat16, device="cuda") - c = torch.zeros(m, n, 1, dtype=torch.bfloat16, device="cuda") - - ref = a[:, :, 0].float() @ b[:, :, 0].float().T - - # Create cute tensors - import cutlass.torch as cutlass_torch - mA = cutlass_torch.from_dlpack(a).mark_layout_dynamic( - leading_dim=cutlass_torch.get_leading_dim(a)) - mB = cutlass_torch.from_dlpack(b).mark_layout_dynamic( - leading_dim=cutlass_torch.get_leading_dim(b)) - mC = cutlass_torch.from_dlpack(c).mark_layout_dynamic( - leading_dim=cutlass_torch.get_leading_dim(c)) - - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - kernel = StageAWithPVParam(mma_tiler_mn=(128, 128), use_2cta_instrs=False, use_tma_store=True) - compiled = cute.compile(kernel, mA, mB, mC, stream) - - # Run with the same tensors - compiled(mA, mB, mC, stream) - torch.cuda.synchronize() - - output = c[:, :, 0].float() - cos = torch.nn.functional.cosine_similarity( - output.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - max_err = (output - ref).abs().max().item() - - print("Stage A: Q({},{}) @ K^T({}, {}) -> S({}, {})".format(m, k, k, n, m, n)) - print(" Cosine: {:.6f}, Max error: {:.6f}".format(cos, max_err)) - print(" {}".format("PASS" if cos >= 0.99 else "FAIL")) - return cos - - -if __name__ == "__main__": - test_stage_a_with_pv_param() diff --git a/tests/archive/test_stage_a_qk.py b/tests/archive/test_stage_a_qk.py deleted file mode 100644 index 41beaf51..00000000 --- a/tests/archive/test_stage_a_qk.py +++ /dev/null @@ -1,632 +0,0 @@ -""" -Stage A: Bare Q@K^T via tcgen05.mma → TMEM → GMEM - -Validates the tcgen05 MMA path with zero attention logic. -Following dense_blockscaled_gemm_persistent.py for all Blackwell idioms. - -Shape: Q(M=128, K=512) @ K^T(512, N=16) → S(128, 16) fp32 output to GMEM. -One CTA. 6 warps: 4 epilogue, 1 MMA, 1 TMA load. - -Pipeline: TMA loads Q and K into SMEM, MMA warp issues tcgen05.mma with TMEM accumulator, -epilogue warps TMEM→reg→SMEM→TMA→GMEM. -""" - -import torch -import math - -try: - import cutlass - import cutlass.cute as cute - import cutlass.torch as cutlass_torch - import cutlass.utils as utils - import cutlass.pipeline as pipeline - import cutlass.utils.blackwell_helpers as sm100_utils - from cutlass.cute.nvgpu import tcgen05, cpasync - from cutlass import BFloat16, Float32 - from cutlass.cute.runtime import make_ptr - from typing import Tuple - import cuda.bindings.driver as cuda - HAS_CUTEDSL = True -except ImportError: - HAS_CUTEDSL = False - print("WARNING: CuTeDSL not available") - -# ── Problem dimensions ──────────────────────────────────────────────── -HG = 128 # query heads per CTA (M dimension) -KT = 64 # KV positions per tile (N dimension) — minimum for tcgen05 is 64 -HD = 512 # head dim (K dimension) - -# ── Warp specialization (mirrors the dense GEMM) ───────────────────── -EPILOGUE_WARP_IDS = (0, 1, 2, 3) -MMA_WARP_ID = 4 -TMA_WARP_ID = 5 -THREADS_PER_WARP = 32 -NUM_WARPS = 6 -NUM_THREADS = THREADS_PER_WARP * NUM_WARPS # 192 - - -class StageAQKTKernel: - """Stage A: Q @ K^T → TMEM → GMEM, no softmax, no PV GEMM. - - This is dense_blockscaled_gemm_persistent.py stripped to attention shapes: - - BF16 inputs, FP32 accumulator - - No block scaling (SFA/SFB) — plain bf16 MMA - - No persistent scheduler — one tile per CTA - - Single AB stage (one KV tile, no double-buffer needed for Stage A) - - TMA load for Q and K, tcgen05.mma, TMEM→reg→SMEM→GMEM epilogue - """ - - def __init__(self, mma_tiler_mn=(HG, KT)): - self.mma_tiler_mn = mma_tiler_mn - self.use_2cta_instrs = mma_tiler_mn[0] == 256 - self.cta_group = tcgen05.CtaGroup.TWO if self.use_2cta_instrs else tcgen05.CtaGroup.ONE - - # Warp IDs and thread count - self.epilog_warp_id = EPILOGUE_WARP_IDS - self.mma_warp_id = MMA_WARP_ID - self.tma_warp_id = TMA_WARP_ID - self.threads_per_warp = THREADS_PER_WARP - self.threads_per_cta = NUM_THREADS - - # Named barriers - self.epilog_sync_barrier = pipeline.NamedBarrier( - barrier_id=1, - num_threads=THREADS_PER_WARP * len(EPILOGUE_WARP_IDS), - ) - self.tmem_alloc_barrier = pipeline.NamedBarrier( - barrier_id=2, - num_threads=THREADS_PER_WARP * (1 + len(EPILOGUE_WARP_IDS)), - ) - - self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") - - def _setup_attributes(self, a_dtype, b_dtype, c_dtype, a_major_mode, b_major_mode, c_layout): - """Setup attributes that depend on input types.""" - self.a_dtype = a_dtype - self.b_dtype = b_dtype - self.c_dtype = c_dtype - self.acc_dtype = Float32 - self.a_major_mode = a_major_mode - self.b_major_mode = b_major_mode - self.c_layout = c_layout - - # Create tiled MMA — plain BF16 (no block scaling) - self.tiled_mma = sm100_utils.make_trivial_tiled_mma( - a_dtype, b_dtype, - a_major_mode, b_major_mode, - Float32, # acc_dtype - self.cta_group, - self.mma_tiler_mn, - ) - - atom_thr_size = cute.size(self.tiled_mma.thr_id.shape) - mma_inst_shape_k = cute.size(self.tiled_mma.shape_mnk, mode=[2]) - mma_inst_tile_k = 4 - self.mma_tiler = ( - self.mma_tiler_mn[0], - self.mma_tiler_mn[1], - mma_inst_shape_k * mma_inst_tile_k, - ) - - self.cta_tile_shape_mnk = ( - self.mma_tiler[0] // atom_thr_size, - self.mma_tiler[1], - self.mma_tiler[2], - ) - - # Cluster shape (1,1) — no clustering for attention - self.cluster_shape_mn = (1, 1) - self.cluster_layout_vmnk = cute.tiled_divide( - cute.make_layout((1, 1, 1)), - (self.tiled_mma.thr_id.shape,), - ) - self.num_mcast_ctas_a = 1 - self.num_mcast_ctas_b = 1 - self.is_a_mcast = False - self.is_b_mcast = False - - # Epilogue tile - self.epi_tile = sm100_utils.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, - self.use_2cta_instrs, - self.c_layout, - self.c_dtype, - ) - self.epi_tile_n = cute.size(self.epi_tile[1]) - - # Stage counts - self.num_ab_stage = 1 - self.num_acc_stage = 1 - self.num_c_stage = 2 - - # SMEM layouts - self.a_smem_layout_staged = sm100_utils.make_smem_layout_a( - self.tiled_mma, self.mma_tiler, a_dtype, self.num_ab_stage, - ) - self.b_smem_layout_staged = sm100_utils.make_smem_layout_b( - self.tiled_mma, self.mma_tiler, b_dtype, self.num_ab_stage, - ) - self.c_smem_layout_staged = sm100_utils.make_smem_layout_epi( - c_dtype, self.c_layout, self.epi_tile, self.num_c_stage, - ) - - # TMEM columns for accumulator - self.num_accumulator_tmem_cols = self.cta_tile_shape_mnk[1] - self.overlapping_accum = False - - @cute.jit - def __call__( - self, - a_ptr: cute.Pointer, - b_ptr: cute.Pointer, - c_ptr: cute.Pointer, - problem_m: cutlass.Int32, - problem_n: cutlass.Int32, - problem_k: cutlass.Int32, - stream: cuda.CUstream, - ): - a_dtype = a_ptr.value_type - b_dtype = b_ptr.value_type - c_dtype = c_ptr.value_type - - a_major_mode, b_major_mode, c_layout = self._layouts - self._setup_attributes(a_dtype, b_dtype, c_dtype, a_major_mode, b_major_mode, c_layout) - - m, n, k = problem_m, problem_n, problem_k - - # Make GMEM tensors — include batch dim (l=1) for local_tile compatibility - a_layout = cute.make_ordered_layout((m, cute.assume(k, 32), 1), order=(1, 0, 2)) - b_layout = cute.make_ordered_layout((n, cute.assume(k, 32), 1), order=(1, 0, 2)) - c_layout_obj = cute.make_ordered_layout((m, cute.assume(n, 32), 1), order=(1, 0, 2)) - - mA = cute.make_tensor(a_ptr, a_layout) - mB = cute.make_tensor(b_ptr, b_layout) - mC = cute.make_tensor(c_ptr, c_layout_obj) - - # TMA descriptors - a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) - b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) - - tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( - sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, self.tiled_mma.thr_id), - mA, a_smem_layout, self.mma_tiler, self.tiled_mma, - self.cluster_layout_vmnk.shape, - ) - tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( - sm100_utils.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, self.tiled_mma.thr_id), - mB, b_smem_layout, self.mma_tiler, self.tiled_mma, - self.cluster_layout_vmnk.shape, - ) - - a_copy_size = cute.size_in_bytes(self.a_dtype, a_smem_layout) - b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout) - self.num_tma_load_bytes = (a_copy_size + b_copy_size) * cute.size(self.tiled_mma.thr_id.shape) - - epi_smem_layout = cute.slice_(self.c_smem_layout_staged, (None, None, 0)) - tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( - cpasync.CopyBulkTensorTileS2GOp(), - mC, epi_smem_layout, self.epi_tile, - ) - - @cute.struct - class SharedStorage: - ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage] - ab_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage] - acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage] - acc_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage] - tmem_dealloc_mbar: cutlass.Int64 - tmem_holding_buf: cutlass.Int32 - sA: cute.struct.Align[ - cute.struct.MemRange[self.a_dtype, cute.cosize(self.a_smem_layout_staged.outer)], - 1024, - ] - sB: cute.struct.Align[ - cute.struct.MemRange[self.b_dtype, cute.cosize(self.b_smem_layout_staged.outer)], - 1024, - ] - sC: cute.struct.Align[ - cute.struct.MemRange[self.c_dtype, cute.cosize(self.c_smem_layout_staged.outer)], - 1024, - ] - - self.shared_storage = SharedStorage - - self._kernel( - self.tiled_mma, - tma_atom_a, tma_tensor_a, - tma_atom_b, tma_tensor_b, - tma_atom_c, tma_tensor_c, - self.cluster_layout_vmnk, - self.a_smem_layout_staged, - self.b_smem_layout_staged, - self.c_smem_layout_staged, - self.epi_tile, - ).launch( - grid=(1, 1, 1), - block=[self.threads_per_cta, 1, 1], - stream=stream, - min_blocks_per_mp=1, - ) - - @cute.kernel - def _kernel( - self, - tiled_mma, - tma_atom_a, mA_mkl, - tma_atom_b, mB_nkl, - tma_atom_c, mC_mnl, - cluster_layout_vmnk, - a_smem_layout_staged, - b_smem_layout_staged, - c_smem_layout_staged, - epi_tile, - ): - warp_idx = cute.arch.warp_idx() - warp_idx = cute.arch.make_warp_uniform(warp_idx) - tidx, _, _ = cute.arch.thread_idx() - - use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2 - is_leader_cta = True - - # ── Shared memory ─────────────────────────────────────── - smem = utils.SmemAllocator() - storage = smem.allocate(self.shared_storage) - - # Init AB pipeline - ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) - num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1 - ab_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, num_tma_producer) - ab_pipeline = pipeline.PipelineTmaUmma.create( - barrier_storage=storage.ab_full_mbar_ptr.data_ptr(), - num_stages=self.num_ab_stage, - producer_group=ab_pipeline_producer_group, - consumer_group=ab_pipeline_consumer_group, - tx_count=self.num_tma_load_bytes, - cta_layout_vmnk=cluster_layout_vmnk, - defer_sync=True, - ) - - # Init accumulator pipeline - acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) - num_acc_consumer_threads = self.threads_per_warp * len(self.epilog_warp_id) * (2 if use_2cta_instrs else 1) - acc_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, num_acc_consumer_threads) - acc_pipeline = pipeline.PipelineUmmaAsync.create( - barrier_storage=storage.acc_full_mbar_ptr.data_ptr(), - num_stages=self.num_acc_stage, - producer_group=acc_pipeline_producer_group, - consumer_group=acc_pipeline_consumer_group, - cta_layout_vmnk=cluster_layout_vmnk, - defer_sync=True, - ) - - # TMEM allocator - tmem = utils.TmemAllocator( - storage.tmem_holding_buf.ptr, - barrier_for_retrieve=self.tmem_alloc_barrier, - allocator_warp_id=self.epilog_warp_id[0], - is_two_cta=use_2cta_instrs, - two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr, - ) - - pipeline.pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) - - # ── SMEM tensors ──────────────────────────────────────── - sA = storage.sA.get_tensor( - a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner - ) - sB = storage.sB.get_tensor( - b_smem_layout_staged.outer, swizzle=b_smem_layout_staged.inner - ) - sC = storage.sC.get_tensor( - c_smem_layout_staged.outer, swizzle=c_smem_layout_staged.inner - ) - - # ── Partition global tensors ──────────────────────────── - gA_mkl = cute.local_tile( - mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) - ) - gB_nkl = cute.local_tile( - mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None) - ) - gC_mnl = cute.local_tile( - mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None) - ) - k_tile_cnt = cute.size(gA_mkl, mode=[3]) - - mma_tile_coord_v = 0 - thr_mma = tiled_mma.get_slice(mma_tile_coord_v) - tCgA = thr_mma.partition_A(gA_mkl) - tCgB = thr_mma.partition_B(gB_nkl) - tCgC = thr_mma.partition_C(gC_mnl) - - # TMA partition for A and B - a_cta_layout = cute.make_layout(1) - tAsA, tAgA = cpasync.tma_partition( - tma_atom_a, 0, a_cta_layout, - cute.group_modes(sA, 0, 3), - cute.group_modes(tCgA, 0, 3), - ) - b_cta_layout = cute.make_layout(1) - tBsB, tBgB = cpasync.tma_partition( - tma_atom_b, 0, b_cta_layout, - cute.group_modes(sB, 0, 3), - cute.group_modes(tCgB, 0, 3), - ) - - # Slice to the single tile coordinate - # (mma_tile_coord_mnl = (0,0,0) for single-tile) - tAgA_slice = tAgA[(None, 0, None, 0)] - tBgB_slice = tBgB[(None, 0, None, 0)] - - # MMA fragments - tCrA = tiled_mma.make_fragment_A(sA) - tCrB = tiled_mma.make_fragment_B(sB) - - # TMEM accumulator - acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) - tCtAcc_fake = tiled_mma.make_fragment_C( - cute.append(acc_shape, self.num_acc_stage) - ) - - pipeline.pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) - - # ══════════════════════════════════════════════════════════ - # TMA LOAD WARP (warp 5) - # ══════════════════════════════════════════════════════════ - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_atom_a) - cpasync.prefetch_descriptor(tma_atom_b) - cpasync.prefetch_descriptor(tma_atom_c) - - ab_producer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Producer, self.num_ab_stage - ) - - for k_tile in cutlass.range(k_tile_cnt, unroll=1): - ab_pipeline.producer_acquire(ab_producer_state) - - cute.copy( - tma_atom_a, - tAgA_slice[(None, ab_producer_state.count)], - tAsA[(None, ab_producer_state.index)], - tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state), - ) - cute.copy( - tma_atom_b, - tBgB_slice[(None, ab_producer_state.count)], - tBsB[(None, ab_producer_state.index)], - tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state), - ) - ab_producer_state.advance() - - ab_pipeline.producer_tail(ab_producer_state) - - # ══════════════════════════════════════════════════════════ - # MMA WARP (warp 4) - # ══════════════════════════════════════════════════════════ - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) - tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) - tCtAcc = tCtAcc_base[(None, None, None, 0)] - - ab_consumer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Consumer, self.num_ab_stage - ) - acc_producer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Producer, self.num_acc_stage - ) - - acc_pipeline.producer_acquire(acc_producer_state) - tiled_mma.set(tcgen05.Field.ACCUMULATE, False) - - for k_tile in range(k_tile_cnt): - if is_leader_cta: - ab_pipeline.consumer_wait(ab_consumer_state, cutlass.Boolean(1)) - - num_kblocks = cute.size(tCrA, mode=[2]) - for kblock_idx in cutlass.range(num_kblocks, unroll_full=True): - kblock_coord = (None, None, kblock_idx, ab_consumer_state.index) - cute.gemm( - tiled_mma, - tCtAcc, - tCrA[kblock_coord], - tCrB[kblock_coord], - tCtAcc, - ) - tiled_mma.set(tcgen05.Field.ACCUMULATE, True) - - if is_leader_cta: - ab_pipeline.consumer_release(ab_consumer_state) - ab_consumer_state.advance() - - if is_leader_cta: - acc_pipeline.producer_commit(acc_producer_state) - acc_producer_state.advance() - acc_pipeline.producer_tail(acc_producer_state) - - # ══════════════════════════════════════════════════════════ - # EPILOGUE WARPS (0..3) - # ══════════════════════════════════════════════════════════ - if warp_idx < self.mma_warp_id: - tmem.allocate(self.num_accumulator_tmem_cols) - tmem.wait_for_alloc() - - acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) - tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) - - epi_tidx = tidx - copy_atom_t2r = sm100_utils.get_tmem_load_op( - self.cta_tile_shape_mnk, - self.c_layout, - self.c_dtype, - self.acc_dtype, - epi_tile, - use_2cta_instrs, - ) - tAcc_epi = cute.flat_divide( - tCtAcc_base[((None, None), 0, 0, None)], - epi_tile, - ) - tiled_copy_t2r = tcgen05.make_tmem_copy( - copy_atom_t2r, tAcc_epi[(None, None, 0, 0, 0)] - ) - thr_copy_t2r = tiled_copy_t2r.get_slice(epi_tidx) - tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi) - - tTR_rAcc = cute.make_rmem_tensor( - thr_copy_t2r.partition_D( - cute.flat_divide( - tCgC[((None, None), 0, 0, None, None, None)], epi_tile - ) - )[(None, None, None, 0, 0, 0, 0, 0)].shape, - self.acc_dtype, - ) - tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, self.c_dtype) - - copy_atom_r2s = sm100_utils.get_smem_store_op( - self.c_layout, self.c_dtype, self.acc_dtype, tiled_copy_t2r - ) - tiled_copy_r2s = cute.make_tiled_copy_D(copy_atom_r2s, tiled_copy_t2r) - thr_copy_r2s = tiled_copy_r2s.get_slice(epi_tidx) - tRS_sC = thr_copy_r2s.partition_D(sC) - tRS_rC = tiled_copy_r2s.retile(tTR_rC) - - gC_epi = cute.flat_divide( - tCgC[((None, None), 0, 0, None, None, None)], epi_tile - ) - sC_for_tma = cute.group_modes(sC, 0, 2) - gC_for_tma = cute.group_modes(gC_epi, 0, 2) - bSG_sC, bSG_gC = cpasync.tma_partition( - tma_atom_c, 0, cute.make_layout(1), - sC_for_tma, gC_for_tma, - ) - - acc_consumer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Consumer, self.num_acc_stage - ) - c_producer_group = pipeline.CooperativeGroup( - pipeline.Agent.Thread, - self.threads_per_warp * len(self.epilog_warp_id), - ) - c_pipeline = pipeline.PipelineTmaStore.create( - num_stages=self.num_c_stage, - producer_group=c_producer_group, - ) - - acc_pipeline.consumer_wait(acc_consumer_state) - - tTR_tAcc_g = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) - bSG_gC_g = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC)) - - subtile_cnt = cute.size(tTR_tAcc_g.shape, mode=[3]) - for subtile_idx in cutlass.range(subtile_cnt): - tTR_tAcc_mn = tTR_tAcc_g[(None, None, None, subtile_idx)] - cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) - - acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load() - tRS_rC.store(acc_vec.to(self.c_dtype)) - - c_buffer = subtile_idx % self.num_c_stage - cute.copy(tiled_copy_r2s, tRS_rC, tRS_sC[(None, None, None, c_buffer)]) - cute.arch.fence_proxy("async.shared", space="cta") - self.epilog_sync_barrier.arrive_and_wait() - - if warp_idx == self.epilog_warp_id[0]: - cute.copy( - tma_atom_c, - bSG_sC[(None, c_buffer)], - bSG_gC_g[(None, subtile_idx)], - ) - c_pipeline.producer_commit() - c_pipeline.producer_acquire() - self.epilog_sync_barrier.arrive_and_wait() - - acc_pipeline.consumer_release(acc_consumer_state) - acc_consumer_state.advance() - - tmem.relinquish_alloc_permit() - self.epilog_sync_barrier.arrive_and_wait() - tmem.free(acc_tmem_ptr) - c_pipeline.producer_tail() - - -def test_stage_a(): - """Test Stage A: Q @ K^T via tcgen05.mma → TMEM → GMEM.""" - if not HAS_CUTEDSL: - print("CuTeDSL not available, skipping") - return - - device = torch.device("cuda") - torch.manual_seed(42) - - prob_m, prob_n, prob_k = HG, KT, HD - - tQ = torch.randn(prob_m, prob_k, dtype=torch.bfloat16, device=device) - tK = torch.randn(prob_n, prob_k, dtype=torch.bfloat16, device=device) - ref = torch.matmul(tQ.to(torch.float32), tK.to(torch.float32).T) - tC = torch.zeros(prob_m, prob_n, dtype=torch.bfloat16, device=device) - - # Compile using make_ptr pattern (like dense GEMM) - a_ptr = make_ptr(BFloat16, 0, cute.AddressSpace.gmem, assumed_align=16) - b_ptr = make_ptr(BFloat16, 0, cute.AddressSpace.gmem, assumed_align=16) - c_ptr = make_ptr(BFloat16, 0, cute.AddressSpace.gmem, assumed_align=16) - - a_major_mode = tcgen05.OperandMajorMode.K - b_major_mode = tcgen05.OperandMajorMode.K - c_layout = utils.LayoutEnum.ROW_MAJOR - - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - # Compile using make_ptr pattern (like dense GEMM) - a_ptr_fake = make_ptr(BFloat16, 0, cute.AddressSpace.gmem, assumed_align=16) - b_ptr_fake = make_ptr(BFloat16, 0, cute.AddressSpace.gmem, assumed_align=16) - c_ptr_fake = make_ptr(BFloat16, 0, cute.AddressSpace.gmem, assumed_align=16) - - a_major_mode = tcgen05.OperandMajorMode.K - b_major_mode = tcgen05.OperandMajorMode.K - c_layout = utils.LayoutEnum.ROW_MAJOR - - kernel = StageAQKTKernel() - kernel._layouts = (a_major_mode, b_major_mode, c_layout) - compiled = cute.compile( - kernel, - a_ptr_fake, b_ptr_fake, c_ptr_fake, - cutlass.Int32(prob_m), cutlass.Int32(prob_n), cutlass.Int32(prob_k), - stream, - ) - - # Create runtime pointers from torch tensor data - a_ptr = make_ptr(BFloat16, tQ.data_ptr(), cute.AddressSpace.gmem, assumed_align=16) - b_ptr = make_ptr(BFloat16, tK.data_ptr(), cute.AddressSpace.gmem, assumed_align=16) - c_ptr = make_ptr(BFloat16, tC.data_ptr(), cute.AddressSpace.gmem, assumed_align=16) - - compiled( - a_ptr, b_ptr, c_ptr, - cutlass.Int32(prob_m), cutlass.Int32(prob_n), cutlass.Int32(prob_k), - stream, - ) - torch.cuda.synchronize() - - output = tC.to(torch.float32) - cos = torch.nn.functional.cosine_similarity( - output.flatten().unsqueeze(0), ref.flatten().unsqueeze(0) - ).item() - max_err = (output - ref).abs().max().item() - mean_err = (output - ref).abs().mean().item() - - print(f"Stage A: Q({prob_m},{prob_k}) @ K^T({prob_k},{prob_n}) → S({prob_m},{prob_n})") - print(f" Cosine similarity: {cos:.6f}") - print(f" Max absolute error: {max_err:.6f}") - print(f" Mean absolute error: {mean_err:.6f}") - - if cos >= 0.99: - print(" ✅ PASS") - else: - print(" ❌ FAIL — cosine < 0.99") - - return cos - - -if __name__ == "__main__": - test_stage_a() diff --git a/tests/archive/test_stage_a_v2.py b/tests/archive/test_stage_a_v2.py deleted file mode 100644 index 07d61b09..00000000 --- a/tests/archive/test_stage_a_v2.py +++ /dev/null @@ -1,372 +0,0 @@ -""" -Stage A: Bare Q@K^T via tcgen05.mma → TMEM → GMEM -Follows the CUTLASS dense_gemm_persistent.py pattern EXACTLY. -BF16 inputs, FP32 accumulator, TMA load/store, warp specialization. -Single tile (no persistent scheduler), cluster (1,1). -""" -import torch -import cutlass -import cutlass.cute as cute -import cutlass.utils as utils -import cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.cute.runtime import make_ptr -import cuda.bindings.driver as cuda - - -class StageAQKTKernel: - def __init__(self, mma_tiler_mn, use_2cta_instrs=False, use_tma_store=True): - self.acc_dtype = Float32 - self.use_2cta_instrs = use_2cta_instrs - self.mma_tiler_mn = mma_tiler_mn - self.mma_tiler = (*mma_tiler_mn, 1) - self.use_tma_store = use_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4 - self.tma_warp_id = 5 - self.threads_per_cta = 32 * 6 # 192 - self.epilog_sync_bar_id = 1 - self.tmem_alloc_sync_bar_id = 2 - self.tmem_dealloc_sync_bar_id = 3 - - def _create_tiled_mma(self): - return utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.a_major_mode, self.b_major_mode, - self.acc_dtype, self.cta_group, self.mma_tiler_mn, - ) - - def _setup_attributes(self): - tiled_mma = self._create_tiled_mma() - mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2]) - mma_inst_tile_k = 4 - self.mma_tiler = (self.mma_tiler[0], self.mma_tiler[1], mma_inst_shape_k * mma_inst_tile_k) - self.cta_tile_shape_mnk = ( - self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), - self.mma_tiler[1], - self.mma_tiler[2], - ) - self.cluster_layout_vmnk = cute.tiled_divide( - cute.make_layout((1, 1, 1)), (tiled_mma.thr_id.shape,)) - self.num_mcast_ctas_a = 1 - self.num_mcast_ctas_b = 1 - self.is_a_mcast = False - self.is_b_mcast = False - - # Epilogue tile - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, self.use_2cta_instrs, self.c_layout, self.c_dtype) - - # Stage counts: 1 AB stage (single tile, no double-buffer), 1 acc stage, 2 C stages - self.num_ab_stage = 1 - self.num_acc_stage = 1 - self.num_c_stage = 2 - - # SMEM layouts - self.a_smem_layout_staged = utils.sm100.make_smem_layout_a( - tiled_mma, self.mma_tiler, self.a_dtype, self.num_ab_stage) - self.b_smem_layout_staged = utils.sm100.make_smem_layout_b( - tiled_mma, self.mma_tiler, self.b_dtype, self.num_ab_stage) - self.c_smem_layout_staged = utils.sm100.make_smem_layout_epi( - self.c_dtype, self.c_layout, self.epi_tile, self.num_c_stage) - - # TMEM alloc cols - acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) - tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(tCtAcc_fake, arch="sm_100") - - # TMA load bytes - a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) - b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.a_dtype, a_smem_layout) + - cute.size_in_bytes(self.b_dtype, b_smem_layout) - ) * cute.size(tiled_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, - stream: cuda.CUstream): - self.a_dtype = a.element_type - self.b_dtype = b.element_type - self.c_dtype = c.element_type - self.a_major_mode = LayoutEnum.from_tensor(a).mma_major_mode() - self.b_major_mode = LayoutEnum.from_tensor(b).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - tiled_mma = self._create_tiled_mma() - self._setup_attributes() - - # TMA load A - a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) - tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma.thr_id), - a, a_smem_layout, self.mma_tiler, tiled_mma, - self.cluster_layout_vmnk.shape, - ) - - # TMA load B - b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) - tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, tiled_mma.thr_id), - b, b_smem_layout, self.mma_tiler, tiled_mma, - self.cluster_layout_vmnk.shape, - ) - - # TMA store C - epi_smem_layout = cute.select(self.c_smem_layout_staged, mode=[0, 1]) - tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( - cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem_layout, self.epi_tile) - - self._kernel( - tiled_mma, tma_atom_a, tma_tensor_a, tma_atom_b, tma_tensor_b, - tma_atom_c, tma_tensor_c, self.cluster_layout_vmnk, - self.a_smem_layout_staged, self.b_smem_layout_staged, - self.c_smem_layout_staged, self.epi_tile, - ).launch(grid=(1, 1, 1), block=[self.threads_per_cta, 1, 1], stream=stream) - - @cute.kernel - def _kernel(self, tiled_mma, tma_atom_a, mA_mkl, tma_atom_b, mB_nkl, - tma_atom_c, mC_mnl, cluster_layout_vmnk, - a_smem_layout_staged, b_smem_layout_staged, c_smem_layout_staged, epi_tile): - warp_idx = cute.arch.warp_idx() - warp_idx = cute.arch.make_warp_uniform(warp_idx) - tidx, _, _ = cute.arch.thread_idx() - use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2 - is_leader_cta = True # single CTA, always leader - - # Prefetch TMA descriptors - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_atom_a) - cpasync.prefetch_descriptor(tma_atom_b) - cpasync.prefetch_descriptor(tma_atom_c) - - # ── Shared storage ─────────────────────────────────── - @cute.struct - class SharedStorage: - ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc_mbar: cutlass.Int64 - tmem_holding_buf: cutlass.Int32 - - smem = utils.SmemAllocator() - storage = smem.allocate(SharedStorage) - - # AB pipeline - ab_producer, ab_consumer = pipeline.PipelineTmaUmma.create( - barrier_storage=storage.ab_full_mbar_ptr.data_ptr(), - num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, - cta_layout_vmnk=cluster_layout_vmnk, - defer_sync=True, - ).make_participants() - - # ACC pipeline - acc_pipeline = pipeline.PipelineUmmaAsync.create( - barrier_storage=storage.acc_full_mbar_ptr.data_ptr(), - num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta_instrs else 1)), - cta_layout_vmnk=cluster_layout_vmnk, - defer_sync=True, - ) - - # TMEM allocator - tmem_alloc_barrier = pipeline.NamedBarrier( - barrier_id=self.tmem_alloc_sync_bar_id, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id)), - ) - tmem = utils.TmemAllocator( - storage.tmem_holding_buf.ptr, - barrier_for_retrieve=tmem_alloc_barrier, - allocator_warp_id=self.epilogue_warp_id[0], - is_two_cta=use_2cta_instrs, - two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr, - ) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cluster_layout_vmnk, is_relaxed=True) - - # SMEM tensors - sA = smem.allocate_tensor( - element_type=self.a_dtype, layout=a_smem_layout_staged.outer, - byte_alignment=128, swizzle=a_smem_layout_staged.inner) - sB = smem.allocate_tensor( - element_type=self.b_dtype, layout=b_smem_layout_staged.outer, - byte_alignment=128, swizzle=b_smem_layout_staged.inner) - sC = smem.allocate_tensor( - element_type=self.c_dtype, layout=c_smem_layout_staged.outer, - byte_alignment=128, swizzle=c_smem_layout_staged.inner) - - # Partition global tensors - gA_mkl = cute.local_tile(mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)) - gB_nkl = cute.local_tile(mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None)) - gC_mnl = cute.local_tile(mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None)) - k_tile_cnt = cute.size(gA_mkl, mode=[3]) - - # Partition for TiledMMA - thr_mma = tiled_mma.get_slice(0) # leader CTA - tCgA = thr_mma.partition_A(gA_mkl) - tCgB = thr_mma.partition_B(gB_nkl) - tCgC = thr_mma.partition_C(gC_mnl) - - # TMA partition A/B - a_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape) - tAsA, tAgA = cpasync.tma_partition( - tma_atom_a, 0, a_cta_layout, - cute.group_modes(sA, 0, 3), cute.group_modes(tCgA, 0, 3)) - b_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape) - tBsB, tBgB = cpasync.tma_partition( - tma_atom_b, 0, b_cta_layout, - cute.group_modes(sB, 0, 3), cute.group_modes(tCgB, 0, 3)) - - # Slice to tile coord (0, 0, 0) - tAgA_slice = tAgA[(None, 0, None, 0)] - tBgB_slice = tBgB[(None, 0, None, 0)] - - # MMA fragments - tCrA = tiled_mma.make_fragment_A(sA) - tCrB = tiled_mma.make_fragment_B(sB) - acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) - tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cluster_layout_vmnk) - - # ══════════════════════════════════════════════════════════ - # TMA LOAD WARP (warp 5) - # ══════════════════════════════════════════════════════════ - if warp_idx == self.tma_warp_id: - ab_producer.reset() - peek_ab_empty_status = ab_producer.try_acquire() - - for k_tile in cutlass.range(k_tile_cnt, unroll=1): - handle = ab_producer.acquire_and_advance(peek_ab_empty_status) - cute.copy(tma_atom_a, tAgA_slice[(None, handle.count)], tAsA[(None, handle.index)], - tma_bar_ptr=handle.barrier) - cute.copy(tma_atom_b, tBgB_slice[(None, handle.count)], tBsB[(None, handle.index)], - tma_bar_ptr=handle.barrier) - peek_ab_empty_status = cutlass.Boolean(1) - if handle.count + 1 < k_tile_cnt: - peek_ab_empty_status = ab_producer.try_acquire() - - ab_producer.tail() - - # ══════════════════════════════════════════════════════════ - # MMA WARP (warp 4) - # ══════════════════════════════════════════════════════════ - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) - tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) - tCtAcc = tCtAcc_base[(None, None, None, 0)] - - ab_consumer.reset() - peek_ab_full_status = cutlass.Boolean(1) - if is_leader_cta: - peek_ab_full_status = ab_consumer.try_wait() - - acc_producer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Producer, self.num_acc_stage) - if is_leader_cta: - acc_pipeline.producer_acquire(acc_producer_state) - tiled_mma.set(tcgen05.Field.ACCUMULATE, False) - - for k_tile in range(k_tile_cnt): - if is_leader_cta: - handle = ab_consumer.wait_and_advance(peek_ab_full_status) - num_kblocks = cute.size(tCrA, mode=[2]) - for kblk_idx in cutlass.range(num_kblocks, unroll_full=True): - kblk_crd = (None, None, kblk_idx, handle.index) - cute.gemm(tiled_mma, tCtAcc, tCrA[kblk_crd], tCrB[kblk_crd], tCtAcc) - tiled_mma.set(tcgen05.Field.ACCUMULATE, True) - handle.release() - peek_ab_full_status = cutlass.Boolean(1) - if handle.count + 1 < k_tile_cnt: - peek_ab_full_status = ab_consumer.try_wait() - - if is_leader_cta: - acc_pipeline.producer_commit(acc_producer_state) - acc_producer_state.advance() - acc_pipeline.producer_tail(acc_producer_state) - - # ══════════════════════════════════════════════════════════ - # EPILOGUE WARPS (0..3) - # ══════════════════════════════════════════════════════════ - if warp_idx < self.mma_warp_id: - tmem.allocate(self.num_tmem_alloc_cols) - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) - tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) - - acc_consumer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Consumer, self.num_acc_stage) - - c_producer_group = pipeline.CooperativeGroup( - pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) - c_pipeline = pipeline.PipelineTmaStore.create( - num_stages=self.num_c_stage, producer_group=c_producer_group) - - # Use the reference epilogue implementation - mma_tile_coord_mnl = (0, 0, 0) - epilogue_op = const_expr(lambda x: x) - num_tiles_executed = 0 - - acc_consumer_state = utils.gemm.sm100.epilogue_tma_store( - self, tidx, warp_idx, tma_atom_c, tCtAcc_base, sC, tCgC, - epi_tile, num_tiles_executed, epilogue_op, - mma_tile_coord_mnl, acc_consumer_state, acc_pipeline, c_pipeline) - - c_pipeline.producer_tail() - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def test_stage_a(): - """Test Stage A: Q @ K^T → TMEM → GMEM""" - device = torch.device("cuda") - torch.manual_seed(42) - - m, n, k = 128, 128, 512 - - # Tensors must be 3D (M, K, L) for the CUTLASS pattern - a = torch.randn(m, k, 1, dtype=torch.bfloat16, device="cuda") - b = torch.randn(n, k, 1, dtype=torch.bfloat16, device="cuda") - c = torch.zeros(m, n, 1, dtype=torch.bfloat16, device="cuda") - - ref = a[:, :, 0].float() @ b[:, :, 0].float().T - - # Create cute tensors - import cutlass.torch as cutlass_torch - mA = cutlass_torch.from_dlpack(a).mark_layout_dynamic( - leading_dim=cutlass_torch.get_leading_dim(a)) - mB = cutlass_torch.from_dlpack(b).mark_layout_dynamic( - leading_dim=cutlass_torch.get_leading_dim(b)) - mC = cutlass_torch.from_dlpack(c).mark_layout_dynamic( - leading_dim=cutlass_torch.get_leading_dim(c)) - - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - kernel = StageAQKTKernel(mma_tiler_mn=(128, 128), use_2cta_instrs=False, use_tma_store=True) - compiled = cute.compile(kernel, mA, mB, mC, stream) - - # Run with the same tensors - compiled(mA, mB, mC, stream) - torch.cuda.synchronize() - - output = c[:, :, 0].float() - cos = torch.nn.functional.cosine_similarity( - output.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - max_err = (output - ref).abs().max().item() - - print("Stage A: Q({},{}) @ K^T({}, {}) -> S({}, {})".format(m, k, k, n, m, n)) - print(" Cosine: {:.6f}, Max error: {:.6f}".format(cos, max_err)) - print(" {}".format("PASS" if cos >= 0.99 else "FAIL")) - return cos - - -if __name__ == "__main__": - test_stage_a() diff --git a/tests/archive/test_stage_a_with_pv_mma.py b/tests/archive/test_stage_a_with_pv_mma.py deleted file mode 100644 index 85f3c0ef..00000000 --- a/tests/archive/test_stage_a_with_pv_mma.py +++ /dev/null @@ -1,374 +0,0 @@ -""" -Stage A: Bare Q@K^T via tcgen05.mma → TMEM → GMEM -Follows the CUTLASS dense_gemm_persistent.py pattern EXACTLY. -BF16 inputs, FP32 accumulator, TMA load/store, warp specialization. -Single tile (no persistent scheduler), cluster (1,1). -""" -import torch -import cutlass -import cutlass.cute as cute -import cutlass.utils as utils -import cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.cute.runtime import make_ptr -import cuda.bindings.driver as cuda - - -class StageAQKTKernel: - def __init__(self, mma_tiler_mn, use_2cta_instrs=False, use_tma_store=True): - self.acc_dtype = Float32 - self.use_2cta_instrs = use_2cta_instrs - self.mma_tiler_mn = mma_tiler_mn - self.mma_tiler = (*mma_tiler_mn, 1) - self.use_tma_store = use_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4 - self.tma_warp_id = 5 - self.threads_per_cta = 32 * 6 # 192 - self.epilog_sync_bar_id = 1 - self.tmem_alloc_sync_bar_id = 2 - self.tmem_dealloc_sync_bar_id = 3 - - def _create_tiled_mma(self): - return utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.a_major_mode, self.b_major_mode, - self.acc_dtype, self.cta_group, self.mma_tiler_mn, - ) - - def _setup_attributes(self): - tiled_mma = self._create_tiled_mma() - pv_mma = utils.sm100.make_trivial_tiled_mma(self.a_dtype, self.b_dtype, cute.nvgpu.OperandMajorMode.K, self.b_major_mode, self.acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2]) - mma_inst_tile_k = 4 - self.mma_tiler = (self.mma_tiler[0], self.mma_tiler[1], mma_inst_shape_k * mma_inst_tile_k) - self.cta_tile_shape_mnk = ( - self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), - self.mma_tiler[1], - self.mma_tiler[2], - ) - self.cluster_layout_vmnk = cute.tiled_divide( - cute.make_layout((1, 1, 1)), (tiled_mma.thr_id.shape,)) - self.num_mcast_ctas_a = 1 - self.num_mcast_ctas_b = 1 - self.is_a_mcast = False - self.is_b_mcast = False - - # Epilogue tile - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, self.use_2cta_instrs, self.c_layout, self.c_dtype) - - # Stage counts: 1 AB stage (single tile, no double-buffer), 1 acc stage, 2 C stages - self.num_ab_stage = 1 - self.num_acc_stage = 1 - self.num_c_stage = 2 - - # SMEM layouts - self.a_smem_layout_staged = utils.sm100.make_smem_layout_a( - tiled_mma, self.mma_tiler, self.a_dtype, self.num_ab_stage) - self.b_smem_layout_staged = utils.sm100.make_smem_layout_b( - tiled_mma, self.mma_tiler, self.b_dtype, self.num_ab_stage) - self.c_smem_layout_staged = utils.sm100.make_smem_layout_epi( - self.c_dtype, self.c_layout, self.epi_tile, self.num_c_stage) - - # TMEM alloc cols - acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) - tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(tCtAcc_fake, arch="sm_100") - - # TMA load bytes - a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) - b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.a_dtype, a_smem_layout) + - cute.size_in_bytes(self.b_dtype, b_smem_layout) - ) * cute.size(tiled_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, - stream: cuda.CUstream): - self.a_dtype = a.element_type - self.b_dtype = b.element_type - self.c_dtype = c.element_type - self.a_major_mode = LayoutEnum.from_tensor(a).mma_major_mode() - self.b_major_mode = LayoutEnum.from_tensor(b).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - tiled_mma = self._create_tiled_mma() - pv_mma = utils.sm100.make_trivial_tiled_mma(self.a_dtype, self.b_dtype, cute.nvgpu.OperandMajorMode.K, self.b_major_mode, self.acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup_attributes() - - # TMA load A - a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) - tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma.thr_id), - a, a_smem_layout, self.mma_tiler, tiled_mma, - self.cluster_layout_vmnk.shape, - ) - - # TMA load B - b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) - tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, tiled_mma.thr_id), - b, b_smem_layout, self.mma_tiler, tiled_mma, - self.cluster_layout_vmnk.shape, - ) - - # TMA store C - epi_smem_layout = cute.select(self.c_smem_layout_staged, mode=[0, 1]) - tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( - cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem_layout, self.epi_tile) - - self._kernel( - tiled_mma, tma_atom_a, tma_tensor_a, tma_atom_b, tma_tensor_b, - tma_atom_c, tma_tensor_c, self.cluster_layout_vmnk, - self.a_smem_layout_staged, self.b_smem_layout_staged, - self.c_smem_layout_staged, self.epi_tile, - ).launch(grid=(1, 1, 1), block=[self.threads_per_cta, 1, 1], stream=stream) - - @cute.kernel - def _kernel(self, tiled_mma, tma_atom_a, mA_mkl, tma_atom_b, mB_nkl, - tma_atom_c, mC_mnl, cluster_layout_vmnk, - a_smem_layout_staged, b_smem_layout_staged, c_smem_layout_staged, epi_tile): - warp_idx = cute.arch.warp_idx() - warp_idx = cute.arch.make_warp_uniform(warp_idx) - tidx, _, _ = cute.arch.thread_idx() - use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2 - is_leader_cta = True # single CTA, always leader - - # Prefetch TMA descriptors - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_atom_a) - cpasync.prefetch_descriptor(tma_atom_b) - cpasync.prefetch_descriptor(tma_atom_c) - - # ── Shared storage ─────────────────────────────────── - @cute.struct - class SharedStorage: - ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc_mbar: cutlass.Int64 - tmem_holding_buf: cutlass.Int32 - - smem = utils.SmemAllocator() - storage = smem.allocate(SharedStorage) - - # AB pipeline - ab_producer, ab_consumer = pipeline.PipelineTmaUmma.create( - barrier_storage=storage.ab_full_mbar_ptr.data_ptr(), - num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, - cta_layout_vmnk=cluster_layout_vmnk, - defer_sync=True, - ).make_participants() - - # ACC pipeline - acc_pipeline = pipeline.PipelineUmmaAsync.create( - barrier_storage=storage.acc_full_mbar_ptr.data_ptr(), - num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta_instrs else 1)), - cta_layout_vmnk=cluster_layout_vmnk, - defer_sync=True, - ) - - # TMEM allocator - tmem_alloc_barrier = pipeline.NamedBarrier( - barrier_id=self.tmem_alloc_sync_bar_id, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id)), - ) - tmem = utils.TmemAllocator( - storage.tmem_holding_buf.ptr, - barrier_for_retrieve=tmem_alloc_barrier, - allocator_warp_id=self.epilogue_warp_id[0], - is_two_cta=use_2cta_instrs, - two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr, - ) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cluster_layout_vmnk, is_relaxed=True) - - # SMEM tensors - sA = smem.allocate_tensor( - element_type=self.a_dtype, layout=a_smem_layout_staged.outer, - byte_alignment=128, swizzle=a_smem_layout_staged.inner) - sB = smem.allocate_tensor( - element_type=self.b_dtype, layout=b_smem_layout_staged.outer, - byte_alignment=128, swizzle=b_smem_layout_staged.inner) - sC = smem.allocate_tensor( - element_type=self.c_dtype, layout=c_smem_layout_staged.outer, - byte_alignment=128, swizzle=c_smem_layout_staged.inner) - - # Partition global tensors - gA_mkl = cute.local_tile(mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)) - gB_nkl = cute.local_tile(mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None)) - gC_mnl = cute.local_tile(mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None)) - k_tile_cnt = cute.size(gA_mkl, mode=[3]) - - # Partition for TiledMMA - thr_mma = tiled_mma.get_slice(0) # leader CTA - tCgA = thr_mma.partition_A(gA_mkl) - tCgB = thr_mma.partition_B(gB_nkl) - tCgC = thr_mma.partition_C(gC_mnl) - - # TMA partition A/B - a_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape) - tAsA, tAgA = cpasync.tma_partition( - tma_atom_a, 0, a_cta_layout, - cute.group_modes(sA, 0, 3), cute.group_modes(tCgA, 0, 3)) - b_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape) - tBsB, tBgB = cpasync.tma_partition( - tma_atom_b, 0, b_cta_layout, - cute.group_modes(sB, 0, 3), cute.group_modes(tCgB, 0, 3)) - - # Slice to tile coord (0, 0, 0) - tAgA_slice = tAgA[(None, 0, None, 0)] - tBgB_slice = tBgB[(None, 0, None, 0)] - - # MMA fragments - tCrA = tiled_mma.make_fragment_A(sA) - tCrB = tiled_mma.make_fragment_B(sB) - acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) - tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cluster_layout_vmnk) - - # ══════════════════════════════════════════════════════════ - # TMA LOAD WARP (warp 5) - # ══════════════════════════════════════════════════════════ - if warp_idx == self.tma_warp_id: - ab_producer.reset() - peek_ab_empty_status = ab_producer.try_acquire() - - for k_tile in cutlass.range(k_tile_cnt, unroll=1): - handle = ab_producer.acquire_and_advance(peek_ab_empty_status) - cute.copy(tma_atom_a, tAgA_slice[(None, handle.count)], tAsA[(None, handle.index)], - tma_bar_ptr=handle.barrier) - cute.copy(tma_atom_b, tBgB_slice[(None, handle.count)], tBsB[(None, handle.index)], - tma_bar_ptr=handle.barrier) - peek_ab_empty_status = cutlass.Boolean(1) - if handle.count + 1 < k_tile_cnt: - peek_ab_empty_status = ab_producer.try_acquire() - - ab_producer.tail() - - # ══════════════════════════════════════════════════════════ - # MMA WARP (warp 4) - # ══════════════════════════════════════════════════════════ - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) - tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) - tCtAcc = tCtAcc_base[(None, None, None, 0)] - - ab_consumer.reset() - peek_ab_full_status = cutlass.Boolean(1) - if is_leader_cta: - peek_ab_full_status = ab_consumer.try_wait() - - acc_producer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Producer, self.num_acc_stage) - if is_leader_cta: - acc_pipeline.producer_acquire(acc_producer_state) - tiled_mma.set(tcgen05.Field.ACCUMULATE, False) - - for k_tile in range(k_tile_cnt): - if is_leader_cta: - handle = ab_consumer.wait_and_advance(peek_ab_full_status) - num_kblocks = cute.size(tCrA, mode=[2]) - for kblk_idx in cutlass.range(num_kblocks, unroll_full=True): - kblk_crd = (None, None, kblk_idx, handle.index) - cute.gemm(tiled_mma, tCtAcc, tCrA[kblk_crd], tCrB[kblk_crd], tCtAcc) - tiled_mma.set(tcgen05.Field.ACCUMULATE, True) - handle.release() - peek_ab_full_status = cutlass.Boolean(1) - if handle.count + 1 < k_tile_cnt: - peek_ab_full_status = ab_consumer.try_wait() - - if is_leader_cta: - acc_pipeline.producer_commit(acc_producer_state) - acc_producer_state.advance() - acc_pipeline.producer_tail(acc_producer_state) - - # ══════════════════════════════════════════════════════════ - # EPILOGUE WARPS (0..3) - # ══════════════════════════════════════════════════════════ - if warp_idx < self.mma_warp_id: - tmem.allocate(self.num_tmem_alloc_cols) - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) - tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) - - acc_consumer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Consumer, self.num_acc_stage) - - c_producer_group = pipeline.CooperativeGroup( - pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) - c_pipeline = pipeline.PipelineTmaStore.create( - num_stages=self.num_c_stage, producer_group=c_producer_group) - - # Use the reference epilogue implementation - mma_tile_coord_mnl = (0, 0, 0) - epilogue_op = const_expr(lambda x: x) - num_tiles_executed = 0 - - acc_consumer_state = utils.gemm.sm100.epilogue_tma_store( - self, tidx, warp_idx, tma_atom_c, tCtAcc_base, sC, tCgC, - epi_tile, num_tiles_executed, epilogue_op, - mma_tile_coord_mnl, acc_consumer_state, acc_pipeline, c_pipeline) - - c_pipeline.producer_tail() - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def test_stage_a(): - """Test Stage A: Q @ K^T → TMEM → GMEM""" - device = torch.device("cuda") - torch.manual_seed(42) - - m, n, k = 128, 128, 512 - - # Tensors must be 3D (M, K, L) for the CUTLASS pattern - a = torch.randn(m, k, 1, dtype=torch.bfloat16, device="cuda") - b = torch.randn(n, k, 1, dtype=torch.bfloat16, device="cuda") - c = torch.zeros(m, n, 1, dtype=torch.bfloat16, device="cuda") - - ref = a[:, :, 0].float() @ b[:, :, 0].float().T - - # Create cute tensors - import cutlass.torch as cutlass_torch - mA = cutlass_torch.from_dlpack(a).mark_layout_dynamic( - leading_dim=cutlass_torch.get_leading_dim(a)) - mB = cutlass_torch.from_dlpack(b).mark_layout_dynamic( - leading_dim=cutlass_torch.get_leading_dim(b)) - mC = cutlass_torch.from_dlpack(c).mark_layout_dynamic( - leading_dim=cutlass_torch.get_leading_dim(c)) - - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - kernel = StageAQKTKernel(mma_tiler_mn=(128, 128), use_2cta_instrs=False, use_tma_store=True) - compiled = cute.compile(kernel, mA, mB, mC, stream) - - # Run with the same tensors - compiled(mA, mB, mC, stream) - torch.cuda.synchronize() - - output = c[:, :, 0].float() - cos = torch.nn.functional.cosine_similarity( - output.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - max_err = (output - ref).abs().max().item() - - print("Stage A: Q({},{}) @ K^T({}, {}) -> S({}, {})".format(m, k, k, n, m, n)) - print(" Cosine: {:.6f}, Max error: {:.6f}".format(cos, max_err)) - print(" {}".format("PASS" if cos >= 0.99 else "FAIL")) - return cos - - -if __name__ == "__main__": - test_stage_a() diff --git a/tests/archive/test_stage_b_afrag.py b/tests/archive/test_stage_b_afrag.py deleted file mode 100644 index 8cce4ad8..00000000 --- a/tests/archive/test_stage_b_afrag.py +++ /dev/null @@ -1,210 +0,0 @@ -"""Stage B: Store P via A-fragment layout, not C-fragment. - -Key insight from Mike: the C-fragment store and A-fragment read use different -physical TMEM address mappings. The fix is to construct the TMEM store -using the A-fragment layout (from p_tmem_s / make_fragment_A). - -Steps: -1. Q@K^T → TMEM (C-fragment, offset 0) -2. ld scores from TMEM (C-fragment) -3. Convert FP32→BF16 -4. st P to TMEM using A-fragment layout (p_tmem_s / tOrP0) -5. PV MMA reads from TMEM using A-fragment (same layout = same physical addresses) -6. Epilogue writes output -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - -class StageBAfrag: - def __init__(self, mma_tiler_mn): - self.qk_acc_dtype = Float32; self.q_dtype = BFloat16; self.o_dtype = BFloat16 - self.c_dtype = BFloat16; self.acc_dtype = Float32 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.num_c_stage = 2; self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - pv_inst_k = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (*self.mma_tiler_mn, pv_inst_k * 4) - self.mma_tiler = self.qk_mma_tiler - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - c_layout = LayoutEnum.ROW_MAJOR; self.c_layout = c_layout - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, c_layout, self.o_dtype) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, c_layout, self.epi_tile, 2) - self.num_ab_stage = 1; self.num_acc_stage = 1 - qk_thr = qk_mma.get_slice(0); qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape); self.s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0); pv_acc_shape = pv_thr.partition_shape_C(self.mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape); self.o_cols = find_tmem_tensor_col_offset(tOtO) - # TMEM layout: S0 (scores) at offset 0, P (A-fragment) at offset 32, O at offset s_cols - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = self.s_cols - self.tmem_alloc_cols = self.s_cols + self.o_cols - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, 1)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, 1)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)); b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = (cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem)) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, LayoutEnum.from_tensor(a).mma_major_mode(), LayoutEnum.from_tensor(b).mma_major_mode(), self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, LayoutEnum.from_tensor(b).mma_major_mode(), self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)); b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - tma_a, tma_ta = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), a, a_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_b, tma_tb = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), b, b_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - self._kernel(qk_mma, pv_mma, tma_a, tma_ta, tma_b, tma_tb, tma_c, tma_tc, self.cluster_layout_vmnk, self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_a, mA, tma_b, mB, tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()); tidx, _, _ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: cpasync.prefetch_descriptor(tma_a); cpasync.prefetch_descriptor(tma_b); cpasync.prefetch_descriptor(tma_c) - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2]; mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2]; tmem_dealloc: cutlass.Int64; holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - ab_p, ab_c = pipeline.PipelineTmaUmma.create(barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True).make_participants() - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), cta_layout_vmnk=cl_vmnk, defer_sync=True).make_participants() - acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, len(self.epilogue_warp_id)), cta_layout_vmnk=cl_vmnk, defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2, num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=False, two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - sA = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sB = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV_ptr = cute.recast_ptr(sB.iterator, v_smem_s.inner); sV = cute.make_tensor(sV_ptr, v_smem_s.outer) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - gA = cute.local_tile(mA, cute.slice_(self.mma_tiler, (None,0,None)), (None,None,None)) - gB = cute.local_tile(mB, cute.slice_(self.mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gA, mode=[3]) - qk_thr = qk_mma.get_slice(0); tCgA = qk_thr.partition_A(gA); tCgB = qk_thr.partition_B(gB); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsA, tAgA = cpasync.tma_partition(tma_a, 0, a_lay, cute.group_modes(sA,0,3), cute.group_modes(tCgA,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsB, tBgB = cpasync.tma_partition(tma_b, 0, b_lay, cute.group_modes(sB,0,3), cute.group_modes(tCgB,0,3)) - tAgA = tAgA[(None,0,None,0)]; tBgB = tBgB[(None,0,None,0)] - tCrA = qk_mma.make_fragment_A(sA); tCrB = qk_mma.make_fragment_B(sB) - tCrV = pv_mma.make_fragment_B(sV) - # TMEM tensors - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - pv_thr = pv_mma.get_slice(0); pv_acc_shape = pv_thr.partition_shape_C(self.mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - # P A-fragment (matching fmha exactly) - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP = pv_thr.make_fragment_A(tP)[None, None, None, 0] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - # ── TMEM LOAD from C-fragment ── - tmem_ld = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_ld = tcgen05.make_tmem_copy(tmem_ld, tStS0) - sfw = tidx % (32 * len(self.epilogue_warp_id)) - thr_ld = tiled_ld.get_slice(sfw) - tLdS = thr_ld.partition_S(tStS0) - cS_id = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS_id) - tLdcS = thr_ld.partition_D(tScS) - # ── TMEM STORE to A-fragment layout ── - # Use St16x128bOp with BF16 for the A-fragment layout - # (St32x32bOp only works with C-fragment layout, not A-fragment) - tmem_st = cute.make_copy_atom(tcgen05.copy.St16x128bOp(tcgen05.copy.Repetition(16)), self.q_dtype) - tiled_st = tcgen05.make_tmem_copy(tmem_st, tP) - thr_st = tiled_st.get_slice(sfw) - tStP = thr_st.partition_D(tP) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, 1)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - # TMA - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek); cute.copy(tma_a, tAgA[(None,h.count)], tAsA[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_b, tBgB[(None,h.count)], tBsB[(None,h.index)], tma_bar_ptr=h.barrier); peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_stage_b_afrag2.py b/tests/archive/test_stage_b_afrag2.py deleted file mode 100644 index 74a2a39f..00000000 --- a/tests/archive/test_stage_b_afrag2.py +++ /dev/null @@ -1,217 +0,0 @@ -"""Stage B: Store P via A-fragment layout with recast C-fragment iterator. - -Matching the backward FMHA pattern exactly: -1. tOrP = pv_thr.make_fragment_A(tP)[None,None,None,0] (A-fragment layout) -2. tdVrP_iter = cute.recast_ptr(tStS.iterator, dtype=BF16) (C-fragment base, recast to BF16) -3. tdVrP = cute.make_tensor(tdVrP_iter + offset, tOrP.layout) -4. make_tmem_copy(St32x32bOp(Repetition(8)), BF16, tdVrP) -5. Store BF16 registers to tdVrP -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - -class StageBAfrag2: - def __init__(self, mma_tiler_mn): - self.qk_acc_dtype = Float32; self.q_dtype = BFloat16; self.o_dtype = BFloat16 - self.c_dtype = BFloat16; self.acc_dtype = Float32 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.num_c_stage = 2; self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - pv_inst_k = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (*self.mma_tiler_mn, pv_inst_k * 4) - self.mma_tiler = self.qk_mma_tiler - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - c_layout = LayoutEnum.ROW_MAJOR; self.c_layout = c_layout - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, c_layout, self.o_dtype) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, c_layout, self.epi_tile, 2) - self.num_ab_stage = 1; self.num_acc_stage = 1 - qk_thr = qk_mma.get_slice(0); qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape); self.s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0); pv_acc_shape = pv_thr.partition_shape_C(self.mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape); self.o_cols = find_tmem_tensor_col_offset(tOtO) - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 0 - self.tmem_o0_offset = self.s_cols * 2 - self.tmem_alloc_cols = 512 - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, 1)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, 1)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)); b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = (cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem)) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, LayoutEnum.from_tensor(a).mma_major_mode(), LayoutEnum.from_tensor(b).mma_major_mode(), self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, LayoutEnum.from_tensor(b).mma_major_mode(), self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)); b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - tma_a, tma_ta = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), a, a_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_b, tma_tb = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), b, b_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - self._kernel(qk_mma, pv_mma, tma_a, tma_ta, tma_b, tma_tb, tma_c, tma_tc, self.cluster_layout_vmnk, self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_a, mA, tma_b, mB, tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()); tidx, _, _ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: cpasync.prefetch_descriptor(tma_a); cpasync.prefetch_descriptor(tma_b); cpasync.prefetch_descriptor(tma_c) - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2]; mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2]; tmem_dealloc: cutlass.Int64; holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - ab_p, ab_c = pipeline.PipelineTmaUmma.create(barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True).make_participants() - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), cta_layout_vmnk=cl_vmnk, defer_sync=True).make_participants() - acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, len(self.epilogue_warp_id)), cta_layout_vmnk=cl_vmnk, defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2, num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=False, two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - sA = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sB = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV_ptr = cute.recast_ptr(sB.iterator, v_smem_s.inner); sV = cute.make_tensor(sV_ptr, v_smem_s.outer) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - gA = cute.local_tile(mA, cute.slice_(self.mma_tiler, (None,0,None)), (None,None,None)) - gB = cute.local_tile(mB, cute.slice_(self.mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gA, mode=[3]) - qk_thr = qk_mma.get_slice(0); tCgA = qk_thr.partition_A(gA); tCgB = qk_thr.partition_B(gB); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsA, tAgA = cpasync.tma_partition(tma_a, 0, a_lay, cute.group_modes(sA,0,3), cute.group_modes(tCgA,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsB, tBgB = cpasync.tma_partition(tma_b, 0, b_lay, cute.group_modes(sB,0,3), cute.group_modes(tCgB,0,3)) - tAgA = tAgA[(None,0,None,0)]; tBgB = tBgB[(None,0,None,0)] - tCrA = qk_mma.make_fragment_A(sA); tCrB = qk_mma.make_fragment_B(sB) - tCrV = pv_mma.make_fragment_B(sV) - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - pv_thr = pv_mma.get_slice(0); pv_acc_shape = pv_thr.partition_shape_C(self.mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - # ── P A-fragment (backward FMHA pattern) ── - # 1. Get A-fragment layout from pv_mma - tP_iter = cute.recast_ptr(tStS.iterator, dtype=self.q_dtype) - tP = cute.make_tensor(tP_iter, p_tmem_s.outer) - tOrP = pv_thr.make_fragment_A(tP)[None, None, None, 0] - # 2. Recast C-fragment iterator to BF16 (matching backward FMHA line 962) - tdVrP_iter = cute.recast_ptr(tStS.iterator, dtype=self.q_dtype) - # 3. Create store target with A-fragment layout + recast iterator - # The offset for P within TMEM: qk_acc_dtype.width / q_dtype.width * tmem_p0_offset - # But since we recast to BF16, the offset should be in BF16 units - tdVrP = cute.make_tensor( - tdVrP_iter + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - # PV MMA's A-fragment (for reading) - tOrP0 = cute.make_tensor(tOrP.iterator, tOrP.layout) - # ── TMEM LOAD from C-fragment ── - tmem_ld = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_ld = tcgen05.make_tmem_copy(tmem_ld, tStS0) - sfw = tidx % (32 * len(self.epilogue_warp_id)) - thr_ld = tiled_ld.get_slice(sfw) - tLdS = thr_ld.partition_S(tStS0) - cS_id = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS_id) - tLdcS = thr_ld.partition_D(tScS) - # ── TMEM STORE via A-fragment layout (backward FMHA pattern) ── - tmem_st = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(8)), self.q_dtype) - tiled_st = tcgen05.make_tmem_copy(tmem_st, tdVrP) - thr_st = tiled_st.get_slice(sfw) - tStP = thr_st.partition_D(tdVrP) - # Source identity for store (A-fragment shape) - cS_P = cute.make_identity_tensor((self.qk_mma_tiler[0], self.pv_mma_tiler[2])) - tScS_P = pv_thr.partition_A(cS_P) - tStcS = thr_st.partition_S(tScS_P) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, 1)) - print(f'[A2] tdVrP.layout: {tdVrP.layout}') - print(f'[A2] tOrP0.layout: {tOrP0.layout}') - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - # TMA - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek); cute.copy(tma_a, tAgA[(None,h.count)], tAsA[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_b, tBgB[(None,h.count)], tBsB[(None,h.index)], tma_bar_ptr=h.barrier); peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_stage_b_debug.py b/tests/archive/test_stage_b_debug.py deleted file mode 100644 index 37c55b61..00000000 --- a/tests/archive/test_stage_b_debug.py +++ /dev/null @@ -1,252 +0,0 @@ -"""Stage B debug: Two MMAs with PipelineUmmaAsync sync, NO softmax. -Just Q@K^T then P@V reading from same TMEM (will be garbage, but should not deadlock).""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -import cuda.bindings.driver as cuda - -class StageBDebug: - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32 - self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16 - self.o_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn - self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.use_2cta_instrs = False - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4 - self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.epilog_sync_bar_id = 1 - self.num_c_stage = 2 - self.tmem_s0_offset = 0 - self.tmem_o0_offset = 256 - self.tmem_p0_offset = 32 - self.tmem_alloc_cols = 512 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - pv_inst_k = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (*self.mma_tiler_mn, pv_inst_k * 4) - self.mma_tiler = self.qk_mma_tiler - self.cta_tile_shape_mnk = tuple(self.qk_mma_tiler) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - self.q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, self.a_dtype, 1) - self.k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, self.b_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - acc_shape = qk_mma.partition_shape_C(self.mma_tiler_mn) - tCtS_fake = qk_mma.make_fragment_C(cute.append(acc_shape, 1)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(tCtS_fake, arch="sm_100") - q_smem = cute.slice_(self.q_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.k_smem_s, (None, None, None, 0)) - self.num_tma_bytes = (cute.size_in_bytes(self.a_dtype, q_smem) + cute.size_in_bytes(self.b_dtype, k_smem)) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, a, b, c, stream): - self.a_dtype = a.element_type; self.b_dtype = b.element_type; self.c_dtype = c.element_type - self.a_major = LayoutEnum.from_tensor(a).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(b).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, self.a_major, self.b_major, self.acc_dtype, self.cta_group, self.mma_tiler_mn, - tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, cute.nvgpu.OperandMajorMode.K, self.b_major, self.acc_dtype, self.cta_group, self.mma_tiler_mn, - tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - q_smem = cute.slice_(self.q_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.k_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - a, q_smem, self.qk_mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - b, k_smem, self.qk_mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_c, tma_tc, - self.cluster_layout_vmnk, self.q_smem_s, self.k_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[192,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_c, mC, cl_vmnk, - q_smem_s, k_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator() - st = smem.allocate(SS) - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 128), - cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 128), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2, num_threads=160) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=0, is_two_cta=False, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=BFloat16, layout=q_smem_s.outer, byte_alignment=128, swizzle=q_smem_s.inner) - sK = smem.allocate_tensor(element_type=BFloat16, layout=k_smem_s.outer, byte_alignment=128, swizzle=k_smem_s.inner) - sC = smem.allocate_tensor(element_type=BFloat16, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tAsK, tAgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tAgK = tAgK[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sK) - - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler_mn) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_mma.partition_shape_C(self.mma_tiler_mn) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_mma.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, 1)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, 1)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # TMA - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tAgK[(None,h.count)], tAsK[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_stage_b_diag.py b/tests/archive/test_stage_b_diag.py deleted file mode 100644 index a2fb975d..00000000 --- a/tests/archive/test_stage_b_diag.py +++ /dev/null @@ -1,384 +0,0 @@ -""" -Stage B v7: Two MMAs + Identity Softmax with COMPUTED TMEM offsets. - -Key fixes over v6: - - TMEM offsets computed via find_tmem_tensor_col_offset (same API as get_num_tmem_alloc_cols) - - P tensor constructed from p_tmem_s.outer (matching fmha.py pattern exactly) - - tilePlikeFP32 computed from qk_mma_tiler and dtype widths - - tmem_alloc_cols from get_num_tmem_alloc_cols with all fragments - - JIT-time diagnostic prints of all TMEM sizes - -Architecture (matches fmha.py exactly): - MMA1: Q @ K^T → tmem_scores (a_source=SMEM, accumulate=False) - Identity softmax: tcgen05.ld C-layout → F32→BF16 → tcgen05.st A-layout - MMA2: P @ V → tmem_output (a_source=TMEM, accumulate=True) -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - - -class StageBIdentitySoftmax: - def __init__(self, mma_tiler_mn, use_2cta_instrs=False, use_tma_store=True): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16 - self.use_2cta_instrs = use_2cta_instrs; self.use_tma_store = use_tma_store - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.epilog_sync_bar_id = 1; self.tmem_alloc_sync_bar_id = 2; self.tmem_dealloc_sync_bar_id = 3 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - pv_inst_k = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (*self.mma_tiler_mn, pv_inst_k * 4) - self.mma_tiler = self.qk_mma_tiler - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], - self.qk_mma_tiler[2], - ) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, self.use_2cta_instrs, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.a_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.b_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.b_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - # ── COMPUTE TMEM OFFSETS USING find_tmem_tensor_col_offset ── - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - # tilePlikeFP32 for the store-side composition - self.tilePlikeFP32 = self.qk_mma_tiler[1] * self.q_dtype.width // 32 - - # ── TMEM LAYOUT (matching fmha.py) ── - # P OVERLAPS S — after softmax, P (A-layout) is written on top of scores (C-layout) - # in the same TMEM region. The A-layout view starts partway through the S region. - # fmha.py: S0=0, P0=32, O0=256 (with S1=128, P1=160 double-buffered) - # The P offset of 32 means the A-layout starts at column 32 within the S region. - # This is because the C-layout and A-layout partition TMEM differently per-thread; - # the first 32 C-layout columns are "dead space" in the A-layout mapping. - # - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 # Same as fmha.py - self.tmem_o0_offset = s_cols # 128 - self.tmem_alloc_cols = s_cols + o_cols # 256 - - # Also compute via get_num_tmem_alloc_cols for the full allocation - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, 1)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, 1)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - print(f"[StageB] s_cols (QK accumulator) = {s_cols}") - print(f"[StageB] o_cols (PV accumulator) = {o_cols}") - print(f"[StageB] tilePlikeFP32 = {self.tilePlikeFP32}") - print(f"[StageB] tmem_s0_offset = {self.tmem_s0_offset}") - print(f"[StageB] tmem_p0_offset = {self.tmem_p0_offset}") - print(f"[StageB] tmem_o0_offset = {self.tmem_o0_offset}") - print(f"[StageB] tmem_alloc_cols (computed) = {self.tmem_alloc_cols}") - print(f"[StageB] num_tmem_alloc_cols (via utils) = {self.num_tmem_alloc_cols}") - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.a_dtype, a_smem) + cute.size_in_bytes(self.b_dtype, b_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - self.a_dtype = a.element_type; self.b_dtype = b.element_type; self.c_dtype = c.element_type - self.a_major = LayoutEnum.from_tensor(a).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(b).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, cute.nvgpu.OperandMajorMode.K, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - tma_a, tma_ta = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - a, a_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_b, tma_tb = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - b, b_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_a, tma_ta, tma_b, tma_tb, tma_c, tma_tc, - self.cluster_layout_vmnk, self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_a, mA, tma_b, mB, tma_c, mC, cl_vmnk, - a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_a); cpasync.prefetch_descriptor(tma_b); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=self.tmem_alloc_sync_bar_id, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sA = smem.allocate_tensor(element_type=self.a_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sB = smem.allocate_tensor(element_type=self.b_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - # V shares the same SMEM as B (same data, different layout for PV MMA) - sV_ptr = cute.recast_ptr(sB.iterator, v_smem_s.inner) - sV = cute.make_tensor(sV_ptr, v_smem_s.outer) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gA = cute.local_tile(mA, cute.slice_(self.mma_tiler, (None,0,None)), (None,None,None)) - gB = cute.local_tile(mB, cute.slice_(self.mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gA, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - tCgA = qk_thr.partition_A(gA); tCgB = qk_thr.partition_B(gB); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsA, tAgA = cpasync.tma_partition(tma_a, 0, a_lay, cute.group_modes(sA,0,3), cute.group_modes(tCgA,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsB, tBgB = cpasync.tma_partition(tma_b, 0, b_lay, cute.group_modes(sB,0,3), cute.group_modes(tCgB,0,3)) - tAgA = tAgA[(None,0,None,0)]; tBgB = tBgB[(None,0,None,0)] - - tCrA = qk_mma.make_fragment_A(sA); tCrB = qk_mma.make_fragment_B(sB) - tCrV = pv_mma.make_fragment_B(sV) # V fragment from V SMEM layout - - # ── TMEM tensors with computed offsets (matching fmha.py pattern) ── - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - # P fragment: construct from p_tmem_s layout (matching fmha.py exactly) - # fmha.py: tP = cute.make_tensor(tStS.iterator, p_tmem_layout_staged.outer) - # tOrP = pv_thr_mma.make_fragment_A(tP)[None, None, None, 0] - # tOrP0 = cute.make_tensor(tOrP.iterator + dtype_width_ratio * tmem_p0_offset, tOrP.layout) - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ── TMA WARP ── - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_a, tAgA[(None,h.count)], tAsA[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_b, tBgB[(None,h.count)], tBsB[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - print(' Output row 0[:8]:', out[0,:8].tolist()) - print(' Output row 1[:8]:', out[1,:8].tolist()) - print(' Output row 63[:8]:', out[63,:8].tolist()) - print(' Ref row 0[:8]:', ref[0,:8].tolist()) - print(' Output col 0[:8]:', out[:8,0].tolist()) - print(' Any zeros:', (out==0).sum().item(), 'of', out.numel()) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_stage_b_final.py b/tests/archive/test_stage_b_final.py deleted file mode 100644 index 01f7ca54..00000000 --- a/tests/archive/test_stage_b_final.py +++ /dev/null @@ -1,207 +0,0 @@ -"""Stage B final: ld FP32 from S0, BF16 recast, st to S1 (offset s_cols), -PV MMA reads A-fragment from S1 at the appropriate offset. - -The recast pattern WORKS when writing to a different TMEM region. -PV MMA A-fragment reads from S1 with tmem_p0_offset adjustment. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - -class StageBFinal: - def __init__(self, mma_tiler_mn): - self.qk_acc_dtype = Float32; self.q_dtype = BFloat16; self.o_dtype = BFloat16 - self.c_dtype = BFloat16; self.acc_dtype = Float32 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.num_c_stage = 2; self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - pv_inst_k = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (*self.mma_tiler_mn, pv_inst_k * 4) - self.mma_tiler = self.qk_mma_tiler - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - c_layout = LayoutEnum.ROW_MAJOR; self.c_layout = c_layout - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, c_layout, self.o_dtype) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, c_layout, self.epi_tile, 2) - self.num_ab_stage = 1; self.num_acc_stage = 1 - qk_thr = qk_mma.get_slice(0); qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape); self.s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0); pv_acc_shape = pv_thr.partition_shape_C(self.mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape); self.o_cols = find_tmem_tensor_col_offset(tOtO) - # S0 at offset 0, P0 (BF16 view) at offset 32 within S0, O0 at offset s_cols - # But for the recast test, we write the FULL C-fragment to S1 (offset s_cols) - # PV MMA A-fragment needs to read from S1. The A-fragment offset is: - # tOrP0 starts at tStS.iterator + (F32_width/BF16_width) * tmem_p0_offset - # For S1, we need A-fragment pointing to the start of S1. - # A-fragment offset for column c: width_ratio * c - # S1 starts at s_cols (128 FP32 columns). A-fragment for P at S1 start: - # offset = (F32_width / BF16_width) * s_cols = 2 * 128 = 256 - # But tOrP layout has stride 64 for M, so base offset is in BF16 column units - # Actually, tOrP0.iterator is the base of the A-fragment in the TMEM address space. - # The A-fragment offset is (qk_acc_dtype.width / q_dtype.width) * tmem_p0_offset. - # For our case, we want P at S1 start, so tmem_p0_offset should map to s_cols. - # But tmem_p0_offset is in FP32 columns, and the A-fragment width ratio converts it. - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 0 # P starts at the beginning of S1 - self.tmem_o0_offset = self.s_cols # O region after S1 - self.tmem_alloc_cols = self.s_cols + self.o_cols # S1 + O - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, 1)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, 1)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)); b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = (cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem)) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, LayoutEnum.from_tensor(a).mma_major_mode(), LayoutEnum.from_tensor(b).mma_major_mode(), self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, LayoutEnum.from_tensor(b).mma_major_mode(), self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)); b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - tma_a, tma_ta = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), a, a_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_b, tma_tb = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), b, b_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - self._kernel(qk_mma, pv_mma, tma_a, tma_ta, tma_b, tma_tb, tma_c, tma_tc, self.cluster_layout_vmnk, self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_a, mA, tma_b, mB, tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()); tidx, _, _ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: cpasync.prefetch_descriptor(tma_a); cpasync.prefetch_descriptor(tma_b); cpasync.prefetch_descriptor(tma_c) - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2]; mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2]; tmem_dealloc: cutlass.Int64; holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - ab_p, ab_c = pipeline.PipelineTmaUmma.create(barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True).make_participants() - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), cta_layout_vmnk=cl_vmnk, defer_sync=True).make_participants() - acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, len(self.epilogue_warp_id)), cta_layout_vmnk=cl_vmnk, defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2, num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=False, two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - sA = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sB = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV_ptr = cute.recast_ptr(sB.iterator, v_smem_s.inner); sV = cute.make_tensor(sV_ptr, v_smem_s.outer) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - gA = cute.local_tile(mA, cute.slice_(self.mma_tiler, (None,0,None)), (None,None,None)) - gB = cute.local_tile(mB, cute.slice_(self.mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gA, mode=[3]) - qk_thr = qk_mma.get_slice(0); tCgA = qk_thr.partition_A(gA); tCgB = qk_thr.partition_B(gB); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsA, tAgA = cpasync.tma_partition(tma_a, 0, a_lay, cute.group_modes(sA,0,3), cute.group_modes(tCgA,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsB, tBgB = cpasync.tma_partition(tma_b, 0, b_lay, cute.group_modes(sB,0,3), cute.group_modes(tCgB,0,3)) - tAgA = tAgA[(None,0,None,0)]; tBgB = tBgB[(None,0,None,0)] - tCrA = qk_mma.make_fragment_A(sA); tCrB = qk_mma.make_fragment_B(sB) - tCrV = pv_mma.make_fragment_B(sV) - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]); tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator, tStS.layout) # S0: MMA output - tStS1 = cute.make_tensor(tStS.iterator + self.s_cols, tStS.layout) # S1: BF16 copy of scores (P) - pv_thr = pv_mma.get_slice(0); pv_acc_shape = pv_thr.partition_shape_C(self.mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - # P A-fragment pointing to S1 (offset s_cols) - tP = cute.make_tensor(tStS.iterator + self.s_cols, p_tmem_s.outer) - tOrP = pv_thr.make_fragment_A(tP)[None, None, None, 0] - # Since P is at the START of S1 (offset s_cols), no additional offset needed - tOrP0 = cute.make_tensor(tOrP.iterator, tOrP.layout) - # TMEM ld/st atoms - tmem_ld = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tmem_st = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_ld = tcgen05.make_tmem_copy(tmem_ld, tStS0); tiled_st = tcgen05.make_tmem_copy(tmem_st, tStS1) - sfw = tidx % (32 * len(self.epilogue_warp_id)); thr_ld = tiled_ld.get_slice(sfw); thr_st = tiled_st.get_slice(sfw) - tLdS = thr_ld.partition_S(tStS0); tStS1_dst = thr_st.partition_D(tStS1) - cS_id = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])); tScS = qk_thr.partition_C(cS_id) - tLdcS = thr_ld.partition_D(tScS); tStcS = thr_st.partition_S(tScS) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, 1)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - # TMA - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek); cute.copy(tma_a, tAgA[(None,h.count)], tAsA[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_b, tBgB[(None,h.count)], tBsB[(None,h.index)], tma_bar_ptr=h.barrier); peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_stage_b_identity.py b/tests/archive/test_stage_b_identity.py deleted file mode 100644 index 4c0d0a14..00000000 --- a/tests/archive/test_stage_b_identity.py +++ /dev/null @@ -1,487 +0,0 @@ -""" -Stage B: Two MMAs + Identity Softmax with Layout Transform - -Following NVIDIA's fmha.py softmax_step pattern exactly. - -Architecture: - MMA1: Q @ K^T → tmem_scores (a_source=SMEM, accumulate=False) - Identity softmax: tcgen05.ld from C-layout → convert F32→BF16 → tcgen05.st to A-layout - MMA2: P @ V → tmem_output (a_source=TMEM, accumulate=True) - -Reference: output = (Q @ K^T) @ V (no softmax, P = raw scores) - -TMEM Layout (following fmha.py for 128x128 MMA tile): - tmem_s0 = 0 (scores, QK accumulator, C-layout) - tmem_p0 = 32 (P, PV A-operand, A-layout — written by identity softmax) - tmem_o0 = 256 (output, PV accumulator, C-layout) - -The identity softmax performs the C-layout → A-layout transform via tcgen05.ld + tcgen05.st. -This is the critical bridge that makes the two-MMA pipeline work on Blackwell. -""" - -import torch -import cutlass -import cutlass.cute as cute -import cutlass.utils as utils -import cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -import cuda.bindings.driver as cuda - - -class StageBIdentitySoftmaxKernel: - def __init__(self, mma_tiler_mn, use_2cta_instrs=False, use_tma_store=True): - self.acc_dtype = Float32 - self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16 - self.o_dtype = BFloat16 - self.use_2cta_instrs = use_2cta_instrs - self.mma_tiler_mn = mma_tiler_mn - self.mma_tiler = (*mma_tiler_mn, 1) - self.use_tma_store = use_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE - - self.softmax_warp_ids = (0, 1, 2, 3) - self.epilogue_warp_id = self.softmax_warp_ids # same warps do softmax + epilogue - self.mma_warp_id = 4 - self.tma_warp_id = 5 - self.threads_per_cta = 32 * 6 - - # TMEM offsets (fmha.py pattern) - self.tmem_s0_offset = 0 - self.tmem_o0_offset = 256 - self.tmem_p0_offset = 32 - - self.tmem_alloc_cols = 512 - - self.epilog_sync_bar_id = 1 - self.tmem_alloc_sync_bar_id = 2 - self.tmem_dealloc_sync_bar_id = 3 - self.scores_full_bar_id = 4 - self.softmax_done_bar_id = 5 - - self.num_c_stage = 2 - - def _setup_attributes(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - pv_inst_k = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (*self.mma_tiler_mn, pv_inst_k * 4) - self.mma_tiler = self.qk_mma_tiler - - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0], self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.cluster_layout_vmnk = cute.tiled_divide( - cute.make_layout((1, 1, 1)), (qk_mma.thr_id.shape,)) - - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, self.use_2cta_instrs, - self.c_layout, self.o_dtype) - - self.num_ab_stage = 1 - self.num_acc_stage = 1 - - self.q_smem_layout_staged = utils.sm100.make_smem_layout_a( - qk_mma, self.qk_mma_tiler, self.a_dtype, self.num_ab_stage) - self.k_smem_layout_staged = utils.sm100.make_smem_layout_b( - qk_mma, self.qk_mma_tiler, self.b_dtype, self.num_ab_stage) - self.v_smem_layout_staged = utils.sm100.make_smem_layout_b( - pv_mma, self.pv_mma_tiler, self.b_dtype, self.num_ab_stage) - self.p_tmem_layout_staged = utils.sm100.make_smem_layout_a( - pv_mma, self.pv_mma_tiler, self.q_dtype, self.num_ab_stage) - self.c_smem_layout_staged = utils.sm100.make_smem_layout_epi( - self.o_dtype, self.c_layout, self.epi_tile, self.num_c_stage) - - # For TMEM allocation - acc_shape_qk = qk_mma.partition_shape_C(self.mma_tiler_mn) - tCtS_fake = qk_mma.make_fragment_C(cute.append(acc_shape_qk, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(tCtS_fake, arch="sm_100") - - q_smem = cute.slice_(self.q_smem_layout_staged, (None, None, None, 0)) - k_smem = cute.slice_(self.k_smem_layout_staged, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.a_dtype, q_smem) + - cute.size_in_bytes(self.b_dtype, k_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, - stream: cuda.CUstream): - self.a_dtype = a.element_type - self.b_dtype = b.element_type - self.c_dtype = c.element_type - self.a_major_mode = LayoutEnum.from_tensor(a).mma_major_mode() - self.b_major_mode = LayoutEnum.from_tensor(b).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, self.a_major_mode, self.b_major_mode, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, - tcgen05.OperandSource.SMEM) - - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, - cute.nvgpu.OperandMajorMode.K, self.b_major_mode, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, - tcgen05.OperandSource.TMEM) - - self._setup_attributes(qk_mma, pv_mma) - - q_smem = cute.slice_(self.q_smem_layout_staged, (None, None, None, 0)) - k_smem = cute.slice_(self.k_smem_layout_staged, (None, None, None, 0)) - - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - a, q_smem, self.qk_mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - b, k_smem, self.qk_mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - - epi_smem = cute.select(self.c_smem_layout_staged, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom( - cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel( - qk_mma, pv_mma, - tma_q, tma_tq, tma_k, tma_tk, - tma_c, tma_tc, - self.cluster_layout_vmnk, - self.q_smem_layout_staged, self.k_smem_layout_staged, - self.v_smem_layout_staged, self.p_tmem_layout_staged, - self.c_smem_layout_staged, self.epi_tile, - ).launch(grid=(1, 1, 1), block=[self.threads_per_cta, 1, 1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, - tma_q, mQ, tma_k, mK, - tma_c, mC, cl_vmnk, - q_smem_staged, k_smem_staged, - v_smem_staged, p_tmem_staged, - c_smem_staged, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta_instrs = cute.size(qk_mma.thr_id.shape) == 2 - is_leader_cta = True - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q) - cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SharedStorage: - ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc_mbar: cutlass.Int64 - tmem_holding_buf: cutlass.Int32 - scores_full_mbar: cutlass.Int64 - softmax_done_mbar: cutlass.Int64 - - smem = utils.SmemAllocator() - storage = smem.allocate(SharedStorage) - - ab_producer, ab_consumer = pipeline.PipelineTmaUmma.create( - barrier_storage=storage.ab_full_mbar_ptr.data_ptr(), - num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, - cta_layout_vmnk=cl_vmnk, - defer_sync=True, - ).make_participants() - - acc_pipeline = pipeline.PipelineUmmaAsync.create( - barrier_storage=storage.acc_full_mbar_ptr.data_ptr(), - num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, - 32 * len(self.softmax_warp_ids) * (2 if use_2cta_instrs else 1)), - cta_layout_vmnk=cl_vmnk, - defer_sync=True, - ) - - tmem_alloc_barrier = pipeline.NamedBarrier( - barrier_id=self.tmem_alloc_sync_bar_id, - num_threads=32 * len((self.mma_warp_id, *self.softmax_warp_ids)), - ) - tmem = utils.TmemAllocator( - storage.tmem_holding_buf.ptr, - barrier_for_retrieve=tmem_alloc_barrier, - allocator_warp_id=self.softmax_warp_ids[0], - is_two_cta=use_2cta_instrs, - two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr, - ) - - scores_full_mbar = pipeline.NamedBarrier( - barrier_id=self.scores_full_bar_id, - num_threads=32 * (1 + len(self.softmax_warp_ids)), - ) - softmax_done_mbar = pipeline.NamedBarrier( - barrier_id=self.softmax_done_bar_id, - num_threads=32 * (1 + len(self.softmax_warp_ids)), - ) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor( - element_type=self.a_dtype, layout=q_smem_staged.outer, - byte_alignment=128, swizzle=q_smem_staged.inner) - sK = smem.allocate_tensor( - element_type=self.b_dtype, layout=k_smem_staged.outer, - byte_alignment=128, swizzle=k_smem_staged.inner) - sC = smem.allocate_tensor( - element_type=self.o_dtype, layout=c_smem_staged.outer, - byte_alignment=128, swizzle=c_smem_staged.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None, 0, None)), (None, None, None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0, None, None)), (None, None, None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None, None, 0)), (None, None, None)) - k_tile_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ) - tCgK = qk_thr.partition_B(gK) - tCgC = qk_thr.partition_C(gC) - - a_cta_layout = cute.make_layout(cute.slice_(cl_vmnk, (0, 0, None, 0)).shape) - tAsQ, tAgQ = cpasync.tma_partition( - tma_q, 0, a_cta_layout, - cute.group_modes(sQ, 0, 3), cute.group_modes(tCgQ, 0, 3)) - b_cta_layout = cute.make_layout(cute.slice_(cl_vmnk, (0, None, 0, 0)).shape) - tAsK, tAgK = cpasync.tma_partition( - tma_k, 0, b_cta_layout, - cute.group_modes(sK, 0, 3), cute.group_modes(tCgK, 0, 3)) - tAgQ = tAgQ[(None, 0, None, 0)] - tAgK = tAgK[(None, 0, None, 0)] - - tCrQ = qk_mma.make_fragment_A(sQ) - tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sK) - - # ── TMEM tensor setup (following fmha.py) ── - # QK accumulator (scores) — 2D C-layout (fmha.py pattern) - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler_mn) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - # PV accumulator (output) — 2D C-layout - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_mma.partition_shape_C(self.mma_tiler_mn) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - # P fragment for PV MMA (a_source=TMEM, A-layout) - tP = cute.make_tensor(tStS.iterator, p_tmem_staged.outer) - tOrP_base = pv_mma.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout, - ) - - # Fake accumulators with stage dim (for epilogue_tma_store + TMEM allocation) - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ══════════════════════════════════════════════════════════ - # TMA LOAD WARP (warp 5) - # ══════════════════════════════════════════════════════════ - if warp_idx == self.tma_warp_id: - ab_producer.reset() - peek_ab_empty_status = ab_producer.try_acquire() - - for k_tile in cutlass.range(k_tile_cnt, unroll=1): - handle = ab_producer.acquire_and_advance(peek_ab_empty_status) - cute.copy(tma_q, tAgQ[(None, handle.count)], tAsQ[(None, handle.index)], - tma_bar_ptr=handle.barrier) - cute.copy(tma_k, tAgK[(None, handle.count)], tAsK[(None, handle.index)], - tma_bar_ptr=handle.barrier) - peek_ab_empty_status = cutlass.Boolean(1) - if handle.count + 1 < k_tile_cnt: - peek_ab_empty_status = ab_producer.try_acquire() - - ab_producer.tail() - - # ══════════════════════════════════════════════════════════ - # MMA WARP (warp 4) - # ══════════════════════════════════════════════════════════ - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - - ab_consumer.reset() - peek_ab_full_status = ab_consumer.try_wait() - - # QK MMA: Q @ K^T → tmem_scores - qk_mma.set(tcgen05.Field.ACCUMULATE, False) - for k_tile in range(k_tile_cnt): - if is_leader_cta: - handle = ab_consumer.wait_and_advance(peek_ab_full_status) - num_kblocks = cute.size(tCrQ, mode=[2]) - for kblk_idx in cutlass.range(num_kblocks, unroll_full=True): - kblk_crd = (None, None, kblk_idx, handle.index) - cute.gemm(qk_mma, tStS0, tCrQ[kblk_crd], tCrK[kblk_crd], tStS0) - handle.release() - peek_ab_full_status = cutlass.Boolean(1) - if handle.count + 1 < k_tile_cnt: - peek_ab_full_status = ab_consumer.try_wait() - - cute.arch.fence_view_async_tmem_store() - scores_full_mbar.arrive() - softmax_done_mbar.wait() - - # PV MMA: P @ V → tmem_output - pv_mma.set(tcgen05.Field.ACCUMULATE, True) - tCrV_s = tCrV[(None, None, None, 0)] - num_pv_kblocks = cute.size(tOrP0, mode=[2]) - for kblk_idx in cutlass.range(num_pv_kblocks, unroll_full=True): - cute.gemm(pv_mma, tOtO0, tOrP0[(None, None, kblk_idx)], - tCrV_s[(None, None, kblk_idx)], tOtO0) - - acc_producer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Producer, self.num_acc_stage) - acc_pipeline.producer_acquire(acc_producer_state) - acc_pipeline.producer_commit(acc_producer_state) - acc_producer_state.advance() - acc_pipeline.producer_tail(acc_producer_state) - - # ══════════════════════════════════════════════════════════ - # SOFTMAX / EPILOGUE WARPS (0..3) - # ══════════════════════════════════════════════════════════ - if warp_idx < self.mma_warp_id: - tmem.allocate(self.tmem_alloc_cols) - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) - - # ── Identity softmax: C-layout → A-layout transform ── - - # 1. LOAD pipeline (reads from QK C-layout) - tmem_load_atom = cute.make_copy_atom( - tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), - self.qk_acc_dtype, - ) - tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS0) - softmax_thread_idx = tidx % (32 * len(self.softmax_warp_ids)) - thr_tmem_load = tiled_tmem_load.get_slice(softmax_thread_idx) - tTMEM_LOADtS = thr_tmem_load.partition_S(tStS0) - cS = cute.make_identity_tensor( - (self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tTMEM_LOADcS = thr_tmem_load.partition_D(tScS) - - # 2. STORE pipeline (writes P in A-layout at tmem_p0_offset) - tilePlikeFP32 = self.qk_mma_tiler[1] // 32 * self.o_dtype.width - tStS_P_layout = cute.composition( - tStS.layout, cute.make_layout((128, tilePlikeFP32))) - tStS_P = cute.make_tensor( - tStS.iterator + self.tmem_p0_offset, tStS_P_layout) - tmem_store_atom = cute.make_copy_atom( - tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), - self.qk_acc_dtype, - ) - tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStS_P) - thr_tmem_store = tiled_tmem_store.get_slice(softmax_thread_idx) - tTMEM_STOREtS_x4 = thr_tmem_store.partition_D(tStS_P) - tScS_P_layout = cute.composition( - tScS.layout, cute.make_layout((128, tilePlikeFP32))) - tScS_P = cute.make_tensor(tScS.iterator, tScS_P_layout) - tTMEM_STOREcS = thr_tmem_store.partition_S(tScS_P) - - # 3. Wait for scores - scores_full_mbar.wait() - - # 4. Load scores from C-layout → registers - tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS) - cute.arch.fence_view_async_tmem_load() - - # 5. IDENTITY: convert F32 → Q dtype, no softmax math - tTMEM_STORErS_x4 = cute.make_rmem_tensor(tTMEM_STOREcS.shape, self.qk_acc_dtype) - tTMEM_STORErS_x4_e = cute.make_tensor( - cute.recast_ptr(tTMEM_STORErS_x4.iterator, dtype=self.q_dtype), - tTMEM_LOADrS.layout, - ) - s_vec = tTMEM_LOADrS.load() - tTMEM_STORErS_x4_e.store(s_vec.to(self.q_dtype)) - - # 6. Store into A-layout (P region) - cute.copy(tiled_tmem_store, tTMEM_STORErS_x4, tTMEM_STOREtS_x4) - cute.arch.fence_view_async_tmem_store() - - # 7. Signal MMA warp - softmax_done_mbar.arrive() - - # ── Epilogue: write output to GMEM ── - tCtO_base = cute.make_tensor( - tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout) - - acc_consumer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Consumer, self.num_acc_stage) - c_producer_group = pipeline.CooperativeGroup( - pipeline.Agent.Thread, 32 * len(self.softmax_warp_ids)) - c_pipeline = pipeline.PipelineTmaStore.create( - num_stages=self.num_c_stage, producer_group=c_producer_group) - - epilogue_op = const_expr(lambda x: x) - acc_consumer_state = utils.gemm.sm100.epilogue_tma_store( - self, tidx, warp_idx, tma_c, tCtO_base, sC, tCgC, - epi_tile, 0, epilogue_op, (0, 0, 0), - acc_consumer_state, acc_pipeline, c_pipeline) - - c_pipeline.producer_tail() - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def test_stage_b_identity_softmax(): - """Test Stage B: (Q @ K^T) @ V with identity softmax layout transform""" - torch.manual_seed(42) - - m, n, k = 128, 128, 128 - - q = torch.randn(m, k, 1, dtype=torch.bfloat16, device="cuda") - kv = torch.randn(n, k, 1, dtype=torch.bfloat16, device="cuda") - c = torch.zeros(m, n, 1, dtype=torch.bfloat16, device="cuda") - - qf = q[:, :, 0].float() - kvf = kv[:, :, 0].float() - scores = qf @ kvf.T - ref = scores @ kvf - - import cutlass.torch as cutlass_torch - mQ = cutlass_torch.from_dlpack(q).mark_layout_dynamic( - leading_dim=cutlass_torch.get_leading_dim(q)) - mK = cutlass_torch.from_dlpack(kv).mark_layout_dynamic( - leading_dim=cutlass_torch.get_leading_dim(kv)) - mC = cutlass_torch.from_dlpack(c).mark_layout_dynamic( - leading_dim=cutlass_torch.get_leading_dim(c)) - - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - kernel = StageBIdentitySoftmaxKernel( - mma_tiler_mn=(128, 128), use_2cta_instrs=False, use_tma_store=True) - print("Compiling Stage B (identity softmax with layout transform)...", flush=True) - compiled = cute.compile(kernel, mQ, mK, mC, stream) - - print("Running...", flush=True) - compiled(mQ, mK, mC, stream) - torch.cuda.synchronize() - - output = c[:, :, 0].float() - cos = torch.nn.functional.cosine_similarity( - output.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - max_err = (output - ref).abs().max().item() - - print("Stage B: (Q @ K^T) @ V with identity softmax layout transform") - print(" Shape: Q({},{}), K/V({},{}), output({},{})".format(m, k, n, k, m, n)) - print(" Cosine: {:.6f}, Max error: {:.6f}".format(cos, max_err)) - print(" {}".format("PASS" if cos >= 0.99 else "FAIL")) - return cos - - -if __name__ == "__main__": - test_stage_b_identity_softmax() diff --git a/tests/archive/test_stage_b_minimal.py b/tests/archive/test_stage_b_minimal.py deleted file mode 100644 index 9d8df640..00000000 --- a/tests/archive/test_stage_b_minimal.py +++ /dev/null @@ -1,271 +0,0 @@ -""" -Stage B Minimal: Two MMAs chained, NO softmax, NO pipeline between them. -QK MMA: Q @ K^T → tmem_scores (SMEM source) -PV MMA: P @ V → tmem_output (TMEM source, P = tmem_scores) - -This tests ONLY the PV MMA with a_source=TMEM. -If this crashes, the bug is in the TMEM A-operand path of PV MMA itself. -If this works with wrong output, the PV MMA works but the softmax pipeline is broken. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -import cuda.bindings.driver as cuda - - -class StageBMinimal: - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn - self.cta_group = tcgen05.CtaGroup.ONE - self.use_2cta_instrs = False; self.use_tma_store = True - self.epilog_sync_bar_id = 1 - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.tmem_alloc_sync_bar_id = 2 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - pv_inst_k = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (*self.mma_tiler_mn, pv_inst_k * 4) - self.mma_tiler = self.qk_mma_tiler - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.a_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.b_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - # TMEM offsets — same as fmha.py - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = 128 - - qk_acc_shape = qk_mma.get_slice(0).partition_shape_C(self.mma_tiler[:2]) - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, 1)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(tCtS_fake, arch="sm_100") - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.a_dtype, a_smem) + cute.size_in_bytes(self.b_dtype, b_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - self.a_dtype = a.element_type; self.b_dtype = b.element_type; self.c_dtype = c.element_type - self.a_major = LayoutEnum.from_tensor(a).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(b).mma_major_mode() - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, cute.nvgpu.OperandMajorMode.K, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - tma_a, tma_ta = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A((1,1), qk_mma.thr_id), - a, a_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_b, tma_tb = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B((1,1), qk_mma.thr_id), - b, b_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_a, tma_ta, tma_b, tma_tb, tma_c, tma_tc, - self.cluster_layout_vmnk, self.a_smem_s, self.b_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_a, mA, tma_b, mB, tma_c, mC, cl_vmnk, - a_smem_s, b_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_a); cpasync.prefetch_descriptor(tma_b); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, 1 * 2] # 1 AB stage - acc_bar: cute.struct.MemRange[cutlass.Int64, 1 * 2] # 1 acc stage - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - # AB pipeline (TMA load) - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - # ACC pipeline (PV output → epilogue) - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=self.tmem_alloc_sync_bar_id, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sA = smem.allocate_tensor(element_type=self.a_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sB = smem.allocate_tensor(element_type=self.b_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gA = cute.local_tile(mA, cute.slice_(self.mma_tiler, (None,0,None)), (None,None,None)) - gB = cute.local_tile(mB, cute.slice_(self.mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gA, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - tCgA = qk_thr.partition_A(gA); tCgB = qk_thr.partition_B(gB); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsA, tAgA = cpasync.tma_partition(tma_a, 0, a_lay, cute.group_modes(sA,0,3), cute.group_modes(tCgA,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsB, tBgB = cpasync.tma_partition(tma_b, 0, b_lay, cute.group_modes(sB,0,3), cute.group_modes(tCgB,0,3)) - tAgA = tAgA[(None,0,None,0)]; tBgB = tBgB[(None,0,None,0)] - - tCrA = qk_mma.make_fragment_A(sA); tCrB = qk_mma.make_fragment_B(sB) - tCrV = pv_mma.make_fragment_B(sB) # V = same as K for our test - - # TMEM tensors - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - # P fragment for PV MMA (TMEM A-operand) - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, 1)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, 1)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ── TMA WARP ── - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_a, tAgA[(None,h.count)], tAsA[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_b, tBgB[(None,h.count)], tBsB[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL (may need softmax for correct P layout)')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_stage_b_ntile_v1.py b/tests/archive/test_stage_b_ntile_v1.py deleted file mode 100644 index 9b1baa68..00000000 --- a/tests/archive/test_stage_b_ntile_v1.py +++ /dev/null @@ -1,399 +0,0 @@ -""" -Stage B with N-tiling: QK + identity softmax + PV with (128,16) PV MMA. -Single subtile test: verify the C→A TMEM bridge works for (128,16) PV. - -pv_mma_tiler = (128, 16, 128). K=128 matches QK's N=128 → C/A bridge preserved. -V is (16, 128) MN-major — first 16 rows of the head_dim. -Output is (128, 16) — should be the first 16 columns of (Q@K^T).bf16(). - -If this passes, the TMEM layout mismatch (Bug 4) is fixed by N-tiling. -Then we extend to the full (128, head_dim) output with the N-subtile loop. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class StageBSingleSubtileKernel: - """Single (128,16) PV subtile test. Output should be (Q@K^T).bf16()[:, :16].""" - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False - self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - - # PV MMA tiler: (M=128, N=16, K=128) - # N=16 = qk_inst_k (one subtile of head_dim) - # K=128 = qk_mma_tiler[1] (QK's N → PV's K, preserves C/A bridge) - self.pv_mma_tiler = ( - self.qk_mma_tiler[0], # M = 128 - qk_inst_k, # N = 16 - self.qk_mma_tiler[1], # K = 128 - ) - - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - - # For the epilogue: use PV's cta_tile_shape_mnk and epi_tile - self.cta_tile_shape_mnk = ( - self.pv_mma_tiler[0] // cute.size(pv_mma.thr_id.shape), - self.pv_mma_tiler[1], self.pv_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width # = 64 - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) - + cute.size_in_bytes(self.q_dtype, b_smem) - + cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - # PV MMA tiler: (M=128, N=16, K=128) — N=16 is one subtile, K=128 matches QK N - # This must be computed before make_trivial_tiled_mma - self.pv_mma_tiler_mn = (128, 16) # (M, N) for PV MMA - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - # PV MMA with (128, 16) output, P from TMEM - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, self.pv_mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - # Output is (128, 16) — one PV subtile - gC = cute.local_tile(mC, cute.select(self.pv_mma_tiler, mode=[0,1]), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - tCgC = pv_thr.partition_C(gC) - - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - # V: pv_mma_tiler = (128, 16, 128). V is (16, 128) MN-major. - # TMA loads V using pv_mma_tiler's B-dimension slice. - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - # QK accumulator - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - # PV accumulator: (128, 16) FP32 in TMEM - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - # P A-fragment for PV MMA - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else "FAIL"}') - return cos - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_stage_b_ntile_v3.py b/tests/archive/test_stage_b_ntile_v3.py deleted file mode 100644 index aa00fd5b..00000000 --- a/tests/archive/test_stage_b_ntile_v3.py +++ /dev/null @@ -1,357 +0,0 @@ -""" -Stage B N-tiling v3: (128,16) PV MMA. Minimal changes from v30 test_pv_diag.py. - -Key changes: -1. pv_mma_tiler = (128, 16, 128) instead of (128, 128, 128) -2. epi_tile from pv_mma_tiler[:2] = (128, 16) -3. cta_tile_shape_mnk from PV for the epilogue -4. V = first 16 rows of identity (16, 128) MN-major -5. Output (128, 16) -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class PvDiag16Kernel: - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False - self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - self.pv_mma_tiler = (self.qk_mma_tiler[0], qk_inst_k, self.qk_mma_tiler[1]) - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - - # CTA tile shape from PV for epilogue (FMHA pattern) - self.cta_tile_shape_mnk = ( - self.pv_mma_tiler[0] // cute.size(pv_mma.thr_id.shape), - self.pv_mma_tiler[1], self.pv_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - # epi_tile from PV tiler, matching FMHA: self.epi_tile = self.pv_mma_tiler[:2] - self.epi_tile = self.pv_mma_tiler[:2] - - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, (128, 16), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.select(self.pv_mma_tiler, mode=[0,1]), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - tCgC = pv_thr.partition_C(gC) - - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else "FAIL"}') - if cos < 0.5: - print(f' out max={out.max().item():.4f} min={out.min().item():.4f}') - print(f' ref max={ref.max().item():.4f} min={ref.min().item():.4f}') - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_stage_b_ntile_v5.py b/tests/archive/test_stage_b_ntile_v5.py deleted file mode 100644 index cfd410e7..00000000 --- a/tests/archive/test_stage_b_ntile_v5.py +++ /dev/null @@ -1,415 +0,0 @@ -""" -Stage B N-tiling v5: (128,16) PV MMA with proper FMHA-style output. - -TMEM → register → SMEM → TMA → GMEM. No shortcuts. Blackwell all the way. - -V = I[:16,:] → output = (Q@K^T).bf16()[:,:16] -pv_mma_tiler = (128, 16, 128). K=128 matches QK's N → C/A bridge preserved. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class PvDiag16V5Kernel: - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False - self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - self.threads_per_warp = 32 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - self.pv_mma_tiler = (self.qk_mma_tiler[0], qk_inst_k, self.qk_mma_tiler[1]) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - - # FMHA: epi_tile = pv_mma_tiler[:2] - self.epi_tile = self.pv_mma_tiler[:2] # (128, 16) - self.c_layout = LayoutEnum.ROW_MAJOR - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, self.num_c_stage) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, (128, 16), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - o_smem = cute.select(self.c_smem_s, mode=[0, 1]) - - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - # TMA store for output — FMHA pattern: cpasync.make_tiled_tma_atom - tma_o, tma_to = cpasync.make_tiled_tma_atom( - cpasync.CopyBulkTensorTileS2GOp(), c, o_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_o, tma_to, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_o, mO, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_o) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sO = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - # Output TMA partition — FMHA pattern: flat_divide + group_modes - gO_qdl = cute.flat_divide(mO, cute.select(self.pv_mma_tiler, mode=[0, 1])) - gO = gO_qdl[None, None, None, 0, 0] # single tile, coord (0,0) - o_smem_layout = cute.select(c_smem_s, mode=[0, 1]) - tOsO, tOgO = cpasync.tma_partition( - tma_o, 0, cute.make_layout(1), - cute.group_modes(sO, 0, 2), - cute.group_modes(gO, 0, 2)) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - tCgV = pv_thr.partition_B(gV) - - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else "FAIL"}') - if cos < 0.5: - print(f' out max={out.max().item():.4f} min={out.min().item():.4f}') - print(f' ref max={ref.max().item():.4f} min={ref.min().item():.4f}') - print(f' out[:4,:4]:\n{out[:4,:4]}') - print(f' ref[:4,:4]:\n{ref[:4,:4]}') - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_stage_b_ntile_v6.py b/tests/archive/test_stage_b_ntile_v6.py deleted file mode 100644 index 16e8e81e..00000000 --- a/tests/archive/test_stage_b_ntile_v6.py +++ /dev/null @@ -1,400 +0,0 @@ -""" -Stage B N-tiling v6: Run BOTH (128,128) and (128,16) PV MMAs in the same kernel. -Use the working (128,128) epilogue for output. Compare first 16 columns. - -This tests whether the (128,16) PV MMA's C→A TMEM bridge is correct -without needing a (128,16) output epilogue (which is broken). -The (128,16) result is written to a separate TMEM region, then -manually padded and copied to the (128,128) output buffer. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class DualPvKernel: - """Run QK → softmax → PV(128,128) and PV(128,16) side by side. - Output PV(128,128) result using working epilogue_tma_store. - """ - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False - self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv128_mma, pv16_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - # PV(128,128): same as v30 - self.pv128_mma_tiler = (self.qk_mma_tiler[0], self.qk_mma_tiler[1], self.qk_mma_tiler[1]) - # PV(128,16): N-tiling with N=16 - self.pv16_mma_tiler = (self.qk_mma_tiler[0], qk_inst_k, self.qk_mma_tiler[1]) - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv128_mma, self.pv128_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv128_mma, self.pv128_mma_tiler, self.q_dtype, 1) - self.p16_tmem_s = utils.sm100.make_smem_layout_a(pv16_mma, self.pv16_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - - pv128_thr = pv128_mma.get_slice(0) - pv128_acc_shape = pv128_thr.partition_shape_C(self.pv128_mma_tiler[:2]) - tOtO = pv128_thr.make_fragment_C(pv128_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - # PV(128,16) accumulator - pv16_thr = pv16_mma.get_slice(0) - pv16_acc_shape = pv16_thr.partition_shape_C(self.pv16_mma_tiler[:2]) - tOtO16 = pv16_thr.make_fragment_C(pv16_acc_shape) - o16_cols = find_tmem_tensor_col_offset(tOtO16) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols # PV(128,128) accumulator - self.tmem_o16_offset = o_cols # PV(128,16) accumulator (separate TMEM) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv128_mma.make_fragment_C(cute.append(pv128_acc_shape, self.num_acc_stage)) - tCtO16_fake = pv16_mma.make_fragment_C(cute.append(pv16_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake, tCtO16_fake], arch="sm_100") - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - pv128_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - pv16_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, (128, 16), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv128_mma, pv16_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv128_mma.thr_id), - v, v_smem, self.pv128_mma_tiler, pv128_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv128_mma, pv16_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.p16_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv128_mma, pv16_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, p16_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv128_thr = pv128_mma.get_slice(0) - pv16_thr = pv16_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv128_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv128_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv128_mma.make_fragment_B(sV) - - # QK accumulator - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - # PV(128,128) accumulator - pv128_acc_shape = pv128_thr.partition_shape_C(self.pv128_mma_tiler[:2]) - tOtO = pv128_thr.make_fragment_C(pv128_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - # PV(128,16) accumulator - pv16_acc_shape = pv16_thr.partition_shape_C(self.pv16_mma_tiler[:2]) - tOtO16 = pv16_thr.make_fragment_C(pv16_acc_shape) - tOtO16_0 = cute.make_tensor(tOtO16.iterator + self.tmem_o16_offset, tOtO16.layout) - - # P A-fragment for PV(128,128) — reads from softmax P in TMEM - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv128_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - # P A-fragment for PV(128,16) — reads from the SAME softmax P in TMEM - tP16 = cute.make_tensor(tStS.iterator, p16_tmem_s.outer) - tOrP16_base = pv16_thr.make_fragment_A(tP16) - tOrP16 = tOrP16_base[(None, None, None, 0)] - tOrP16_0 = cute.make_tensor( - tOrP16.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP16.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv128_mma.make_fragment_C(cute.append(pv128_acc_shape, self.num_acc_stage)) - tCtO16_fake = pv16_mma.make_fragment_C(cute.append(pv16_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else "FAIL"}') - # If PV(128,128) still works, the QK+softmax pipeline is fine - # The PV(128,16) result is in TMEM at tmem_o16_offset but not yet output - # We need to read it to compare - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_stage_b_ntile_v7.py b/tests/archive/test_stage_b_ntile_v7.py deleted file mode 100644 index 7d3fadb8..00000000 --- a/tests/archive/test_stage_b_ntile_v7.py +++ /dev/null @@ -1,376 +0,0 @@ -""" -Stage B N-tiling v7: PV(128,16) single N-tile test (head_dim=16). - -Based on test_pv_diag (128,128 PV) with minimal changes: -- PV MMA tiler: (128, 16, 128) instead of (128, 128, 128) -- V: (16, 128) identity MN-major instead of (128, 128) -- Output: (128, 16) instead of (128, 128) -- epi_tile uses PV's cta_tile_shape_mnk for compute_epilogue_tile_shape - -With V=I(16,128): O = P[:, :16] = (Q@K^T).bf16()[:, :16] -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class NTileV7Kernel: - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False - self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - # PV with N=16 (single N-tile, head_dim=16) - self.pv_mma_tiler = (self.qk_mma_tiler[0], self.qk_mma_tiler[1], self.qk_mma_tiler[1]) - # pv_mma_tiler = (128, 128, 128) — same as QK, but PV MMA is (128, 16) with TMEM source - # Wait, pv_mma_tiler must match the PV MMA shape, not QK's. - # For PV(128,16): pv_mma_tiler = (M_pv, N_pv, K_pv) = (128, 16, 128) - # In fmha: pv_mma_tiler = (qk_mma_tiler[0], qk_mma_tiler[2], qk_mma_tiler[1]) - # = (128, inst_k*4, 128) — but we want N=16, not inst_k*4. - # The pv_mma_tiler is used for V SMEM layout and TMA, NOT for the MMA atom. - # The actual MMA atom is (128, 16) set in make_trivial_tiled_mma. - # pv_mma_tiler must match: (128, 16, 128) for the B-operand to tile V as (16, 128). - self.pv_mma_tiler = (self.qk_mma_tiler[0], 16, self.qk_mma_tiler[1]) - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - # epi_tile must match PV's output shape (128, 16) - self.pv_cta_tile_shape_mnk = ( - self.pv_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.pv_mma_tiler[1], self.pv_mma_tiler[2]) - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.pv_cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - # PV with (128, 16) output, P from TMEM - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, (128, 16), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - if cos < 0.99: - print(' out[0,:8] =', out[0,:8].tolist()) - print(' ref[0,:8] =', ref_out[0,:8].tolist()) - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_stage_b_ntile_v8.py b/tests/archive/test_stage_b_ntile_v8.py deleted file mode 100644 index 5e0693dd..00000000 --- a/tests/archive/test_stage_b_ntile_v8.py +++ /dev/null @@ -1,387 +0,0 @@ -""" -Minimal PV-only test: Load P from GMEM to TMEM via QK-style MMA, then PV from TMEM. -Step 1: QK MMA writes FP32 S to TMEM (we know this works) -Step 2: Softmax packing writes BF16 P to TMEM (test this) -Step 3: PV MMA reads BF16 P from TMEM and V from SMEM, produces O - -But to isolate the bug, let me test just the PV MMA in isolation. -I'll write known BF16 values to TMEM using the softmax packing path, -then immediately read them back using the PV A-fragment path, -and compare. - -Actually, the simplest isolation test: -1. Do QK MMA to get S in TMEM (cosine 0.999999 verified) -2. Do softmax packing: S → P in TMEM (at offset 32) -3. Skip PV entirely — read P from TMEM using the C-fragment composition LOAD path -4. Output P to GMEM and compare against S.to(BF16) - -This tests whether the softmax packing writes P correctly to the same TMEM -that the PV would read from. - -But we can't easily read P from TMEM using the standard epilogue path -because the epilogue expects FP32 accumulator data. - -Alternative: Use the PV MMA with V=I (identity). If P is correct, -then P @ I = P. But V needs to be MN-major and (128, 128), not (128, 64). -The output would be (128, 128) which doesn't match our (128, 64) c tensor. - -Let me use V that selects the first 64 columns: V[k, n] = delta(k, n) for k in [0,63]. -This gives P @ V = P[:, :64], and the output is (128, 64). -But V is (128, 128) in the MMA K,N dims. V[k, n] for k in [0,127], n in [0,63]. -Hmm, this is getting complicated. Let me just do the identity approach with a (128, 128) output. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class NTileV8Kernel: - """QK + softmax packing + PV with V=I to isolate PV MMA correctness. - Output should be P = S.to(BF16), i.e. (Q@K^T).bfloat16() - With V=I, O = P @ I = P. - But V is (K=128, N=128) in the MMA. We need a 128x128 identity in MN-major. - Output tensor is (128, 128). - """ - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False # needed by epilogue_tma_store - self.epilog_sync_bar_id = 1 # needed by epilogue_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = int(cute.size(qk_mma.shape_mnk, mode=[2])) - self.qk_mma_tiler = (*self.mma_tiler_mn, int(qk_inst_k * 4)) - # PV with V=I: output is (128, 128), same as QK - self.pv_mma_tiler = (self.qk_mma_tiler[0], qk_inst_k, self.qk_mma_tiler[1]) - # pv_mma_tiler = (128, 16, 128) FMHA: (qk_M, qk_K_inst, qk_N) - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - pv_cta_tile = (self.pv_mma_tiler[0], self.pv_mma_tiler[1], self.pv_mma_tiler[2]) - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - pv_cta_tile, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - # ⛔⛔⛔ CRITICAL: num_tma_load_bytes MUST include ALL TMA-loaded tensors (Q + K + V). Missing V → DEADLOCK. See FOOTGUN #0 in README. - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - # PV with 128x128 output (V=I) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, (128, 16), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.pv_mma_tiler, (None,0,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1 O = P[:,:16] = (Q@K^T).bf16()[:,:16] - ref = (qf @ kf.T).bfloat16().float()[:, :16] - - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - kernel = NTileV8Kernel(mma_tiler_mn=(128, 128)) - print('Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - print('Running...', flush=True) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - out = c[:,:,0].float() - cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - print('PV(128,16) v8: cosine {:.6f} {}'.format(cos, 'PASS' if cos >= 0.99 else 'FAIL')) - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_stage_b_pipeline_only.py b/tests/archive/test_stage_b_pipeline_only.py deleted file mode 100644 index 63462f42..00000000 --- a/tests/archive/test_stage_b_pipeline_only.py +++ /dev/null @@ -1,281 +0,0 @@ -""" -Stage B Pipeline-Only: Two MMAs with PipelineUmmaAsync between them, -but IDENTITY transform (no tcgen05.ld/st — P is just tmem_scores). -Tests whether the PipelineUmmaAsync barrier ordering causes the crash. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -import cuda.bindings.driver as cuda - - -class StageBPipelineOnly: - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn - self.cta_group = tcgen05.CtaGroup.ONE - self.use_2cta_instrs = False; self.use_tma_store = True - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.epilog_sync_bar_id = 1; self.tmem_alloc_sync_bar_id = 2 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - pv_inst_k = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (*self.mma_tiler_mn, pv_inst_k * 4) - self.mma_tiler = self.qk_mma_tiler - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.a_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.b_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = 128 - - qk_acc_shape = qk_mma.get_slice(0).partition_shape_C(self.mma_tiler[:2]) - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, 1)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(tCtS_fake, arch="sm_100") - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.a_dtype, a_smem) + cute.size_in_bytes(self.b_dtype, b_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - self.a_dtype = a.element_type; self.b_dtype = b.element_type; self.c_dtype = c.element_type - self.a_major = LayoutEnum.from_tensor(a).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(b).mma_major_mode() - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, cute.nvgpu.OperandMajorMode.K, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - tma_a, tma_ta = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A((1,1), qk_mma.thr_id), - a, a_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_b, tma_tb = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B((1,1), qk_mma.thr_id), - b, b_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_a, tma_ta, tma_b, tma_tb, tma_c, tma_tc, - self.cluster_layout_vmnk, self.a_smem_s, self.b_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_a, mA, tma_b, mB, tma_c, mC, cl_vmnk, - a_smem_s, b_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_a); cpasync.prefetch_descriptor(tma_b); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, 2] # 1 stage - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] # 1 stage MMA↔softmax - acc_bar: cute.struct.MemRange[cutlass.Int64, 2] # 1 stage - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - # MMA↔softmax pipeline - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=self.tmem_alloc_sync_bar_id, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sA = smem.allocate_tensor(element_type=self.a_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sB = smem.allocate_tensor(element_type=self.b_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gA = cute.local_tile(mA, cute.slice_(self.mma_tiler, (None,0,None)), (None,None,None)) - gB = cute.local_tile(mB, cute.slice_(self.mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gA, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - tCgA = qk_thr.partition_A(gA); tCgB = qk_thr.partition_B(gB); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsA, tAgA = cpasync.tma_partition(tma_a, 0, a_lay, cute.group_modes(sA,0,3), cute.group_modes(tCgA,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsB, tBgB = cpasync.tma_partition(tma_b, 0, b_lay, cute.group_modes(sB,0,3), cute.group_modes(tCgB,0,3)) - tAgA = tAgA[(None,0,None,0)]; tBgB = tBgB[(None,0,None,0)] - - tCrA = qk_mma.make_fragment_A(sA); tCrB = qk_mma.make_fragment_B(sB) - tCrV = pv_mma.make_fragment_B(sB) - - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - # P fragment for PV MMA (TMEM A-operand) - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, 1)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, 1)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ── TMA WARP ── - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_a, tAgA[(None,h.count)], tAsA[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_b, tBgB[(None,h.count)], tBsB[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1 tmem_scores (a_source=SMEM, accumulate=False) - Identity softmax: tcgen05.ld -> fill 1.0 -> tcgen05.st back to TMEM - MMA2: P @ V -> tmem_output (a_source=TMEM, accumulate=True) - Two barriers: scores_full (MMA->epi), softmax_done (epi->MMA) - Two TMEM regions: scores (offset 0), output (offset N) - -With identity softmax (P = 1.0), output = ones(M,N) @ V = column sums of V tiled. -""" -import torch -import cutlass -import cutlass.cute as cute -import cutlass.utils as utils -import cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -import cuda.bindings.driver as cuda - - -class StageBKernel: - def __init__(self, mma_tiler_mn, use_2cta_instrs=False, use_tma_store=True): - self.acc_dtype = Float32 - self.use_2cta_instrs = use_2cta_instrs - self.mma_tiler_mn = mma_tiler_mn - self.mma_tiler = (*mma_tiler_mn, 1) - self.use_tma_store = use_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4 - self.tma_warp_id = 5 - self.threads_per_cta = 32 * 6 - self.epilog_sync_bar_id = 1 - self.tmem_alloc_sync_bar_id = 2 - self.tmem_dealloc_sync_bar_id = 3 - self.scores_full_bar_id = 5 - self.softmax_done_bar_id = 6 - - def _setup_attributes(self, tiled_mma1, tiled_mma2): - mma_inst_shape_k = cute.size(tiled_mma1.shape_mnk, mode=[2]) - mma_inst_tile_k = 4 - self.mma_tiler = (self.mma_tiler[0], self.mma_tiler[1], - mma_inst_shape_k * mma_inst_tile_k) - self.cta_tile_shape_mnk = ( - self.mma_tiler[0] // cute.size(tiled_mma1.thr_id.shape), - self.mma_tiler[1], - self.mma_tiler[2], - ) - self.cluster_layout_vmnk = cute.tiled_divide( - cute.make_layout((1, 1, 1)), (tiled_mma1.thr_id.shape,)) - - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, self.use_2cta_instrs, self.c_layout, self.c_dtype) - - self.num_ab_stage = 1 - self.num_acc_stage = 1 - self.num_c_stage = 2 - - self.a_smem_layout_staged = utils.sm100.make_smem_layout_a( - tiled_mma1, self.mma_tiler, self.a_dtype, self.num_ab_stage) - self.b_smem_layout_staged = utils.sm100.make_smem_layout_b( - tiled_mma1, self.mma_tiler, self.b_dtype, self.num_ab_stage) - self.c_smem_layout_staged = utils.sm100.make_smem_layout_epi( - self.c_dtype, self.c_layout, self.epi_tile, self.num_c_stage) - - acc_shape = tiled_mma1.partition_shape_C(self.mma_tiler_mn) - tCtAcc_fake = tiled_mma1.make_fragment_C(cute.append(acc_shape, self.num_acc_stage)) - self.num_tmem_cols_per_region = utils.get_num_tmem_alloc_cols(tCtAcc_fake, arch="sm_100") - total = self.num_tmem_cols_per_region * 2 - self.total_tmem_cols = 256 - if total > 256: - self.total_tmem_cols = 512 - - a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) - b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.a_dtype, a_smem_layout) + - cute.size_in_bytes(self.b_dtype, b_smem_layout) - ) * cute.size(tiled_mma1.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, - stream: cuda.CUstream): - self.a_dtype = a.element_type - self.b_dtype = b.element_type - self.c_dtype = c.element_type - self.a_major_mode = LayoutEnum.from_tensor(a).mma_major_mode() - self.b_major_mode = LayoutEnum.from_tensor(b).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - tiled_mma1 = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, self.a_major_mode, self.b_major_mode, - self.acc_dtype, self.cta_group, self.mma_tiler_mn, - tcgen05.OperandSource.SMEM, - ) - tiled_mma2 = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, self.a_major_mode, self.b_major_mode, - self.acc_dtype, self.cta_group, self.mma_tiler_mn, - tcgen05.OperandSource.TMEM, - ) - - self._setup_attributes(tiled_mma1, tiled_mma2) - - a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) - b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) - - tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A( - self.cluster_shape_mn, tiled_mma1.thr_id), - a, a_smem_layout, self.mma_tiler, tiled_mma1, - self.cluster_layout_vmnk.shape, - ) - tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B( - self.cluster_shape_mn, tiled_mma1.thr_id), - b, b_smem_layout, self.mma_tiler, tiled_mma1, - self.cluster_layout_vmnk.shape, - ) - - epi_smem_layout = cute.select(self.c_smem_layout_staged, mode=[0, 1]) - tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( - cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem_layout, self.epi_tile) - - self._kernel( - tiled_mma1, tiled_mma2, - tma_atom_a, tma_tensor_a, tma_atom_b, tma_tensor_b, - tma_atom_c, tma_tensor_c, self.cluster_layout_vmnk, - self.a_smem_layout_staged, self.b_smem_layout_staged, - self.c_smem_layout_staged, self.epi_tile, - ).launch(grid=(1, 1, 1), block=[self.threads_per_cta, 1, 1], stream=stream) - - @cute.kernel - def _kernel(self, tiled_mma1, tiled_mma2, - tma_atom_a, mA_mkl, tma_atom_b, mB_nkl, - tma_atom_c, mC_mnl, cluster_layout_vmnk, - a_smem_layout_staged, b_smem_layout_staged, - c_smem_layout_staged, epi_tile): - warp_idx = cute.arch.warp_idx() - warp_idx = cute.arch.make_warp_uniform(warp_idx) - tidx, _, _ = cute.arch.thread_idx() - use_2cta_instrs = cute.size(tiled_mma1.thr_id.shape) == 2 - is_leader_cta = True - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_atom_a) - cpasync.prefetch_descriptor(tma_atom_b) - cpasync.prefetch_descriptor(tma_atom_c) - - @cute.struct - class SharedStorage: - ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc_mbar: cutlass.Int64 - tmem_holding_buf: cutlass.Int32 - - smem = utils.SmemAllocator() - storage = smem.allocate(SharedStorage) - - ab_producer, ab_consumer = pipeline.PipelineTmaUmma.create( - barrier_storage=storage.ab_full_mbar_ptr.data_ptr(), - num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, - cta_layout_vmnk=cluster_layout_vmnk, - defer_sync=True, - ).make_participants() - - acc_pipeline = pipeline.PipelineUmmaAsync.create( - barrier_storage=storage.acc_full_mbar_ptr.data_ptr(), - num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta_instrs else 1)), - cta_layout_vmnk=cluster_layout_vmnk, - defer_sync=True, - ) - - tmem_alloc_barrier = pipeline.NamedBarrier( - barrier_id=self.tmem_alloc_sync_bar_id, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id)), - ) - tmem = utils.TmemAllocator( - storage.tmem_holding_buf.ptr, - barrier_for_retrieve=tmem_alloc_barrier, - allocator_warp_id=self.epilogue_warp_id[0], - is_two_cta=use_2cta_instrs, - two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr, - ) - - # MMA <-> softmax handshake barriers - scores_full_barrier = pipeline.NamedBarrier( - barrier_id=self.scores_full_bar_id, - num_threads=32 * (1 + len(self.epilogue_warp_id)), - ) - softmax_done_barrier = pipeline.NamedBarrier( - barrier_id=self.softmax_done_bar_id, - num_threads=32 * (len(self.epilogue_warp_id) + 1), - ) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cluster_layout_vmnk, is_relaxed=True) - - sA = smem.allocate_tensor( - element_type=self.a_dtype, layout=a_smem_layout_staged.outer, - byte_alignment=128, swizzle=a_smem_layout_staged.inner) - sB = smem.allocate_tensor( - element_type=self.b_dtype, layout=b_smem_layout_staged.outer, - byte_alignment=128, swizzle=b_smem_layout_staged.inner) - sC = smem.allocate_tensor( - element_type=self.c_dtype, layout=c_smem_layout_staged.outer, - byte_alignment=128, swizzle=c_smem_layout_staged.inner) - - gA_mkl = cute.local_tile(mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)) - gB_nkl = cute.local_tile(mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None)) - gC_mnl = cute.local_tile(mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None)) - k_tile_cnt = cute.size(gA_mkl, mode=[3]) - - thr_mma1 = tiled_mma1.get_slice(0) - tCgA = thr_mma1.partition_A(gA_mkl) - tCgB = thr_mma1.partition_B(gB_nkl) - tCgC = thr_mma1.partition_C(gC_mnl) - - a_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape) - tAsA, tAgA = cpasync.tma_partition( - tma_atom_a, 0, a_cta_layout, - cute.group_modes(sA, 0, 3), cute.group_modes(tCgA, 0, 3)) - b_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape) - tBsB, tBgB = cpasync.tma_partition( - tma_atom_b, 0, b_cta_layout, - cute.group_modes(sB, 0, 3), cute.group_modes(tCgB, 0, 3)) - - tAgA_slice = tAgA[(None, 0, None, 0)] - tBgB_slice = tBgB[(None, 0, None, 0)] - - tCrA = tiled_mma1.make_fragment_A(sA) - tCrB = tiled_mma1.make_fragment_B(sB) - tCrA_mma2 = tiled_mma2.make_fragment_A(sA) - tCrB_mma2 = tiled_mma2.make_fragment_B(sB) - - acc_shape = tiled_mma1.partition_shape_C(self.mma_tiler_mn) - tCtAcc_fake = tiled_mma1.make_fragment_C(cute.append(acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cluster_layout_vmnk) - - # TMA LOAD WARP - if warp_idx == self.tma_warp_id: - ab_producer.reset() - peek_ab_empty_status = ab_producer.try_acquire() - for k_tile in cutlass.range(k_tile_cnt, unroll=1): - handle = ab_producer.acquire_and_advance(peek_ab_empty_status) - cute.copy(tma_atom_a, tAgA_slice[(None, handle.count)], tAsA[(None, handle.index)], - tma_bar_ptr=handle.barrier) - cute.copy(tma_atom_b, tBgB_slice[(None, handle.count)], tBsB[(None, handle.index)], - tma_bar_ptr=handle.barrier) - peek_ab_empty_status = cutlass.Boolean(1) - if handle.count + 1 < k_tile_cnt: - peek_ab_empty_status = ab_producer.try_acquire() - ab_producer.tail() - - # MMA WARP - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) - - tCtScores_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) - tCtScores = tCtScores_base[(None, None, None, 0)] - - output_tmem_ptr = cute.recast_ptr( - tmem_ptr + self.num_tmem_cols_per_region, dtype=self.acc_dtype) - tCtOutput_base = cute.make_tensor(output_tmem_ptr, tCtAcc_fake.layout) - tCtOutput = tCtOutput_base[(None, None, None, 0)] - - ab_consumer.reset() - peek_ab_full_status = cutlass.Boolean(1) - if is_leader_cta: - peek_ab_full_status = ab_consumer.try_wait() - - # MMA1: Q @ K^T -> tmem_scores - tiled_mma1.set(tcgen05.Field.ACCUMULATE, False) - for k_tile in range(k_tile_cnt): - if is_leader_cta: - handle = ab_consumer.wait_and_advance(peek_ab_full_status) - num_kblocks = cute.size(tCrA, mode=[2]) - for kblk_idx in cutlass.range(num_kblocks, unroll_full=True): - kblk_crd = (None, None, kblk_idx, handle.index) - cute.gemm(tiled_mma1, tCtScores, tCrA[kblk_crd], tCrB[kblk_crd], tCtScores) - handle.release() - peek_ab_full_status = cutlass.Boolean(1) - if handle.count + 1 < k_tile_cnt: - peek_ab_full_status = ab_consumer.try_wait() - - # Signal scores ready - scores_full_barrier.arrive() - # Wait for softmax done - softmax_done_barrier.arrive_and_wait() - - # MMA2: P @ V -> tmem_output - tiled_mma2.set(tcgen05.Field.ACCUMULATE, True) - num_kblocks_mma2 = cute.size(tCrB_mma2, mode=[2]) - for kblk_idx in cutlass.range(num_kblocks_mma2, unroll_full=True): - kblk_crd = (None, None, kblk_idx, 0) - cute.gemm(tiled_mma2, tCtOutput, tCrA_mma2[kblk_crd], tCrB_mma2[kblk_crd], tCtOutput) - - acc_producer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Producer, self.num_acc_stage) - if is_leader_cta: - acc_pipeline.producer_acquire(acc_producer_state) - acc_pipeline.producer_commit(acc_producer_state) - acc_producer_state.advance() - acc_pipeline.producer_tail(acc_producer_state) - - # EPILOGUE WARPS - if warp_idx < self.mma_warp_id: - tmem.allocate(self.total_tmem_cols) - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) - - tCtScores_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) - tCtScores = tCtScores_base[(None, None, None, 0)] - - output_tmem_ptr = cute.recast_ptr( - tmem_ptr + self.num_tmem_cols_per_region, dtype=self.acc_dtype) - tCtOutput_base = cute.make_tensor(output_tmem_ptr, tCtAcc_fake.layout) - - # Wait for scores - scores_full_barrier.arrive_and_wait() - - # Identity softmax: load, fill 1.0, store back - tiled_copy_t2r, tTR_tScores, tTR_rScores = utils.gemm.sm100.epilogue_tmem_copy_and_partition( - self, tidx, tCtScores, tCgC, epi_tile, self.use_2cta_instrs) - - cute.copy(tiled_copy_t2r, tTR_tScores, tTR_rScores) - - for idx in cutlass.range(cute.size(tTR_rScores), unroll_full=True): - tTR_rScores[idx] = self.acc_dtype(1.0) - - copy_atom_r2t = cute.make_copy_atom( - tcgen05.St16x128bOp(tcgen05.Repetition.x32, tcgen05.Unpack.NONE), - self.acc_dtype, - ) - tiled_copy_r2t = tcgen05.make_tmem_copy(copy_atom_r2t, tCtScores) - thr_copy_r2t = tiled_copy_r2t.get_slice(tidx) - tRT_tScores = thr_copy_r2t.partition_D(tCtScores) - tRT_rP = thr_copy_r2t.partition_S(tTR_rScores) - cute.copy(tiled_copy_r2t, tRT_rP, tRT_tScores) - cute.arch.fence_view_async_tmem_load() - - # Signal softmax done - softmax_done_barrier.arrive() - - # Store output - acc_consumer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Consumer, self.num_acc_stage) - - c_producer_group = pipeline.CooperativeGroup( - pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) - c_pipeline = pipeline.PipelineTmaStore.create( - num_stages=self.num_c_stage, producer_group=c_producer_group) - - mma_tile_coord_mnl = (0, 0, 0) - epilogue_op = const_expr(lambda x: x) - num_tiles_executed = 0 - - acc_consumer_state = utils.gemm.sm100.epilogue_tma_store( - self, tidx, warp_idx, tma_atom_c, tCtOutput_base, sC, tCgC, - epi_tile, num_tiles_executed, epilogue_op, - mma_tile_coord_mnl, acc_consumer_state, acc_pipeline, c_pipeline) - - c_pipeline.producer_tail() - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def test_stage_b(): - device = torch.device("cuda") - torch.manual_seed(42) - - m, n, k = 128, 128, 128 - - a = torch.randn(m, k, 1, dtype=torch.bfloat16, device="cuda") - b = torch.randn(n, k, 1, dtype=torch.bfloat16, device="cuda") - c = torch.zeros(m, n, 1, dtype=torch.bfloat16, device="cuda") - - # Identity softmax: P = 1.0, output = ones(M,N) @ V(N,K) - v = b[:, :, 0].float() # (128, 128) - ref = torch.ones(m, n, dtype=torch.float32) @ v - - import cutlass.torch as cutlass_torch - mA = cutlass_torch.from_dlpack(a).mark_layout_dynamic( - leading_dim=cutlass_torch.get_leading_dim(a)) - mB = cutlass_torch.from_dlpack(b).mark_layout_dynamic( - leading_dim=cutlass_torch.get_leading_dim(b)) - mC = cutlass_torch.from_dlpack(c).mark_layout_dynamic( - leading_dim=cutlass_torch.get_leading_dim(c)) - - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - kernel = StageBKernel(mma_tiler_mn=(128, 128), use_2cta_instrs=False, use_tma_store=True) - compiled = cute.compile(kernel, mA, mB, mC, stream) - - compiled(mA, mB, mC, stream) - torch.cuda.synchronize() - - output = c[:, :, 0].float() - cos = torch.nn.functional.cosine_similarity( - output.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - max_err = (output - ref).abs().max().item() - - print("Stage B: Q @ K^T -> identity_softmax(P=1) @ V -> output") - print(" Cosine: {:.6f}, Max error: {:.6f}".format(cos, max_err)) - print(" {}".format("PASS" if cos >= 0.99 else "FAIL")) - return cos - - -if __name__ == "__main__": - test_stage_b() diff --git a/tests/archive/test_stage_b_v10.py b/tests/archive/test_stage_b_v10.py deleted file mode 100644 index 83169757..00000000 --- a/tests/archive/test_stage_b_v10.py +++ /dev/null @@ -1,321 +0,0 @@ -"""Stage B v10: Use FP32 ld→st on FULL (128,128) layout, not subview. -The ld reads S0 (full C-fragment), the st writes S1 (full C-fragment at offset 128), -then PV MMA reads S1 via A-fragment. - -Key insight: The FP32 ld→st on the SAME layout works (cosine 0.999999). -The BF16 recast pattern with DIFFERENT layout sizes is broken. -For identity softmax, we need to write the data back to TMEM in a format -that the PV MMA's A-fragment can read. Since the C-fragment and A-fragment -both access the same physical TMEM, writing via C-fragment (store) and -reading via A-fragment (PV MMA) should work IF the data is in the right place. - -The pv_mma_tiler has N=128 (from mma_tiler_mn[1]) but P's K dimension = 128. -Wait, pv_mma_tiler = (M, K, K_inner*4). For the PV MMA, A=P with K=128 (full KV dim). -But p_tmem_s is make_smem_layout_a(pv_mma, pv_mma_tiler, BF16, 1) which gives the -A-fragment layout for (128, 128, 64). The A-fragment covers K=64 at CTA level... -Hmm, pv_mma_tiler[1] should be the V dimension (N in MMA terms), not K. - -Actually, for PV MMA: P (A operand, M×K) × V (B operand, K×N) → O (C operand, M×N) -With a_source=TMEM, the P operand's K dimension matches the CTA tile K. -pv_mma_tiler = (128, 128, 64): M=128, N=128, K=64. -So P is M×K = 128×64. That's the A-fragment shape. - -So the A-fragment only reads 64 columns of TMEM (K=64 in 4 blocks of 16). -The C-fragment for S has 128 columns. The P region starts at offset 32 (tmem_p0_offset). -With s_cols=128 and tmem_p0_offset=32, the P data starts at column 32 and spans 64 columns -(columns 32-95), which is within the S region (columns 0-127). - -So the store should write to columns 32-95 of the S region. The C-fragment composition -(128, tilePlikeFP32=64) with offset 32 does exactly this. - -But we showed the subview store + recast doesn't work. Let me try: write the FULL -C-fragment (128×128) at offset 128 (tStS1), and adjust the PV MMA's A-fragment -offset so it reads from the right columns within that region. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - -class StageBv10: - def __init__(self, mma_tiler_mn): - self.qk_acc_dtype = Float32; self.q_dtype = BFloat16; self.o_dtype = BFloat16 - self.c_dtype = BFloat16; self.acc_dtype = Float32 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2; self.use_2cta_instrs = False - self.epilog_sync_bar_id = 1 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - pv_inst_k = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (*self.mma_tiler_mn, pv_inst_k * 4) - self.mma_tiler = self.qk_mma_tiler - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - c_layout = LayoutEnum.ROW_MAJOR; self.c_layout = c_layout - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, False, c_layout, self.o_dtype) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, c_layout, self.epi_tile, 2) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - self.s_cols = find_tmem_tensor_col_offset(tStS) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - self.o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = self.s_cols - self.tmem_alloc_cols = self.s_cols + self.o_cols # 256 - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, 1)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, 1)) - self.num_tmem_alloc_cols = utils.get_num_tmem_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, - LayoutEnum.from_tensor(a).mma_major_mode(), - LayoutEnum.from_tensor(b).mma_major_mode(), - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, - tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, - cute.nvgpu.OperandMajorMode.K, - LayoutEnum.from_tensor(b).mma_major_mode(), - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, - tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - tma_a, tma_ta = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - a, a_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_b, tma_tb = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - b, b_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - self._kernel(qk_mma, pv_mma, tma_a, tma_ta, tma_b, tma_tb, tma_c, tma_tc, - self.cluster_layout_vmnk, self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_a, mA, tma_b, mB, tma_c, mC, cl_vmnk, - a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_a); cpasync.prefetch_descriptor(tma_b); cpasync.prefetch_descriptor(tma_c) - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, len(self.epilogue_warp_id)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=False, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - sA = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sB = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV_ptr = cute.recast_ptr(sB.iterator, v_smem_s.inner) - sV = cute.make_tensor(sV_ptr, v_smem_s.outer) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - gA = cute.local_tile(mA, cute.slice_(self.mma_tiler, (None,0,None)), (None,None,None)) - gB = cute.local_tile(mB, cute.slice_(self.mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gA, mode=[3]) - qk_thr = qk_mma.get_slice(0) - tCgA = qk_thr.partition_A(gA); tCgB = qk_thr.partition_B(gB); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsA, tAgA = cpasync.tma_partition(tma_a, 0, a_lay, cute.group_modes(sA,0,3), cute.group_modes(tCgA,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsB, tBgB = cpasync.tma_partition(tma_b, 0, b_lay, cute.group_modes(sB,0,3), cute.group_modes(tCgB,0,3)) - tAgA = tAgA[(None,0,None,0)]; tBgB = tBgB[(None,0,None,0)] - tCrA = qk_mma.make_fragment_A(sA); tCrB = qk_mma.make_fragment_B(sB) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - # P A-fragment - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP = pv_thr.make_fragment_A(tP)[None, None, None, 0] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - # LD and ST copy atoms on the SAME layout (full C-fragment) - tmem_ld = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tmem_st = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_ld = tcgen05.make_tmem_copy(tmem_ld, tStS0) - tiled_st = tcgen05.make_tmem_copy(tmem_st, tStS0) # SAME layout for ld and st - sfw = tidx % (32 * len(self.epilogue_warp_id)) - thr_ld = tiled_ld.get_slice(sfw) - thr_st = tiled_st.get_slice(sfw) - tLdS = thr_ld.partition_S(tStS0) - tStS_dst = thr_st.partition_D(tStS0) # St writes to S0 (same as MMA output) - cS_id = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS_id) - tLdcS = thr_ld.partition_D(tScS) - tStcS = thr_st.partition_S(tScS) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, 1)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, 1)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # TMA - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_a, tAgA[(None,h.count)], tAsA[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_b, tBgB[(None,h.count)], tBsB[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_stage_b_v11.py b/tests/archive/test_stage_b_v11.py deleted file mode 100644 index 25cb4445..00000000 --- a/tests/archive/test_stage_b_v11.py +++ /dev/null @@ -1,210 +0,0 @@ -"""Stage B v11: Backward FMHA pattern exactly. - -1. ld FP32 from S0 (C-fragment) -2. Quantize FP32→BF16 (same register shape, .load()/.store()) -3. Reshape BF16 register to store partition shape -4. St32x32bOp BF16 store to tdVrP (A-fragment layout, offset s_cols) -5. PV MMA reads from tOrP0 (A-fragment, offset s_cols) -6. Epilogue writes output -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - -class StageBv11: - def __init__(self, mma_tiler_mn): - self.qk_acc_dtype = Float32; self.q_dtype = BFloat16; self.o_dtype = BFloat16 - self.c_dtype = BFloat16; self.acc_dtype = Float32 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.num_c_stage = 2; self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - pv_inst_k = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (*self.mma_tiler_mn, pv_inst_k * 4) - self.mma_tiler = self.qk_mma_tiler - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - c_layout = LayoutEnum.ROW_MAJOR; self.c_layout = c_layout - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, c_layout, self.o_dtype) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, c_layout, self.epi_tile, 2) - self.num_ab_stage = 1; self.num_acc_stage = 1 - qk_thr = qk_mma.get_slice(0); qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape); self.s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0); pv_acc_shape = pv_thr.partition_shape_C(self.mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape); self.o_cols = find_tmem_tensor_col_offset(tOtO) - self.tmem_s0_offset = 0 - self.tmem_o0_offset = self.s_cols * 2 # After S0 (128) + P (128 BF16 = 64 FP32, but A-frag offset uses width ratio) - self.tmem_alloc_cols = 512 # Enough for S0 + P + O - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, 1)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, 1)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)); b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = (cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem)) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, LayoutEnum.from_tensor(a).mma_major_mode(), LayoutEnum.from_tensor(b).mma_major_mode(), self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, LayoutEnum.from_tensor(b).mma_major_mode(), self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)); b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - tma_a, tma_ta = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), a, a_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_b, tma_tb = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), b, b_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - self._kernel(qk_mma, pv_mma, tma_a, tma_ta, tma_b, tma_tb, tma_c, tma_tc, self.cluster_layout_vmnk, self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_a, mA, tma_b, mB, tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()); tidx, _, _ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: cpasync.prefetch_descriptor(tma_a); cpasync.prefetch_descriptor(tma_b); cpasync.prefetch_descriptor(tma_c) - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2]; mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2]; tmem_dealloc: cutlass.Int64; holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - ab_p, ab_c = pipeline.PipelineTmaUmma.create(barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True).make_participants() - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), cta_layout_vmnk=cl_vmnk, defer_sync=True).make_participants() - acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, len(self.epilogue_warp_id)), cta_layout_vmnk=cl_vmnk, defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2, num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=False, two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - sA = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sB = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV_ptr = cute.recast_ptr(sB.iterator, v_smem_s.inner); sV = cute.make_tensor(sV_ptr, v_smem_s.outer) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - gA = cute.local_tile(mA, cute.slice_(self.mma_tiler, (None,0,None)), (None,None,None)) - gB = cute.local_tile(mB, cute.slice_(self.mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gA, mode=[3]) - qk_thr = qk_mma.get_slice(0); tCgA = qk_thr.partition_A(gA); tCgB = qk_thr.partition_B(gB); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsA, tAgA = cpasync.tma_partition(tma_a, 0, a_lay, cute.group_modes(sA,0,3), cute.group_modes(tCgA,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsB, tBgB = cpasync.tma_partition(tma_b, 0, b_lay, cute.group_modes(sB,0,3), cute.group_modes(tCgB,0,3)) - tAgA = tAgA[(None,0,None,0)]; tBgB = tBgB[(None,0,None,0)] - tCrA = qk_mma.make_fragment_A(sA); tCrB = qk_mma.make_fragment_B(sB) - tCrV = pv_mma.make_fragment_B(sV) - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator, tStS.layout) - pv_thr = pv_mma.get_slice(0); pv_acc_shape = pv_thr.partition_shape_C(self.mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - # ── P A-fragment (backward FMHA pattern) ── - tP_iter = cute.recast_ptr(tStS.iterator, dtype=self.q_dtype) - tP = cute.make_tensor(tP_iter, p_tmem_s.outer) - tOrP = pv_thr.make_fragment_A(tP)[None, None, None, 0] - # Apply offset for P region (after S0) - p_bf16_offset = self.qk_acc_dtype.width // self.q_dtype.width * self.s_cols - tdVrP = cute.make_tensor(tOrP.iterator + p_bf16_offset, tOrP.layout) - tOrP0 = cute.make_tensor(tOrP.iterator + p_bf16_offset, tOrP.layout) - # ── TMEM LOAD from C-fragment ── - tmem_ld = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_ld = tcgen05.make_tmem_copy(tmem_ld, tStS0) - sfw = tidx % (32 * len(self.epilogue_warp_id)) - thr_ld = tiled_ld.get_slice(sfw) - tLdS = thr_ld.partition_S(tStS0) - cS_id = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS_id) - tLdcS = thr_ld.partition_D(tScS) - # ── TMEM STORE via A-fragment layout ── - tmem_st = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(8)), self.q_dtype) - tiled_st = tcgen05.make_tmem_copy(tmem_st, tdVrP) - thr_st = tiled_st.get_slice(sfw) - tStP = thr_st.partition_D(tdVrP) - tdVcST = pv_thr.partition_A(cS_id) # A-operand partition of identity - tStcS = thr_st.partition_S(tdVcST) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, 1)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - # TMA - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek); cute.copy(tma_a, tAgA[(None,h.count)], tAsA[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_b, tBgB[(None,h.count)], tBsB[(None,h.index)], tma_bar_ptr=h.barrier); peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_stage_b_v11b.py b/tests/archive/test_stage_b_v11b.py deleted file mode 100644 index 91cd25ae..00000000 --- a/tests/archive/test_stage_b_v11b.py +++ /dev/null @@ -1,398 +0,0 @@ -""" -Stage B v11b: Identity Softmax - store to composition(tP, (128,128)) - -Key fixes over v6: - - TMEM offsets computed via find_tmem_tensor_col_offset (same API as get_num_tmem_alloc_cols) - - P tensor constructed from p_tmem_s.outer (matching fmha.py pattern exactly) - - tilePlikeFP32 computed from qk_mma_tiler and dtype widths - - tmem_alloc_cols from get_num_tmem_alloc_cols with all fragments - - JIT-time diagnostic prints of all TMEM sizes - -Architecture (matches fmha.py exactly): - MMA1: Q @ K^T → tmem_scores (a_source=SMEM, accumulate=False) - Identity softmax: tcgen05.ld C-layout → F32→BF16 → tcgen05.st A-layout - MMA2: P @ V → tmem_output (a_source=TMEM, accumulate=True) -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - - -class StageBIdentitySoftmax: - def __init__(self, mma_tiler_mn, use_2cta_instrs=False, use_tma_store=True): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16 - self.use_2cta_instrs = use_2cta_instrs; self.use_tma_store = use_tma_store - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.epilog_sync_bar_id = 1; self.tmem_alloc_sync_bar_id = 2; self.tmem_dealloc_sync_bar_id = 3 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - pv_inst_k = cute.size(pv_mma.shape_mnk, mode=[2]) - # FIX: PV tiler swaps N and K from QK (fmha.py: pv_mma_tiler = (M, QK_K, QK_N)) - # P is (M, QK_N) and V is (QK_N, D), so PV MMA has K=QK_N, N=D - self.pv_mma_tiler = (self.qk_mma_tiler[0], self.qk_mma_tiler[2], self.qk_mma_tiler[1]) - self.mma_tiler = self.qk_mma_tiler - print(f"[StageB] qk_mma.shape_mnk = {qk_mma.shape_mnk}") - print(f"[StageB] pv_mma.shape_mnk = {pv_mma.shape_mnk}") - print(f"[StageB] qk_mma_tiler = {self.qk_mma_tiler}") - print(f"[StageB] pv_mma_tiler = {self.pv_mma_tiler}") - print(f"[StageB] qk_inst_k = {qk_inst_k}, pv_inst_k = {pv_inst_k}") - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], - self.qk_mma_tiler[2], - ) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, self.use_2cta_instrs, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.a_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.b_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.b_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - # ── COMPUTE TMEM OFFSETS USING find_tmem_tensor_col_offset ── - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - # tilePlikeFP32 for the store-side composition - self.tilePlikeFP32 = self.qk_mma_tiler[1] * self.q_dtype.width // 32 - - # ── TMEM LAYOUT (matching fmha.py) ── - # P OVERLAPS S — after softmax, P (A-layout) is written on top of scores (C-layout) - # in the same TMEM region. The A-layout view starts partway through the S region. - # fmha.py: S0=0, P0=32, O0=256 (with S1=128, P1=160 double-buffered) - # The P offset of 32 means the A-layout starts at column 32 within the S region. - # This is because the C-layout and A-layout partition TMEM differently per-thread; - # the first 32 C-layout columns are "dead space" in the A-layout mapping. - # - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 # Same as fmha.py - self.tmem_o0_offset = s_cols # 128 - self.tmem_alloc_cols = s_cols + o_cols # 256 - - # Also compute via get_num_tmem_alloc_cols for the full allocation - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, 1)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, 1)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - print(f"[StageB] s_cols (QK accumulator) = {s_cols}") - print(f"[StageB] o_cols (PV accumulator) = {o_cols}") - print(f"[StageB] tilePlikeFP32 = {self.tilePlikeFP32}") - print(f"[StageB] tmem_s0_offset = {self.tmem_s0_offset}") - print(f"[StageB] tmem_p0_offset = {self.tmem_p0_offset}") - print(f"[StageB] tmem_o0_offset = {self.tmem_o0_offset}") - print(f"[StageB] tmem_alloc_cols (computed) = {self.tmem_alloc_cols}") - print(f"[StageB] num_tmem_alloc_cols (via utils) = {self.num_tmem_alloc_cols}") - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.a_dtype, a_smem) + cute.size_in_bytes(self.b_dtype, b_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - self.a_dtype = a.element_type; self.b_dtype = b.element_type; self.c_dtype = c.element_type - self.a_major = LayoutEnum.from_tensor(a).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(b).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, cute.nvgpu.OperandMajorMode.K, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - tma_a, tma_ta = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - a, a_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_b, tma_tb = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - b, b_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_a, tma_ta, tma_b, tma_tb, tma_c, tma_tc, - self.cluster_layout_vmnk, self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_a, mA, tma_b, mB, tma_c, mC, cl_vmnk, - a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_a); cpasync.prefetch_descriptor(tma_b); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=self.tmem_alloc_sync_bar_id, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sA = smem.allocate_tensor(element_type=self.a_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sB = smem.allocate_tensor(element_type=self.b_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - # V shares the same SMEM as B (same data, different layout for PV MMA) - sV_ptr = cute.recast_ptr(sB.iterator, v_smem_s.inner) - sV = cute.make_tensor(sV_ptr, v_smem_s.outer) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gA = cute.local_tile(mA, cute.slice_(self.mma_tiler, (None,0,None)), (None,None,None)) - gB = cute.local_tile(mB, cute.slice_(self.mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gA, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - tCgA = qk_thr.partition_A(gA); tCgB = qk_thr.partition_B(gB); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsA, tAgA = cpasync.tma_partition(tma_a, 0, a_lay, cute.group_modes(sA,0,3), cute.group_modes(tCgA,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsB, tBgB = cpasync.tma_partition(tma_b, 0, b_lay, cute.group_modes(sB,0,3), cute.group_modes(tCgB,0,3)) - tAgA = tAgA[(None,0,None,0)]; tBgB = tBgB[(None,0,None,0)] - - tCrA = qk_mma.make_fragment_A(sA); tCrB = qk_mma.make_fragment_B(sB) - tCrV = pv_mma.make_fragment_B(sV) # V fragment from V SMEM layout - print(f"[DIAG] tCrV.size = {cute.size(tCrV)}") - print(f"[DIAG] tCrA.size = {cute.size(tCrA)}") - print(f"[DIAG] tCrB.size = {cute.size(tCrB)}") - print(f"[DIAG] nblk_qk (tCrA mode 2) = {cute.size(tCrA, mode=[2])}") - - # ── TMEM tensors with computed offsets (matching fmha.py pattern) ── - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - # P fragment: construct from p_tmem_s layout (matching fmha.py exactly) - # fmha.py: tP = cute.make_tensor(tStS.iterator, p_tmem_layout_staged.outer) - # tOrP = pv_thr_mma.make_fragment_A(tP)[None, None, None, 0] - # tOrP0 = cute.make_tensor(tOrP.iterator + dtype_width_ratio * tmem_p0_offset, tOrP.layout) - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - # DIAGNOSTIC: Compare tP (A-layout) vs tStS_P (composition) - tilePlikeFP32 = self.qk_mma_tiler[1] * self.q_dtype.width // 32 - tStS_P_layout = cute.composition(tStS.layout, cute.make_layout((128, tilePlikeFP32))) - tStS_P = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStS_P_layout) - print(f'[DIAG] tP.layout: {tP.layout}') - print(f'[DIAG] tP.size: {cute.size(tP)}') - print(f'[DIAG] tP.element_type: {tP.element_type if hasattr(tP, 'element_type') else 'N/A'}') - print(f'[DIAG] tStS_P.layout: {tStS_P.layout}') - print(f'[DIAG] tStS_P.size: {cute.size(tStS_P)}') - print(f'[DIAG] tStS_P.element_type: {tStS_P.element_type if hasattr(tStS_P, 'element_type') else 'N/A'}') - print(f'[DIAG] tilePlikeFP32: {tilePlikeFP32}') - print(f'[DIAG] tP and tStS_P same iterator? {tP.iterator == tStS_P.iterator if hasattr(tP, 'iterator') else 'cant compare'}') - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ── TMA WARP ── - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_a, tAgA[(None,h.count)], tAsA[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_b, tBgB[(None,h.count)], tBsB[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_stage_b_v12.py b/tests/archive/test_stage_b_v12.py deleted file mode 100644 index e392fa3d..00000000 --- a/tests/archive/test_stage_b_v12.py +++ /dev/null @@ -1,408 +0,0 @@ -""" -Stage B v12: BF16 St16x128bOp store to tP A-layout, P=all-ones - -Key fixes over v6: - - TMEM offsets computed via find_tmem_tensor_col_offset (same API as get_num_tmem_alloc_cols) - - P tensor constructed from p_tmem_s.outer (matching fmha.py pattern exactly) - - tilePlikeFP32 computed from qk_mma_tiler and dtype widths - - tmem_alloc_cols from get_num_tmem_alloc_cols with all fragments - - JIT-time diagnostic prints of all TMEM sizes - -Architecture (matches fmha.py exactly): - MMA1: Q @ K^T → tmem_scores (a_source=SMEM, accumulate=False) - Identity softmax: tcgen05.ld C-layout → F32→BF16 → tcgen05.st A-layout - MMA2: P @ V → tmem_output (a_source=TMEM, accumulate=True) -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - - -class StageBIdentitySoftmax: - def __init__(self, mma_tiler_mn, use_2cta_instrs=False, use_tma_store=True): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16 - self.use_2cta_instrs = use_2cta_instrs; self.use_tma_store = use_tma_store - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.epilog_sync_bar_id = 1; self.tmem_alloc_sync_bar_id = 2; self.tmem_dealloc_sync_bar_id = 3 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - pv_inst_k = cute.size(pv_mma.shape_mnk, mode=[2]) - # FIX: PV tiler swaps N and K from QK (fmha.py: pv_mma_tiler = (M, QK_K, QK_N)) - # P is (M, QK_N) and V is (QK_N, D), so PV MMA has K=QK_N, N=D - self.pv_mma_tiler = (self.qk_mma_tiler[0], self.qk_mma_tiler[2], self.qk_mma_tiler[1]) - self.mma_tiler = self.qk_mma_tiler - print(f"[StageB] qk_mma.shape_mnk = {qk_mma.shape_mnk}") - print(f"[StageB] pv_mma.shape_mnk = {pv_mma.shape_mnk}") - print(f"[StageB] qk_mma_tiler = {self.qk_mma_tiler}") - print(f"[StageB] pv_mma_tiler = {self.pv_mma_tiler}") - print(f"[StageB] qk_inst_k = {qk_inst_k}, pv_inst_k = {pv_inst_k}") - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], - self.qk_mma_tiler[2], - ) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, self.use_2cta_instrs, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.a_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.b_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.b_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - # ── COMPUTE TMEM OFFSETS USING find_tmem_tensor_col_offset ── - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - # tilePlikeFP32 for the store-side composition - self.tilePlikeFP32 = self.qk_mma_tiler[1] * self.q_dtype.width // 32 - - # ── TMEM LAYOUT (matching fmha.py) ── - # P OVERLAPS S — after softmax, P (A-layout) is written on top of scores (C-layout) - # in the same TMEM region. The A-layout view starts partway through the S region. - # fmha.py: S0=0, P0=32, O0=256 (with S1=128, P1=160 double-buffered) - # The P offset of 32 means the A-layout starts at column 32 within the S region. - # This is because the C-layout and A-layout partition TMEM differently per-thread; - # the first 32 C-layout columns are "dead space" in the A-layout mapping. - # - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 # Same as fmha.py - self.tmem_o0_offset = s_cols # 128 - self.tmem_alloc_cols = s_cols + o_cols # 256 - - # Also compute via get_num_tmem_alloc_cols for the full allocation - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, 1)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, 1)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - print(f"[StageB] s_cols (QK accumulator) = {s_cols}") - print(f"[StageB] o_cols (PV accumulator) = {o_cols}") - print(f"[StageB] tilePlikeFP32 = {self.tilePlikeFP32}") - print(f"[StageB] tmem_s0_offset = {self.tmem_s0_offset}") - print(f"[StageB] tmem_p0_offset = {self.tmem_p0_offset}") - print(f"[StageB] tmem_o0_offset = {self.tmem_o0_offset}") - print(f"[StageB] tmem_alloc_cols (computed) = {self.tmem_alloc_cols}") - print(f"[StageB] num_tmem_alloc_cols (via utils) = {self.num_tmem_alloc_cols}") - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.a_dtype, a_smem) + cute.size_in_bytes(self.b_dtype, b_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - self.a_dtype = a.element_type; self.b_dtype = b.element_type; self.c_dtype = c.element_type - self.a_major = LayoutEnum.from_tensor(a).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(b).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, cute.nvgpu.OperandMajorMode.K, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - tma_a, tma_ta = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - a, a_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_b, tma_tb = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - b, b_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_a, tma_ta, tma_b, tma_tb, tma_c, tma_tc, - self.cluster_layout_vmnk, self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_a, mA, tma_b, mB, tma_c, mC, cl_vmnk, - a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_a); cpasync.prefetch_descriptor(tma_b); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=self.tmem_alloc_sync_bar_id, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sA = smem.allocate_tensor(element_type=self.a_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sB = smem.allocate_tensor(element_type=self.b_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - # V shares the same SMEM as B (same data, different layout for PV MMA) - sV_ptr = cute.recast_ptr(sB.iterator, v_smem_s.inner) - sV = cute.make_tensor(sV_ptr, v_smem_s.outer) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gA = cute.local_tile(mA, cute.slice_(self.mma_tiler, (None,0,None)), (None,None,None)) - gB = cute.local_tile(mB, cute.slice_(self.mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gA, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - tCgA = qk_thr.partition_A(gA); tCgB = qk_thr.partition_B(gB); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsA, tAgA = cpasync.tma_partition(tma_a, 0, a_lay, cute.group_modes(sA,0,3), cute.group_modes(tCgA,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsB, tBgB = cpasync.tma_partition(tma_b, 0, b_lay, cute.group_modes(sB,0,3), cute.group_modes(tCgB,0,3)) - tAgA = tAgA[(None,0,None,0)]; tBgB = tBgB[(None,0,None,0)] - - tCrA = qk_mma.make_fragment_A(sA); tCrB = qk_mma.make_fragment_B(sB) - tCrV = pv_mma.make_fragment_B(sV) # V fragment from V SMEM layout - print(f"[DIAG] tCrV.size = {cute.size(tCrV)}") - print(f"[DIAG] tCrA.size = {cute.size(tCrA)}") - print(f"[DIAG] tCrB.size = {cute.size(tCrB)}") - print(f"[DIAG] nblk_qk (tCrA mode 2) = {cute.size(tCrA, mode=[2])}") - - # ── TMEM tensors with computed offsets (matching fmha.py pattern) ── - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - # P fragment: construct from p_tmem_s layout (matching fmha.py exactly) - # fmha.py: tP = cute.make_tensor(tStS.iterator, p_tmem_layout_staged.outer) - # tOrP = pv_thr_mma.make_fragment_A(tP)[None, None, None, 0] - # tOrP0 = cute.make_tensor(tOrP.iterator + dtype_width_ratio * tmem_p0_offset, tOrP.layout) - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - print(f'[DIAG] p_tmem_s.outer: {p_tmem_s.outer}') - print(f'[DIAG] p_tmem_s.inner: {p_tmem_s.inner}') - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - # DIAGNOSTIC: Compare tP (A-layout) vs tStS_P (composition) - tilePlikeFP32 = self.qk_mma_tiler[1] * self.q_dtype.width // 32 - tStS_P_layout = cute.composition(tStS.layout, cute.make_layout((128, tilePlikeFP32))) - tStS_P = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStS_P_layout) - print(f'[DIAG] tP.layout: {tP.layout}') - print(f'[DIAG] tP.size: {cute.size(tP)}') - print(f'[DIAG] tP.element_type: {tP.element_type if hasattr(tP, 'element_type') else 'N/A'}') - print(f'[DIAG] tStS_P.layout: {tStS_P.layout}') - print(f'[DIAG] tStS_P.size: {cute.size(tStS_P)}') - print(f'[DIAG] tStS_P.element_type: {tStS_P.element_type if hasattr(tStS_P, 'element_type') else 'N/A'}') - print(f'[DIAG] tilePlikeFP32: {tilePlikeFP32}') - print(f'[DIAG] tP and tStS_P same iterator? {tP.iterator == tStS_P.iterator if hasattr(tP, 'iterator') else 'cant compare'}') - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ── TMA WARP ── - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_a, tAgA[(None,h.count)], tAsA[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_b, tBgB[(None,h.count)], tBsB[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_stage_b_v13.py b/tests/archive/test_stage_b_v13.py deleted file mode 100644 index a6fd803b..00000000 --- a/tests/archive/test_stage_b_v13.py +++ /dev/null @@ -1,401 +0,0 @@ -""" -Stage B v13: Two MMAs + Identity Softmax using FMHA's C-fragment packing pattern. - -The key insight: the C→A "transform" is NOT a reordering — it's a PACKING. -When you write 128 BF16 values packed into 64 FP32 words via the C-fragment -composition (128, tilePlikeFP32) with St32x32bOp as FP32, the physical TMEM -locations used are exactly the ones the PV MMA's A-fragment reads from. - -FMHA pattern: -1. Load S from TMEM via C-fragment layout (FP32, 128×128) -2. Convert to BF16 and pack: FP32 backing tensor + BF16 recast view -3. Store packed FP32 backing to TMEM via C-fragment composition with St32x32bOp(FP32) -4. PV MMA reads from A-fragment TMEM — same physical locations as the packed BF16 - -Architecture: - MMA1: Q @ K^T → tmem_scores (a_source=SMEM, accumulate=False) - Identity softmax: tcgen05.ld C-layout → F32→BF16 packed → tcgen05.st C-fragment composition - MMA2: P @ V → tmem_output (a_source=TMEM, accumulate=True) -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - - -class StageBIdentitySoftmax: - def __init__(self, mma_tiler_mn, use_2cta_instrs=False, use_tma_store=True): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16 - self.use_2cta_instrs = use_2cta_instrs; self.use_tma_store = use_tma_store - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.epilog_sync_bar_id = 1; self.tmem_alloc_sync_bar_id = 2; self.tmem_dealloc_sync_bar_id = 3 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - # FMHA pattern: pv_mma_tiler = (M, QK_K, QK_N) — K=QK_N, N=head_dim - self.pv_mma_tiler = (self.qk_mma_tiler[0], self.qk_mma_tiler[2], self.qk_mma_tiler[1]) - self.mma_tiler = self.qk_mma_tiler - print(f"[StageB] qk_mma.shape_mnk = {qk_mma.shape_mnk}") - print(f"[StageB] pv_mma.shape_mnk = {pv_mma.shape_mnk}") - print(f"[StageB] qk_mma_tiler = {self.qk_mma_tiler}") - print(f"[StageB] pv_mma_tiler = {self.pv_mma_tiler}") - - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], - self.qk_mma_tiler[2], - ) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, self.use_2cta_instrs, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.a_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.b_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.b_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - # TMEM offsets (matching fmha.py) - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - print(f"[StageB] tilePlikeFP32 = {self.tilePlikeFP32}") - - # FMHA TMEM layout: S0=0, P0=32, O0=256 - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - self.tmem_alloc_cols = s_cols + o_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - print(f"[StageB] s_cols = {s_cols}, o_cols = {o_cols}") - print(f"[StageB] tmem_alloc_cols = {self.tmem_alloc_cols}") - print(f"[StageB] num_tmem_alloc_cols = {self.num_tmem_alloc_cols}") - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.a_dtype, a_smem) + cute.size_in_bytes(self.b_dtype, b_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - self.a_dtype = a.element_type; self.b_dtype = b.element_type; self.c_dtype = c.element_type - self.a_major = LayoutEnum.from_tensor(a).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(b).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, cute.nvgpu.OperandMajorMode.K, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - tma_a, tma_ta = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - a, a_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_b, tma_tb = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - b, b_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_a, tma_ta, tma_b, tma_tb, tma_c, tma_tc, - self.cluster_layout_vmnk, self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_a, mA, tma_b, mB, tma_c, mC, cl_vmnk, - a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_a); cpasync.prefetch_descriptor(tma_b); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=self.tmem_alloc_sync_bar_id, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sA = smem.allocate_tensor(element_type=self.a_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sB = smem.allocate_tensor(element_type=self.b_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - # V shares SMEM with B (same data, different layout for PV MMA) - sV_ptr = cute.recast_ptr(sB.iterator, v_smem_s.inner) - sV = cute.make_tensor(sV_ptr, v_smem_s.outer) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gA = cute.local_tile(mA, cute.slice_(self.mma_tiler, (None,0,None)), (None,None,None)) - gB = cute.local_tile(mB, cute.slice_(self.mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gA, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - tCgA = qk_thr.partition_A(gA); tCgB = qk_thr.partition_B(gB); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsA, tAgA = cpasync.tma_partition(tma_a, 0, a_lay, cute.group_modes(sA,0,3), cute.group_modes(tCgA,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsB, tBgB = cpasync.tma_partition(tma_b, 0, b_lay, cute.group_modes(sB,0,3), cute.group_modes(tCgB,0,3)) - tAgA = tAgA[(None,0,None,0)]; tBgB = tBgB[(None,0,None,0)] - - tCrA = qk_mma.make_fragment_A(sA); tCrB = qk_mma.make_fragment_B(sB) - tCrV = pv_mma.make_fragment_B(sV) - - # ── TMEM tensors ── - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - # ── P A-fragment for PV MMA (matching fmha.py exactly) ── - # tP uses C-fragment iterator but A-fragment (p_tmem_s) layout - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - # ── Softmax store uses C-fragment composition (FMHA pattern) ── - # tStS_P: C-fragment layout composed with (128, tilePlikeFP32) - # This is where we store the packed BF16 P values - tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - tStS_P_layout = cute.composition(tStS.layout, cute.make_layout((128, tilePlikeFP32))) - tStS_P = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStS_P_layout) - - print(f'[DIAG] tStS.layout: {tStS.layout}') - print(f'[DIAG] tStS_P.layout: {tStS_P.layout}') - print(f'[DIAG] tP.layout: {tP.layout}') - print(f'[DIAG] tOrP.layout: {tOrP.layout}') - print(f'[DIAG] tOrP0.layout: {tOrP0.layout}') - print(f'[DIAG] tilePlikeFP32: {tilePlikeFP32}') - print(f'[DIAG] qk_mma_tiler: {self.qk_mma_tiler}') - print(f'[DIAG] pv_mma_tiler: {self.pv_mma_tiler}') - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ── TMA WARP ── - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_a, tAgA[(None,h.count)], tAsA[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_b, tBgB[(None,h.count)], tBsB[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_stage_b_v14.py b/tests/archive/test_stage_b_v14.py deleted file mode 100644 index ecd90772..00000000 --- a/tests/archive/test_stage_b_v14.py +++ /dev/null @@ -1,352 +0,0 @@ -""" -Stage B v14: Two MMAs + Identity Softmax using backward FMHA's A-fragment store pattern. - -The backward FMHA writes P to TMEM using the A-fragment layout directly: -1. Load S from TMEM via C-fragment layout (FP32) -2. Quantize FP32 -> BF16 (make_rmem_tensor with BF16, load/store) -3. Reshape quantized to match store coordinate shape -4. Store via St32x32bOp(BF16) to A-fragment TMEM layout -5. PV MMA reads from the same A-fragment addresses -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - - -class StageBIdentitySoftmax: - def __init__(self, mma_tiler_mn, use_2cta_instrs=False, use_tma_store=True): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16 - self.use_2cta_instrs = use_2cta_instrs; self.use_tma_store = use_tma_store - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.epilog_sync_bar_id = 1; self.tmem_alloc_sync_bar_id = 2; self.tmem_dealloc_sync_bar_id = 3 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - self.pv_mma_tiler = (self.qk_mma_tiler[0], self.qk_mma_tiler[2], self.qk_mma_tiler[1]) - self.mma_tiler = self.qk_mma_tiler - print(f"[StageB] qk_mma_tiler = {self.qk_mma_tiler}") - print(f"[StageB] pv_mma_tiler = {self.pv_mma_tiler}") - - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], - self.qk_mma_tiler[2], - ) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, self.use_2cta_instrs, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.a_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.b_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.b_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - self.tmem_alloc_cols = s_cols + o_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.a_dtype, a_smem) + cute.size_in_bytes(self.b_dtype, b_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - self.a_dtype = a.element_type; self.b_dtype = b.element_type; self.c_dtype = c.element_type - self.a_major = LayoutEnum.from_tensor(a).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(b).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, cute.nvgpu.OperandMajorMode.K, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - tma_a, tma_ta = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - a, a_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_b, tma_tb = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - b, b_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_a, tma_ta, tma_b, tma_tb, tma_c, tma_tc, - self.cluster_layout_vmnk, self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_a, mA, tma_b, mB, tma_c, mC, cl_vmnk, - a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_a); cpasync.prefetch_descriptor(tma_b); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=self.tmem_alloc_sync_bar_id, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sA = smem.allocate_tensor(element_type=self.a_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sB = smem.allocate_tensor(element_type=self.b_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV_ptr = cute.recast_ptr(sB.iterator, v_smem_s.inner) - sV = cute.make_tensor(sV_ptr, v_smem_s.outer) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gA = cute.local_tile(mA, cute.slice_(self.mma_tiler, (None,0,None)), (None,None,None)) - gB = cute.local_tile(mB, cute.slice_(self.mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gA, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - tCgA = qk_thr.partition_A(gA); tCgB = qk_thr.partition_B(gB); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsA, tAgA = cpasync.tma_partition(tma_a, 0, a_lay, cute.group_modes(sA,0,3), cute.group_modes(tCgA,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsB, tBgB = cpasync.tma_partition(tma_b, 0, b_lay, cute.group_modes(sB,0,3), cute.group_modes(tCgB,0,3)) - tAgA = tAgA[(None,0,None,0)]; tBgB = tBgB[(None,0,None,0)] - - tCrA = qk_mma.make_fragment_A(sA); tCrB = qk_mma.make_fragment_B(sB) - tCrV = pv_mma.make_fragment_B(sV) - - # TMEM tensors - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - # P A-fragment (backward FMHA pattern) - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tdVrP_base = pv_thr.make_fragment_A(tP) - tdVrP = tdVrP_base[(None, None, None, 0)] - tdVrP_iter = cute.recast_ptr(tStS.iterator, dtype=self.q_dtype) - tdVrP = cute.make_tensor(tdVrP_iter, tdVrP.layout) - tdVrP0 = cute.make_tensor( - tdVrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tdVrP.layout) - - print(f'[DIAG] tStS.layout: {tStS.layout}') - print(f'[DIAG] tdVrP.layout: {tdVrP.layout}') - print(f'[DIAG] tdVrP0.layout: {tdVrP0.layout}') - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # TMA WARP - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_a, tAgA[(None,h.count)], tAsA[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_b, tBgB[(None,h.count)], tBsB[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1 BF16 (backward FMHA pattern) - tRT_rST_bf16 = cute.make_rmem_tensor(tTR_rST.shape, self.q_dtype) - frg_cnt = 4 - frg_tile = cute.size(tTR_rST) // frg_cnt - tTR_rST_frg = cute.logical_divide(tTR_rST, cute.make_layout(frg_tile)) - tRT_rST_bf16_frg = cute.make_tensor(tRT_rST_bf16.iterator, tTR_rST_frg.layout) - for j in range(frg_cnt): - frg_vec = tTR_rST_frg[None, j].load() - tRT_rST_bf16_frg[None, j].store(frg_vec.to(self.q_dtype)) - - # 5. Reshape to match store coordinate shape - tRT_rST_reshaped = cute.make_tensor( - tRT_rST_bf16.iterator, cute.make_layout(tRT_cS.shape)) - - # 6. STORE to A-fragment TMEM - cute.copy(tiled_tmem_store, tRT_rST_reshaped, tRT_tP) - cute.arch.fence_view_async_tmem_store() - - si_handle.release() - - # Epilogue - tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout) - acc_cons_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage) - c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) - c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp) - acc_cons_st = utils.gemm.sm100.epilogue_tma_store( - self, tidx, warp_idx, tma_c, tCtO_base, sC, tCgC, - epi_tile, 0, const_expr(lambda x: x), (0,0,0), acc_cons_st, acc_pipe, c_pipe) - c_pipe.producer_tail() - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def test(): - torch.manual_seed(42) - m, n, k = 128, 128, 128 - q = torch.randn(m, k, 1, dtype=torch.bfloat16, device='cuda') - kv = torch.randn(n, k, 1, dtype=torch.bfloat16, device='cuda') - c = torch.zeros(m, n, 1, dtype=torch.bfloat16, device='cuda') - qf = q[:,:,0].float(); kvf = kv[:,:,0].float() - ref = qf @ kvf.T @ kvf - import cutlass.torch as ct - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(kv).mark_layout_dynamic(leading_dim=ct.get_leading_dim(kv)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - kernel = StageBIdentitySoftmax(mma_tiler_mn=(128, 128), use_2cta_instrs=False, use_tma_store=True) - print('Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mC, stream) - print('Running...', flush=True) - compiled(mQ, mK, mC, stream) - torch.cuda.synchronize() - out = c[:,:,0].float() - cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - max_err = (out - ref).abs().max().item() - print('Stage B v14: backward FMHA A-fragment store pattern (identity softmax)') - print(' Cosine: {:.6f}, Max error: {:.6f}'.format(cos, max_err)) - print(' {}'.format('PASS' if cos >= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_stage_b_v16.py b/tests/archive/test_stage_b_v16.py deleted file mode 100644 index 6b3a95be..00000000 --- a/tests/archive/test_stage_b_v16.py +++ /dev/null @@ -1,453 +0,0 @@ -""" -Stage B v7: Two MMAs + Identity Softmax with COMPUTED TMEM offsets. - -Key fixes over v6: - - TMEM offsets computed via find_tmem_tensor_col_offset (same API as get_num_tmem_alloc_cols) - - P tensor constructed from p_tmem_s.outer (matching fmha.py pattern exactly) - - tilePlikeFP32 computed from qk_mma_tiler and dtype widths - - tmem_alloc_cols from get_num_tmem_alloc_cols with all fragments - - JIT-time diagnostic prints of all TMEM sizes - -Architecture (matches fmha.py exactly): - MMA1: Q @ K^T → tmem_scores (a_source=SMEM, accumulate=False) - Identity softmax: tcgen05.ld C-layout → F32→BF16 → tcgen05.st A-layout - MMA2: P @ V → tmem_output (a_source=TMEM, accumulate=True) -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - - -class StageBIdentitySoftmax: - def __init__(self, mma_tiler_mn, use_2cta_instrs=False, use_tma_store=True): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16 - self.use_2cta_instrs = use_2cta_instrs; self.use_tma_store = use_tma_store - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.epilog_sync_bar_id = 1; self.tmem_alloc_sync_bar_id = 2; self.tmem_dealloc_sync_bar_id = 3 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - pv_inst_k = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (self.qk_mma_tiler[0], self.qk_mma_tiler[2], self.qk_mma_tiler[1]) - self.mma_tiler = self.qk_mma_tiler - print(f"[StageB] qk_mma.shape_mnk = {qk_mma.shape_mnk}") - print(f"[StageB] pv_mma.shape_mnk = {pv_mma.shape_mnk}") - print(f"[StageB] qk_mma_tiler = {self.qk_mma_tiler}") - print(f"[StageB] pv_mma_tiler = {self.pv_mma_tiler}") - print(f"[StageB] qk_inst_k = {qk_inst_k}, pv_inst_k = {pv_inst_k}") - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], - self.qk_mma_tiler[2], - ) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, self.use_2cta_instrs, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.a_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.b_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.b_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - # ── COMPUTE TMEM OFFSETS USING find_tmem_tensor_col_offset ── - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - # tilePlikeFP32 for the store-side composition - self.tilePlikeFP32 = self.qk_mma_tiler[1] * self.q_dtype.width // 32 - - # ── TMEM LAYOUT (matching fmha.py) ── - # P OVERLAPS S — after softmax, P (A-layout) is written on top of scores (C-layout) - # in the same TMEM region. The A-layout view starts partway through the S region. - # fmha.py: S0=0, P0=32, O0=256 (with S1=128, P1=160 double-buffered) - # The P offset of 32 means the A-layout starts at column 32 within the S region. - # This is because the C-layout and A-layout partition TMEM differently per-thread; - # the first 32 C-layout columns are "dead space" in the A-layout mapping. - # - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 # Original - self.tmem_o0_offset = s_cols # 128 - self.tmem_alloc_cols = 512 # FMHA: allocate max TMEM - - # Also compute via get_num_tmem_alloc_cols for the full allocation - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, 1)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, 1)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - print(f"[StageB] s_cols (QK accumulator) = {s_cols}") - print(f"[StageB] o_cols (PV accumulator) = {o_cols}") - print(f"[StageB] tilePlikeFP32 = {self.tilePlikeFP32}") - print(f"[StageB] tmem_s0_offset = {self.tmem_s0_offset}") - print(f"[StageB] tmem_p0_offset = {self.tmem_p0_offset}") - print(f"[StageB] tmem_o0_offset = {self.tmem_o0_offset}") - print(f"[StageB] tmem_alloc_cols (computed) = {self.tmem_alloc_cols}") - print(f"[StageB] num_tmem_alloc_cols (via utils) = {self.num_tmem_alloc_cols}") - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.a_dtype, a_smem) + cute.size_in_bytes(self.b_dtype, b_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - self.a_dtype = a.element_type; self.b_dtype = b.element_type; self.c_dtype = c.element_type - self.a_major = LayoutEnum.from_tensor(a).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(b).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, cute.nvgpu.OperandMajorMode.K, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - # Introspect PV MMA atom - print(f"[ATOM] PV MMA type: {type(pv_mma)}") - print(f"[ATOM] PV MMA op: {pv_mma.op if hasattr(pv_mma, "op") else "no op"}") - print(f"[ATOM] PV MMA trait: {pv_mma._trait if hasattr(pv_mma, "_trait") else "no trait"}") - print(f"[ATOM] PV MMA shape_mnk: {pv_mma.shape_mnk}") - print(f"[ATOM] QK MMA shape_mnk: {qk_mma.shape_mnk}") - # Check a_src - print(f"[ATOM] PV MMA op.a_src: {pv_mma.op.a_src}") - print(f"[ATOM] QK MMA op.a_src: {qk_mma.op.a_src}") - print(f"[ATOM] PV MMA op: {pv_mma.op}") - self._setup(qk_mma, pv_mma) - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - tma_a, tma_ta = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - a, a_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_b, tma_tb = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - b, b_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_a, tma_ta, tma_b, tma_tb, tma_c, tma_tc, - self.cluster_layout_vmnk, self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_a, mA, tma_b, mB, tma_c, mC, cl_vmnk, - a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_a); cpasync.prefetch_descriptor(tma_b); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=self.tmem_alloc_sync_bar_id, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sA = smem.allocate_tensor(element_type=self.a_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sB = smem.allocate_tensor(element_type=self.b_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - # V shares the same SMEM as B (same data, different layout for PV MMA) - sV_ptr = cute.recast_ptr(sB.iterator, v_smem_s.inner) - sV = cute.make_tensor(sV_ptr, v_smem_s.outer) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gA = cute.local_tile(mA, cute.slice_(self.mma_tiler, (None,0,None)), (None,None,None)) - gB = cute.local_tile(mB, cute.slice_(self.mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gA, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - tCgA = qk_thr.partition_A(gA); tCgB = qk_thr.partition_B(gB); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsA, tAgA = cpasync.tma_partition(tma_a, 0, a_lay, cute.group_modes(sA,0,3), cute.group_modes(tCgA,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsB, tBgB = cpasync.tma_partition(tma_b, 0, b_lay, cute.group_modes(sB,0,3), cute.group_modes(tCgB,0,3)) - tAgA = tAgA[(None,0,None,0)]; tBgB = tBgB[(None,0,None,0)] - - tCrA = qk_mma.make_fragment_A(sA); tCrB = qk_mma.make_fragment_B(sB) - tCrV = pv_mma.make_fragment_B(sV) # V fragment from V SMEM layout - print(f"[DIAG] tCrV.size = {cute.size(tCrV)}") - print(f"[DIAG] tCrA.size = {cute.size(tCrA)}") - print(f"[DIAG] tCrB.size = {cute.size(tCrB)}") - print(f"[DIAG] nblk_qk (tCrA mode 2) = {cute.size(tCrA, mode=[2])}") - - # ── TMEM tensors with computed offsets (matching fmha.py pattern) ── - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - # P fragment: construct from p_tmem_s layout (matching fmha.py exactly) - # fmha.py: tP = cute.make_tensor(tStS.iterator, p_tmem_layout_staged.outer) - # tOrP = pv_thr_mma.make_fragment_A(tP)[None, None, None, 0] - # tdVrP0 = cute.make_tensor(tdVrP.iterator + dtype_width_ratio * tmem_p0_offset, tdVrP.layout) - print(f'[TMEM] p_tmem_s: {p_tmem_s}') - print(f'[TMEM] p_tmem_s.outer: {p_tmem_s.outer}') - print(f'[TMEM] p_tmem_s.inner: {p_tmem_s.inner}') - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - print(f'[DIAG] tStS.layout: {tStS.layout}') - print(f'[DIAG] tStS.size: {cute.size(tStS)}') - print(f'[DIAG] p_tmem_s.outer: {p_tmem_s.outer}') - print(f'[DIAG] p_tmem_s.inner: {p_tmem_s.inner}') - tdVrP_base = pv_thr.make_fragment_A(tP) - tdVrP = tdVrP_base[(None, None, None, 0)] - tdVrP_iter = cute.recast_ptr(tStS.iterator, dtype=self.q_dtype) - tdVrP = cute.make_tensor(tdVrP_iter, tdVrP.layout) - tdVrP0 = cute.make_tensor( - tdVrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tdVrP.layout) - - # Compute nblk_pv for diagnostics - nblk_pv = cute.size(tdVrP0, mode=[2]) - nblk_qk = cute.size(tCrA, mode=[2]) - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - # COMPREHENSIVE LAYOUT DUMP - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tilePlikeFP32 = self.qk_mma_tiler[1] * self.q_dtype.width // 32 - tStS_P_layout = cute.composition(tStS.layout, cute.make_layout((128, tilePlikeFP32))) - tStS_P = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStS_P_layout) - tScS_P_layout = cute.composition(tScS.layout, cute.make_layout((128, tilePlikeFP32))) - tScS_P = cute.make_tensor(tScS.iterator, tScS_P_layout) - - print(f'[LAYOUT] QK C-fragment tStS.layout: {tStS.layout}') - print(f'[LAYOUT] QK C-fragment tStS cosize: {cute.cosize(tStS.layout)}') - print(f'[LAYOUT] QK C-fragment tStS.size: {cute.size(tStS)}') - print(f'[LAYOUT] QK C-fragment tScS.layout: {tScS.layout}') - print(f'[LAYOUT] QK C-fragment tScS cosize: {cute.cosize(tScS.layout)}') - print(f'[LAYOUT] PV A-fragment tdVrP.layout: {tdVrP.layout}') - print(f'[LAYOUT] PV A-fragment tdVrP cosize: {cute.cosize(tdVrP.layout)}') - print(f'[LAYOUT] PV A-fragment tdVrP.size: {cute.size(tdVrP)}') - print(f'[LAYOUT] PV A-fragment tdVrP0.layout: {tdVrP0.layout}') - print(f'[LAYOUT] PV A-fragment tdVrP0 cosize: {cute.cosize(tdVrP0.layout)}') - print(f'[LAYOUT] tP.layout: {tP.layout}') - print(f'[LAYOUT] tP cosize: {cute.cosize(tP.layout)}') - print(f'[LAYOUT] tStS_P (composed) layout: {tStS_P.layout}') - print(f'[LAYOUT] tStS_P (composed) cosize: {cute.cosize(tStS_P.layout)}') - print(f'[LAYOUT] tScS_P (composed) layout: {tScS_P.layout}') - print(f'[LAYOUT] tScS_P (composed) cosize: {cute.cosize(tScS_P.layout)}') - print(f'[LAYOUT] tOtO.layout: {tOtO.layout}') - print(f'[LAYOUT] tOtO cosize: {cute.cosize(tOtO.layout)}') - print(f'[LAYOUT] pv_mma_tiler: {self.pv_mma_tiler}') - print(f'[LAYOUT] qk_mma_tiler: {self.qk_mma_tiler}') - print(f'[LAYOUT] tilePlikeFP32: {tilePlikeFP32}') - - # DIAGNOSTIC: Compare tP (A-layout) vs tStS_P (composition) - tilePlikeFP32 = self.qk_mma_tiler[1] * self.q_dtype.width // 32 - tStS_P_layout = cute.composition(tStS.layout, cute.make_layout((128, tilePlikeFP32))) - tStS_P = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStS_P_layout) - print(f'[DIAG] tP.layout: {tP.layout}') - print(f'[DIAG] tP.size: {cute.size(tP)}') - print(f'[DIAG] tP.element_type: {tP.element_type if hasattr(tP, 'element_type') else 'N/A'}') - print(f'[DIAG] tStS_P.layout: {tStS_P.layout}') - print(f'[DIAG] tStS_P.size: {cute.size(tStS_P)}') - print(f'[DIAG] tStS_P.element_type: {tStS_P.element_type if hasattr(tStS_P, 'element_type') else 'N/A'}') - print(f'[DIAG] tilePlikeFP32: {tilePlikeFP32}') - print(f'[DIAG] tP and tStS_P same iterator? {tP.iterator == tStS_P.iterator if hasattr(tP, 'iterator') else 'cant compare'}') - - print(f'[DIAG] nblk_pv = {nblk_pv}, nblk_qk = {nblk_qk}') - print(f'[DIAG] tCrV.size = {cute.size(tCrV)}') - print(f'[DIAG] tdVrP0.size = {cute.size(tdVrP0)}') - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ── TMA WARP ── - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_a, tAgA[(None,h.count)], tAsA[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_b, tBgB[(None,h.count)], tBsB[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1 BF16 (backward FMHA pattern) - tRT_rST_bf16 = cute.make_rmem_tensor(tTMEM_LOADrS.shape, self.q_dtype) - frg_cnt = 4 - frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt - tTR_rST_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile)) - tRT_rST_bf16_frg = cute.make_tensor(tRT_rST_bf16.iterator, tTR_rST_frg.layout) - for j in range(frg_cnt): - frg_vec = tTR_rST_frg[None, j].load() - tRT_rST_bf16_frg[None, j].store(frg_vec.to(self.q_dtype)) - # 6. Reshape and store to A-fragment TMEM - tRT_rST_reshaped = cute.make_tensor( - tRT_rST_bf16.iterator, cute.make_layout(tRT_cS.shape)) - cute.copy(tiled_tmem_store, tRT_rST_reshaped, tRT_tP) - cute.arch.fence_view_async_tmem_store() - - # 7. Release back to MMA warp - si_handle.release() - - # ── Epilogue ── - tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout) - acc_cons_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage) - c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) - c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp) - acc_cons_st = utils.gemm.sm100.epilogue_tma_store( - self, tidx, warp_idx, tma_c, tCtO_base, sC, tCgC, - epi_tile, 0, const_expr(lambda x: x), (0,0,0), acc_cons_st, acc_pipe, c_pipe) - c_pipe.producer_tail() - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def test(): - torch.manual_seed(42) - m, n, k = 128, 128, 128 - q = torch.randn(m, k, 1, dtype=torch.bfloat16, device='cuda') - kv = torch.randn(n, k, 1, dtype=torch.bfloat16, device='cuda') - c = torch.zeros(m, n, 1, dtype=torch.bfloat16, device='cuda') - qf = q[:,:,0].float(); kvf = kv[:,:,0].float() - ref = qf @ kvf.T @ kvf - import cutlass.torch as ct - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(kv).mark_layout_dynamic(leading_dim=ct.get_leading_dim(kv)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - kernel = StageBIdentitySoftmax(mma_tiler_mn=(128, 128), use_2cta_instrs=False, use_tma_store=True) - print('Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mC, stream) - print('Running...', flush=True) - compiled(mQ, mK, mC, stream) - torch.cuda.synchronize() - out = c[:,:,0].float() - cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - max_err = (out - ref).abs().max().item() - print('Stage B v7: (Q @ K^T) @ V with identity softmax (computed TMEM offsets)') - print(' Cosine: {:.6f}, Max error: {:.6f}'.format(cos, max_err)) - print(' {}'.format('PASS' if cos >= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_stage_b_v17.py b/tests/archive/test_stage_b_v17.py deleted file mode 100644 index 6bff0212..00000000 --- a/tests/archive/test_stage_b_v17.py +++ /dev/null @@ -1,450 +0,0 @@ -""" -Stage B v7: Two MMAs + Identity Softmax with COMPUTED TMEM offsets. - -Key fixes over v6: - - TMEM offsets computed via find_tmem_tensor_col_offset (same API as get_num_tmem_alloc_cols) - - P tensor constructed from p_tmem_s.outer (matching fmha.py pattern exactly) - - tilePlikeFP32 computed from qk_mma_tiler and dtype widths - - tmem_alloc_cols from get_num_tmem_alloc_cols with all fragments - - JIT-time diagnostic prints of all TMEM sizes - -Architecture (matches fmha.py exactly): - MMA1: Q @ K^T → tmem_scores (a_source=SMEM, accumulate=False) - Identity softmax: tcgen05.ld C-layout → F32→BF16 → tcgen05.st A-layout - MMA2: P @ V → tmem_output (a_source=TMEM, accumulate=True) -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - - -class StageBIdentitySoftmax: - def __init__(self, mma_tiler_mn, use_2cta_instrs=False, use_tma_store=True): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16 - self.use_2cta_instrs = use_2cta_instrs; self.use_tma_store = use_tma_store - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.epilog_sync_bar_id = 1; self.tmem_alloc_sync_bar_id = 2; self.tmem_dealloc_sync_bar_id = 3 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - pv_inst_k = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (self.qk_mma_tiler[0], self.qk_mma_tiler[2], self.qk_mma_tiler[1]) - self.mma_tiler = self.qk_mma_tiler - print(f"[StageB] qk_mma.shape_mnk = {qk_mma.shape_mnk}") - print(f"[StageB] pv_mma.shape_mnk = {pv_mma.shape_mnk}") - print(f"[StageB] qk_mma_tiler = {self.qk_mma_tiler}") - print(f"[StageB] pv_mma_tiler = {self.pv_mma_tiler}") - print(f"[StageB] qk_inst_k = {qk_inst_k}, pv_inst_k = {pv_inst_k}") - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], - self.qk_mma_tiler[2], - ) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, self.use_2cta_instrs, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.a_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.b_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.b_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - # ── COMPUTE TMEM OFFSETS USING find_tmem_tensor_col_offset ── - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - # tilePlikeFP32 for the store-side composition - self.tilePlikeFP32 = self.qk_mma_tiler[1] * self.q_dtype.width // 32 - - # ── TMEM LAYOUT (matching fmha.py) ── - # P OVERLAPS S — after softmax, P (A-layout) is written on top of scores (C-layout) - # in the same TMEM region. The A-layout view starts partway through the S region. - # fmha.py: S0=0, P0=32, O0=256 (with S1=128, P1=160 double-buffered) - # The P offset of 32 means the A-layout starts at column 32 within the S region. - # This is because the C-layout and A-layout partition TMEM differently per-thread; - # the first 32 C-layout columns are "dead space" in the A-layout mapping. - # - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 # Original - self.tmem_o0_offset = s_cols # 128 - self.tmem_alloc_cols = s_cols + o_cols # 256 - - # Also compute via get_num_tmem_alloc_cols for the full allocation - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, 1)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, 1)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - print(f"[StageB] s_cols (QK accumulator) = {s_cols}") - print(f"[StageB] o_cols (PV accumulator) = {o_cols}") - print(f"[StageB] tilePlikeFP32 = {self.tilePlikeFP32}") - print(f"[StageB] tmem_s0_offset = {self.tmem_s0_offset}") - print(f"[StageB] tmem_p0_offset = {self.tmem_p0_offset}") - print(f"[StageB] tmem_o0_offset = {self.tmem_o0_offset}") - print(f"[StageB] tmem_alloc_cols (computed) = {self.tmem_alloc_cols}") - print(f"[StageB] num_tmem_alloc_cols (via utils) = {self.num_tmem_alloc_cols}") - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.a_dtype, a_smem) + cute.size_in_bytes(self.b_dtype, b_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - self.a_dtype = a.element_type; self.b_dtype = b.element_type; self.c_dtype = c.element_type - self.a_major = LayoutEnum.from_tensor(a).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(b).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, cute.nvgpu.OperandMajorMode.K, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - # Introspect PV MMA atom - print(f"[ATOM] PV MMA type: {type(pv_mma)}") - print(f"[ATOM] PV MMA op: {pv_mma.op if hasattr(pv_mma, "op") else "no op"}") - print(f"[ATOM] PV MMA trait: {pv_mma._trait if hasattr(pv_mma, "_trait") else "no trait"}") - print(f"[ATOM] PV MMA shape_mnk: {pv_mma.shape_mnk}") - print(f"[ATOM] QK MMA shape_mnk: {qk_mma.shape_mnk}") - # Check a_src - print(f"[ATOM] PV MMA op.a_src: {pv_mma.op.a_src}") - print(f"[ATOM] QK MMA op.a_src: {qk_mma.op.a_src}") - print(f"[ATOM] PV MMA op: {pv_mma.op}") - self._setup(qk_mma, pv_mma) - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - tma_a, tma_ta = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - a, a_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_b, tma_tb = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - b, b_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_a, tma_ta, tma_b, tma_tb, tma_c, tma_tc, - self.cluster_layout_vmnk, self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_a, mA, tma_b, mB, tma_c, mC, cl_vmnk, - a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_a); cpasync.prefetch_descriptor(tma_b); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=self.tmem_alloc_sync_bar_id, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sA = smem.allocate_tensor(element_type=self.a_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sB = smem.allocate_tensor(element_type=self.b_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - # V shares the same SMEM as B (same data, different layout for PV MMA) - sV_ptr = cute.recast_ptr(sB.iterator, v_smem_s.inner) - sV = cute.make_tensor(sV_ptr, v_smem_s.outer) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gA = cute.local_tile(mA, cute.slice_(self.mma_tiler, (None,0,None)), (None,None,None)) - gB = cute.local_tile(mB, cute.slice_(self.mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gA, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - tCgA = qk_thr.partition_A(gA); tCgB = qk_thr.partition_B(gB); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsA, tAgA = cpasync.tma_partition(tma_a, 0, a_lay, cute.group_modes(sA,0,3), cute.group_modes(tCgA,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsB, tBgB = cpasync.tma_partition(tma_b, 0, b_lay, cute.group_modes(sB,0,3), cute.group_modes(tCgB,0,3)) - tAgA = tAgA[(None,0,None,0)]; tBgB = tBgB[(None,0,None,0)] - - tCrA = qk_mma.make_fragment_A(sA); tCrB = qk_mma.make_fragment_B(sB) - tCrV = pv_mma.make_fragment_B(sV) # V fragment from V SMEM layout - print(f"[DIAG] tCrV.size = {cute.size(tCrV)}") - print(f"[DIAG] tCrA.size = {cute.size(tCrA)}") - print(f"[DIAG] tCrB.size = {cute.size(tCrB)}") - print(f"[DIAG] nblk_qk (tCrA mode 2) = {cute.size(tCrA, mode=[2])}") - - # ── TMEM tensors with computed offsets (matching fmha.py pattern) ── - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - # P fragment: construct from p_tmem_s layout (matching fmha.py exactly) - # fmha.py: tP = cute.make_tensor(tStS.iterator, p_tmem_layout_staged.outer) - # tOrP = pv_thr_mma.make_fragment_A(tP)[None, None, None, 0] - # tOrP0 = cute.make_tensor(tOrP.iterator + dtype_width_ratio * tmem_p0_offset, tOrP.layout) - print(f'[TMEM] p_tmem_s: {p_tmem_s}') - print(f'[TMEM] p_tmem_s.outer: {p_tmem_s.outer}') - print(f'[TMEM] p_tmem_s.inner: {p_tmem_s.inner}') - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - print(f'[DIAG] tStS.layout: {tStS.layout}') - print(f'[DIAG] tStS.size: {cute.size(tStS)}') - print(f'[DIAG] p_tmem_s.outer: {p_tmem_s.outer}') - print(f'[DIAG] p_tmem_s.inner: {p_tmem_s.inner}') - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - # Compute nblk_pv for diagnostics - nblk_pv = cute.size(tOrP0, mode=[2]) - nblk_qk = cute.size(tCrA, mode=[2]) - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - # COMPREHENSIVE LAYOUT DUMP - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tilePlikeFP32 = self.qk_mma_tiler[1] * self.q_dtype.width // 32 - tStS_P_layout = cute.composition(tStS.layout, cute.make_layout((128, tilePlikeFP32))) - tStS_P = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStS_P_layout) - tScS_P_layout = cute.composition(tScS.layout, cute.make_layout((128, tilePlikeFP32))) - tScS_P = cute.make_tensor(tScS.iterator, tScS_P_layout) - - print(f'[LAYOUT] QK C-fragment tStS.layout: {tStS.layout}') - print(f'[LAYOUT] QK C-fragment tStS cosize: {cute.cosize(tStS.layout)}') - print(f'[LAYOUT] QK C-fragment tStS.size: {cute.size(tStS)}') - print(f'[LAYOUT] QK C-fragment tScS.layout: {tScS.layout}') - print(f'[LAYOUT] QK C-fragment tScS cosize: {cute.cosize(tScS.layout)}') - print(f'[LAYOUT] PV A-fragment tOrP.layout: {tOrP.layout}') - print(f'[LAYOUT] PV A-fragment tOrP cosize: {cute.cosize(tOrP.layout)}') - print(f'[LAYOUT] PV A-fragment tOrP.size: {cute.size(tOrP)}') - print(f'[LAYOUT] PV A-fragment tOrP0.layout: {tOrP0.layout}') - print(f'[LAYOUT] PV A-fragment tOrP0 cosize: {cute.cosize(tOrP0.layout)}') - print(f'[LAYOUT] tP.layout: {tP.layout}') - print(f'[LAYOUT] tP cosize: {cute.cosize(tP.layout)}') - print(f'[LAYOUT] tStS_P (composed) layout: {tStS_P.layout}') - print(f'[LAYOUT] tStS_P (composed) cosize: {cute.cosize(tStS_P.layout)}') - print(f'[LAYOUT] tScS_P (composed) layout: {tScS_P.layout}') - print(f'[LAYOUT] tScS_P (composed) cosize: {cute.cosize(tScS_P.layout)}') - print(f'[LAYOUT] tOtO.layout: {tOtO.layout}') - print(f'[LAYOUT] tOtO cosize: {cute.cosize(tOtO.layout)}') - print(f'[LAYOUT] pv_mma_tiler: {self.pv_mma_tiler}') - print(f'[LAYOUT] qk_mma_tiler: {self.qk_mma_tiler}') - print(f'[LAYOUT] tilePlikeFP32: {tilePlikeFP32}') - - # DIAGNOSTIC: Compare tP (A-layout) vs tStS_P (composition) - tilePlikeFP32 = self.qk_mma_tiler[1] * self.q_dtype.width // 32 - tStS_P_layout = cute.composition(tStS.layout, cute.make_layout((128, tilePlikeFP32))) - tStS_P = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStS_P_layout) - print(f'[DIAG] tP.layout: {tP.layout}') - print(f'[DIAG] tP.size: {cute.size(tP)}') - print(f'[DIAG] tP.element_type: {tP.element_type if hasattr(tP, 'element_type') else 'N/A'}') - print(f'[DIAG] tStS_P.layout: {tStS_P.layout}') - print(f'[DIAG] tStS_P.size: {cute.size(tStS_P)}') - print(f'[DIAG] tStS_P.element_type: {tStS_P.element_type if hasattr(tStS_P, 'element_type') else 'N/A'}') - print(f'[DIAG] tilePlikeFP32: {tilePlikeFP32}') - print(f'[DIAG] tP and tStS_P same iterator? {tP.iterator == tStS_P.iterator if hasattr(tP, 'iterator') else 'cant compare'}') - - print(f'[DIAG] nblk_pv = {nblk_pv}, nblk_qk = {nblk_qk}') - print(f'[DIAG] tCrV.size = {cute.size(tCrV)}') - print(f'[DIAG] tOrP0.size = {cute.size(tOrP0)}') - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ── TMA WARP ── - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_a, tAgA[(None,h.count)], tAsA[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_b, tBgB[(None,h.count)], tBsB[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_stage_b_v18.py b/tests/archive/test_stage_b_v18.py deleted file mode 100644 index faffa755..00000000 --- a/tests/archive/test_stage_b_v18.py +++ /dev/null @@ -1,452 +0,0 @@ -""" -Stage B v7: Two MMAs + Identity Softmax with COMPUTED TMEM offsets. - -Key fixes over v6: - - TMEM offsets computed via find_tmem_tensor_col_offset (same API as get_num_tmem_alloc_cols) - - P tensor constructed from p_tmem_s.outer (matching fmha.py pattern exactly) - - tilePlikeFP32 computed from qk_mma_tiler and dtype widths - - tmem_alloc_cols from get_num_tmem_alloc_cols with all fragments - - JIT-time diagnostic prints of all TMEM sizes - -Architecture (matches fmha.py exactly): - MMA1: Q @ K^T → tmem_scores (a_source=SMEM, accumulate=False) - Identity softmax: tcgen05.ld C-layout → F32→BF16 → tcgen05.st A-layout - MMA2: P @ V → tmem_output (a_source=TMEM, accumulate=True) -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - - -class StageBIdentitySoftmax: - def __init__(self, mma_tiler_mn, use_2cta_instrs=False, use_tma_store=True): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16 - self.use_2cta_instrs = use_2cta_instrs; self.use_tma_store = use_tma_store - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.epilog_sync_bar_id = 1; self.tmem_alloc_sync_bar_id = 2; self.tmem_dealloc_sync_bar_id = 3 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - pv_inst_k = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (self.qk_mma_tiler[0], self.qk_mma_tiler[2], self.qk_mma_tiler[1]) - self.mma_tiler = self.qk_mma_tiler - print(f"[StageB] qk_mma.shape_mnk = {qk_mma.shape_mnk}") - print(f"[StageB] pv_mma.shape_mnk = {pv_mma.shape_mnk}") - print(f"[StageB] qk_mma_tiler = {self.qk_mma_tiler}") - print(f"[StageB] pv_mma_tiler = {self.pv_mma_tiler}") - print(f"[StageB] qk_inst_k = {qk_inst_k}, pv_inst_k = {pv_inst_k}") - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], - self.qk_mma_tiler[2], - ) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, self.use_2cta_instrs, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.a_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.b_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.b_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - # ── COMPUTE TMEM OFFSETS USING find_tmem_tensor_col_offset ── - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - # tilePlikeFP32 for the store-side composition - self.tilePlikeFP32 = self.qk_mma_tiler[1] * self.q_dtype.width // 32 - - # ── TMEM LAYOUT (matching fmha.py) ── - # P OVERLAPS S — after softmax, P (A-layout) is written on top of scores (C-layout) - # in the same TMEM region. The A-layout view starts partway through the S region. - # fmha.py: S0=0, P0=32, O0=256 (with S1=128, P1=160 double-buffered) - # The P offset of 32 means the A-layout starts at column 32 within the S region. - # This is because the C-layout and A-layout partition TMEM differently per-thread; - # the first 32 C-layout columns are "dead space" in the A-layout mapping. - # - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 # Original - self.tmem_o0_offset = s_cols # 128 - self.tmem_alloc_cols = 512 # FMHA-style: allocate max TMEM # 256 - - # Also compute via get_num_tmem_alloc_cols for the full allocation - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, 1)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, 1)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - print(f"[StageB] s_cols (QK accumulator) = {s_cols}") - print(f"[StageB] o_cols (PV accumulator) = {o_cols}") - print(f"[StageB] tilePlikeFP32 = {self.tilePlikeFP32}") - print(f"[StageB] tmem_s0_offset = {self.tmem_s0_offset}") - print(f"[StageB] tmem_p0_offset = {self.tmem_p0_offset}") - print(f"[StageB] tmem_o0_offset = {self.tmem_o0_offset}") - print(f"[StageB] tmem_alloc_cols (computed) = {self.tmem_alloc_cols}") - print(f"[StageB] num_tmem_alloc_cols (via utils) = {self.num_tmem_alloc_cols}") - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.a_dtype, a_smem) + cute.size_in_bytes(self.b_dtype, b_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - self.a_dtype = a.element_type; self.b_dtype = b.element_type; self.c_dtype = c.element_type - self.a_major = LayoutEnum.from_tensor(a).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(b).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, cute.nvgpu.OperandMajorMode.K, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - # Introspect PV MMA atom - print(f"[ATOM] PV MMA type: {type(pv_mma)}") - print(f"[ATOM] PV MMA op: {pv_mma.op if hasattr(pv_mma, "op") else "no op"}") - print(f"[ATOM] PV MMA trait: {pv_mma._trait if hasattr(pv_mma, "_trait") else "no trait"}") - print(f"[ATOM] PV MMA shape_mnk: {pv_mma.shape_mnk}") - print(f"[ATOM] QK MMA shape_mnk: {qk_mma.shape_mnk}") - # Check a_src - print(f"[ATOM] PV MMA op.a_src: {pv_mma.op.a_src}") - print(f"[ATOM] QK MMA op.a_src: {qk_mma.op.a_src}") - print(f"[ATOM] PV MMA op: {pv_mma.op}") - self._setup(qk_mma, pv_mma) - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - tma_a, tma_ta = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - a, a_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_b, tma_tb = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - b, b_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_a, tma_ta, tma_b, tma_tb, tma_c, tma_tc, - self.cluster_layout_vmnk, self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_a, mA, tma_b, mB, tma_c, mC, cl_vmnk, - a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_a); cpasync.prefetch_descriptor(tma_b); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=self.tmem_alloc_sync_bar_id, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sA = smem.allocate_tensor(element_type=self.a_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sB = smem.allocate_tensor(element_type=self.b_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - # V shares the same SMEM as B (same data, different layout for PV MMA) - sV_ptr = cute.recast_ptr(sB.iterator, v_smem_s.inner) - sV = cute.make_tensor(sV_ptr, v_smem_s.outer) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gA = cute.local_tile(mA, cute.slice_(self.mma_tiler, (None,0,None)), (None,None,None)) - gB = cute.local_tile(mB, cute.slice_(self.mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gA, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - tCgA = qk_thr.partition_A(gA); tCgB = qk_thr.partition_B(gB); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsA, tAgA = cpasync.tma_partition(tma_a, 0, a_lay, cute.group_modes(sA,0,3), cute.group_modes(tCgA,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsB, tBgB = cpasync.tma_partition(tma_b, 0, b_lay, cute.group_modes(sB,0,3), cute.group_modes(tCgB,0,3)) - tAgA = tAgA[(None,0,None,0)]; tBgB = tBgB[(None,0,None,0)] - - tCrA = qk_mma.make_fragment_A(sA); tCrB = qk_mma.make_fragment_B(sB) - tCrV = pv_mma.make_fragment_B(sV) # V fragment from V SMEM layout - print(f"[DIAG] tCrV.size = {cute.size(tCrV)}") - print(f"[DIAG] tCrA.size = {cute.size(tCrA)}") - print(f"[DIAG] tCrB.size = {cute.size(tCrB)}") - print(f"[DIAG] nblk_qk (tCrA mode 2) = {cute.size(tCrA, mode=[2])}") - - # ── TMEM tensors with computed offsets (matching fmha.py pattern) ── - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - # P fragment: construct from p_tmem_s layout (matching fmha.py exactly) - # fmha.py: tP = cute.make_tensor(tStS.iterator, p_tmem_layout_staged.outer) - # tOrP = pv_thr_mma.make_fragment_A(tP)[None, None, None, 0] - # tOrP0 = cute.make_tensor(tOrP.iterator + dtype_width_ratio * tmem_p0_offset, tOrP.layout) - print(f'[TMEM] p_tmem_s: {p_tmem_s}') - print(f'[TMEM] p_tmem_s.outer: {p_tmem_s.outer}') - print(f'[TMEM] p_tmem_s.inner: {p_tmem_s.inner}') - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - print(f'[DIAG] tStS.layout: {tStS.layout}') - print(f'[DIAG] tStS.size: {cute.size(tStS)}') - print(f'[DIAG] p_tmem_s.outer: {p_tmem_s.outer}') - print(f'[DIAG] p_tmem_s.inner: {p_tmem_s.inner}') - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - # Compute nblk_pv for diagnostics - nblk_pv = cute.size(tOrP0, mode=[2]) - nblk_qk = cute.size(tCrA, mode=[2]) - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - # COMPREHENSIVE LAYOUT DUMP - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tilePlikeFP32 = self.qk_mma_tiler[1] * self.q_dtype.width // 32 - tStS_P_layout = cute.composition(tStS.layout, cute.make_layout((128, tilePlikeFP32))) - tStS_P = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStS_P_layout) - tScS_P_layout = cute.composition(tScS.layout, cute.make_layout((128, tilePlikeFP32))) - tScS_P = cute.make_tensor(tScS.iterator, tScS_P_layout) - - print(f'[LAYOUT] QK C-fragment tStS.layout: {tStS.layout}') - print(f'[LAYOUT] QK C-fragment tStS cosize: {cute.cosize(tStS.layout)}') - print(f'[LAYOUT] QK C-fragment tStS.size: {cute.size(tStS)}') - print(f'[LAYOUT] QK C-fragment tScS.layout: {tScS.layout}') - print(f'[LAYOUT] QK C-fragment tScS cosize: {cute.cosize(tScS.layout)}') - print(f'[LAYOUT] PV A-fragment tOrP.layout: {tOrP.layout}') - print(f'[LAYOUT] PV A-fragment tOrP cosize: {cute.cosize(tOrP.layout)}') - print(f'[LAYOUT] PV A-fragment tOrP.size: {cute.size(tOrP)}') - print(f'[LAYOUT] PV A-fragment tOrP0.layout: {tOrP0.layout}') - print(f'[LAYOUT] PV A-fragment tOrP0 cosize: {cute.cosize(tOrP0.layout)}') - print(f'[LAYOUT] tP.layout: {tP.layout}') - print(f'[LAYOUT] tP cosize: {cute.cosize(tP.layout)}') - print(f'[LAYOUT] tStS_P (composed) layout: {tStS_P.layout}') - print(f'[LAYOUT] tStS_P (composed) cosize: {cute.cosize(tStS_P.layout)}') - print(f'[LAYOUT] tScS_P (composed) layout: {tScS_P.layout}') - print(f'[LAYOUT] tScS_P (composed) cosize: {cute.cosize(tScS_P.layout)}') - print(f'[LAYOUT] tOtO.layout: {tOtO.layout}') - print(f'[LAYOUT] tOtO cosize: {cute.cosize(tOtO.layout)}') - print(f'[LAYOUT] pv_mma_tiler: {self.pv_mma_tiler}') - print(f'[LAYOUT] qk_mma_tiler: {self.qk_mma_tiler}') - print(f'[LAYOUT] tilePlikeFP32: {tilePlikeFP32}') - - # DIAGNOSTIC: Compare tP (A-layout) vs tStS_P (composition) - tilePlikeFP32 = self.qk_mma_tiler[1] * self.q_dtype.width // 32 - tStS_P_layout = cute.composition(tStS.layout, cute.make_layout((128, tilePlikeFP32))) - tStS_P = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStS_P_layout) - print(f'[DIAG] tP.layout: {tP.layout}') - print(f'[DIAG] tP.size: {cute.size(tP)}') - print(f'[DIAG] tP.element_type: {tP.element_type if hasattr(tP, 'element_type') else 'N/A'}') - print(f'[DIAG] tStS_P.layout: {tStS_P.layout}') - print(f'[DIAG] tStS_P.size: {cute.size(tStS_P)}') - print(f'[DIAG] tStS_P.element_type: {tStS_P.element_type if hasattr(tStS_P, 'element_type') else 'N/A'}') - print(f'[DIAG] tilePlikeFP32: {tilePlikeFP32}') - print(f'[DIAG] tP and tStS_P same iterator? {tP.iterator == tStS_P.iterator if hasattr(tP, 'iterator') else 'cant compare'}') - - print(f'[DIAG] nblk_pv = {nblk_pv}, nblk_qk = {nblk_qk}') - print(f'[DIAG] tCrV.size = {cute.size(tCrV)}') - print(f'[DIAG] tOrP0.size = {cute.size(tOrP0)}') - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ── TMA WARP ── - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_a, tAgA[(None,h.count)], tAsA[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_b, tBgB[(None,h.count)], tBsB[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1 BF16 (backward FMHA pattern) - tRT_rST_bf16 = cute.make_rmem_tensor(tTMEM_LOADrS.shape, self.q_dtype) - frg_cnt = 4 - frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt - tTR_rST_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile)) - tRT_rST_bf16_frg = cute.make_tensor(tRT_rST_bf16.iterator, tTR_rST_frg.layout) - for j in range(frg_cnt): - frg_vec = tTR_rST_frg[None, j].load() - tRT_rST_bf16_frg[None, j].store(frg_vec.to(self.q_dtype)) - - # 6. Reshape and store to A-fragment TMEM - tRT_rST_reshaped = cute.make_tensor( - tRT_rST_bf16.iterator, cute.make_layout(tRT_cS.shape)) - cute.copy(tiled_tmem_store, tRT_rST_reshaped, tRT_tP) - cute.arch.fence_view_async_tmem_store() - - # 7. Release back to MMA warp - si_handle.release() - - # ── Epilogue ── - tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout) - acc_cons_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage) - c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) - c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp) - acc_cons_st = utils.gemm.sm100.epilogue_tma_store( - self, tidx, warp_idx, tma_c, tCtO_base, sC, tCgC, - epi_tile, 0, const_expr(lambda x: x), (0,0,0), acc_cons_st, acc_pipe, c_pipe) - c_pipe.producer_tail() - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def test(): - torch.manual_seed(42) - m, n, k = 128, 128, 128 - q = torch.randn(m, k, 1, dtype=torch.bfloat16, device='cuda') - kv = torch.randn(n, k, 1, dtype=torch.bfloat16, device='cuda') - c = torch.zeros(m, n, 1, dtype=torch.bfloat16, device='cuda') - qf = q[:,:,0].float(); kvf = kv[:,:,0].float() - ref = qf @ kvf.T @ kvf - import cutlass.torch as ct - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(kv).mark_layout_dynamic(leading_dim=ct.get_leading_dim(kv)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - kernel = StageBIdentitySoftmax(mma_tiler_mn=(128, 128), use_2cta_instrs=False, use_tma_store=True) - print('Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mC, stream) - print('Running...', flush=True) - compiled(mQ, mK, mC, stream) - torch.cuda.synchronize() - out = c[:,:,0].float() - cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - max_err = (out - ref).abs().max().item() - print('Stage B v7: (Q @ K^T) @ V with identity softmax (computed TMEM offsets)') - print(' Cosine: {:.6f}, Max error: {:.6f}'.format(cos, max_err)) - print(' {}'.format('PASS' if cos >= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_stage_b_v19.py b/tests/archive/test_stage_b_v19.py deleted file mode 100644 index a523ee6e..00000000 --- a/tests/archive/test_stage_b_v19.py +++ /dev/null @@ -1,450 +0,0 @@ -""" -Stage B v7: Two MMAs + Identity Softmax with COMPUTED TMEM offsets. - -Key fixes over v6: - - TMEM offsets computed via find_tmem_tensor_col_offset (same API as get_num_tmem_alloc_cols) - - P tensor constructed from p_tmem_s.outer (matching fmha.py pattern exactly) - - tilePlikeFP32 computed from qk_mma_tiler and dtype widths - - tmem_alloc_cols from get_num_tmem_alloc_cols with all fragments - - JIT-time diagnostic prints of all TMEM sizes - -Architecture (matches fmha.py exactly): - MMA1: Q @ K^T → tmem_scores (a_source=SMEM, accumulate=False) - Identity softmax: tcgen05.ld C-layout → F32→BF16 → tcgen05.st A-layout - MMA2: P @ V → tmem_output (a_source=TMEM, accumulate=True) -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - - -class StageBIdentitySoftmax: - def __init__(self, mma_tiler_mn, use_2cta_instrs=False, use_tma_store=True): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16 - self.use_2cta_instrs = use_2cta_instrs; self.use_tma_store = use_tma_store - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.epilog_sync_bar_id = 1; self.tmem_alloc_sync_bar_id = 2; self.tmem_dealloc_sync_bar_id = 3 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - pv_inst_k = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (*self.mma_tiler_mn, pv_inst_k * 4) - self.mma_tiler = self.qk_mma_tiler - print(f"[StageB] qk_mma.shape_mnk = {qk_mma.shape_mnk}") - print(f"[StageB] pv_mma.shape_mnk = {pv_mma.shape_mnk}") - print(f"[StageB] qk_mma_tiler = {self.qk_mma_tiler}") - print(f"[StageB] pv_mma_tiler = {self.pv_mma_tiler}") - print(f"[StageB] qk_inst_k = {qk_inst_k}, pv_inst_k = {pv_inst_k}") - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], - self.qk_mma_tiler[2], - ) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, self.use_2cta_instrs, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.a_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.b_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.b_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - # ── COMPUTE TMEM OFFSETS USING find_tmem_tensor_col_offset ── - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - # tilePlikeFP32 for the store-side composition - self.tilePlikeFP32 = self.qk_mma_tiler[1] * self.q_dtype.width // 32 - - # ── TMEM LAYOUT (matching fmha.py) ── - # P OVERLAPS S — after softmax, P (A-layout) is written on top of scores (C-layout) - # in the same TMEM region. The A-layout view starts partway through the S region. - # fmha.py: S0=0, P0=32, O0=256 (with S1=128, P1=160 double-buffered) - # The P offset of 32 means the A-layout starts at column 32 within the S region. - # This is because the C-layout and A-layout partition TMEM differently per-thread; - # the first 32 C-layout columns are "dead space" in the A-layout mapping. - # - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 # Original - self.tmem_o0_offset = s_cols # 128 - self.tmem_alloc_cols = s_cols + o_cols # 256 - - # Also compute via get_num_tmem_alloc_cols for the full allocation - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, 1)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, 1)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - print(f"[StageB] s_cols (QK accumulator) = {s_cols}") - print(f"[StageB] o_cols (PV accumulator) = {o_cols}") - print(f"[StageB] tilePlikeFP32 = {self.tilePlikeFP32}") - print(f"[StageB] tmem_s0_offset = {self.tmem_s0_offset}") - print(f"[StageB] tmem_p0_offset = {self.tmem_p0_offset}") - print(f"[StageB] tmem_o0_offset = {self.tmem_o0_offset}") - print(f"[StageB] tmem_alloc_cols (computed) = {self.tmem_alloc_cols}") - print(f"[StageB] num_tmem_alloc_cols (via utils) = {self.num_tmem_alloc_cols}") - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.a_dtype, a_smem) + cute.size_in_bytes(self.b_dtype, b_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - self.a_dtype = a.element_type; self.b_dtype = b.element_type; self.c_dtype = c.element_type - self.a_major = LayoutEnum.from_tensor(a).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(b).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, cute.nvgpu.OperandMajorMode.K, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - # Introspect PV MMA atom - print(f"[ATOM] PV MMA type: {type(pv_mma)}") - print(f"[ATOM] PV MMA op: {pv_mma.op if hasattr(pv_mma, "op") else "no op"}") - print(f"[ATOM] PV MMA trait: {pv_mma._trait if hasattr(pv_mma, "_trait") else "no trait"}") - print(f"[ATOM] PV MMA shape_mnk: {pv_mma.shape_mnk}") - print(f"[ATOM] QK MMA shape_mnk: {qk_mma.shape_mnk}") - # Check a_src - print(f"[ATOM] PV MMA op.a_src: {pv_mma.op.a_src}") - print(f"[ATOM] QK MMA op.a_src: {qk_mma.op.a_src}") - print(f"[ATOM] PV MMA op: {pv_mma.op}") - self._setup(qk_mma, pv_mma) - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - tma_a, tma_ta = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - a, a_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_b, tma_tb = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - b, b_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_a, tma_ta, tma_b, tma_tb, tma_c, tma_tc, - self.cluster_layout_vmnk, self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_a, mA, tma_b, mB, tma_c, mC, cl_vmnk, - a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_a); cpasync.prefetch_descriptor(tma_b); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=self.tmem_alloc_sync_bar_id, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sA = smem.allocate_tensor(element_type=self.a_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sB = smem.allocate_tensor(element_type=self.b_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - # V shares the same SMEM as B (same data, different layout for PV MMA) - sV_ptr = cute.recast_ptr(sB.iterator, v_smem_s.inner) - sV = cute.make_tensor(sV_ptr, v_smem_s.outer) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gA = cute.local_tile(mA, cute.slice_(self.mma_tiler, (None,0,None)), (None,None,None)) - gB = cute.local_tile(mB, cute.slice_(self.mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gA, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - tCgA = qk_thr.partition_A(gA); tCgB = qk_thr.partition_B(gB); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsA, tAgA = cpasync.tma_partition(tma_a, 0, a_lay, cute.group_modes(sA,0,3), cute.group_modes(tCgA,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsB, tBgB = cpasync.tma_partition(tma_b, 0, b_lay, cute.group_modes(sB,0,3), cute.group_modes(tCgB,0,3)) - tAgA = tAgA[(None,0,None,0)]; tBgB = tBgB[(None,0,None,0)] - - tCrA = qk_mma.make_fragment_A(sA); tCrB = qk_mma.make_fragment_B(sB) - tCrV = pv_mma.make_fragment_B(sV) # V fragment from V SMEM layout - print(f"[DIAG] tCrV.size = {cute.size(tCrV)}") - print(f"[DIAG] tCrA.size = {cute.size(tCrA)}") - print(f"[DIAG] tCrB.size = {cute.size(tCrB)}") - print(f"[DIAG] nblk_qk (tCrA mode 2) = {cute.size(tCrA, mode=[2])}") - - # ── TMEM tensors with computed offsets (matching fmha.py pattern) ── - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - # P fragment: construct from p_tmem_s layout (matching fmha.py exactly) - # fmha.py: tP = cute.make_tensor(tStS.iterator, p_tmem_layout_staged.outer) - # tOrP = pv_thr_mma.make_fragment_A(tP)[None, None, None, 0] - # tOrP0 = cute.make_tensor(tOrP.iterator + dtype_width_ratio * tmem_p0_offset, tOrP.layout) - print(f'[TMEM] p_tmem_s: {p_tmem_s}') - print(f'[TMEM] p_tmem_s.outer: {p_tmem_s.outer}') - print(f'[TMEM] p_tmem_s.inner: {p_tmem_s.inner}') - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - print(f'[DIAG] tStS.layout: {tStS.layout}') - print(f'[DIAG] tStS.size: {cute.size(tStS)}') - print(f'[DIAG] p_tmem_s.outer: {p_tmem_s.outer}') - print(f'[DIAG] p_tmem_s.inner: {p_tmem_s.inner}') - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - # Compute nblk_pv for diagnostics - nblk_pv = cute.size(tOrP0, mode=[2]) - nblk_qk = cute.size(tCrA, mode=[2]) - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - # COMPREHENSIVE LAYOUT DUMP - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tilePlikeFP32 = self.qk_mma_tiler[1] * self.q_dtype.width // 32 - tStS_P_layout = cute.composition(tStS.layout, cute.make_layout((128, tilePlikeFP32))) - tStS_P = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStS_P_layout) - tScS_P_layout = cute.composition(tScS.layout, cute.make_layout((128, tilePlikeFP32))) - tScS_P = cute.make_tensor(tScS.iterator, tScS_P_layout) - - print(f'[LAYOUT] QK C-fragment tStS.layout: {tStS.layout}') - print(f'[LAYOUT] QK C-fragment tStS cosize: {cute.cosize(tStS.layout)}') - print(f'[LAYOUT] QK C-fragment tStS.size: {cute.size(tStS)}') - print(f'[LAYOUT] QK C-fragment tScS.layout: {tScS.layout}') - print(f'[LAYOUT] QK C-fragment tScS cosize: {cute.cosize(tScS.layout)}') - print(f'[LAYOUT] PV A-fragment tOrP.layout: {tOrP.layout}') - print(f'[LAYOUT] PV A-fragment tOrP cosize: {cute.cosize(tOrP.layout)}') - print(f'[LAYOUT] PV A-fragment tOrP.size: {cute.size(tOrP)}') - print(f'[LAYOUT] PV A-fragment tOrP0.layout: {tOrP0.layout}') - print(f'[LAYOUT] PV A-fragment tOrP0 cosize: {cute.cosize(tOrP0.layout)}') - print(f'[LAYOUT] tP.layout: {tP.layout}') - print(f'[LAYOUT] tP cosize: {cute.cosize(tP.layout)}') - print(f'[LAYOUT] tStS_P (composed) layout: {tStS_P.layout}') - print(f'[LAYOUT] tStS_P (composed) cosize: {cute.cosize(tStS_P.layout)}') - print(f'[LAYOUT] tScS_P (composed) layout: {tScS_P.layout}') - print(f'[LAYOUT] tScS_P (composed) cosize: {cute.cosize(tScS_P.layout)}') - print(f'[LAYOUT] tOtO.layout: {tOtO.layout}') - print(f'[LAYOUT] tOtO cosize: {cute.cosize(tOtO.layout)}') - print(f'[LAYOUT] pv_mma_tiler: {self.pv_mma_tiler}') - print(f'[LAYOUT] qk_mma_tiler: {self.qk_mma_tiler}') - print(f'[LAYOUT] tilePlikeFP32: {tilePlikeFP32}') - - # DIAGNOSTIC: Compare tP (A-layout) vs tStS_P (composition) - tilePlikeFP32 = self.qk_mma_tiler[1] * self.q_dtype.width // 32 - tStS_P_layout = cute.composition(tStS.layout, cute.make_layout((128, tilePlikeFP32))) - tStS_P = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStS_P_layout) - print(f'[DIAG] tP.layout: {tP.layout}') - print(f'[DIAG] tP.size: {cute.size(tP)}') - print(f'[DIAG] tP.element_type: {tP.element_type if hasattr(tP, 'element_type') else 'N/A'}') - print(f'[DIAG] tStS_P.layout: {tStS_P.layout}') - print(f'[DIAG] tStS_P.size: {cute.size(tStS_P)}') - print(f'[DIAG] tStS_P.element_type: {tStS_P.element_type if hasattr(tStS_P, 'element_type') else 'N/A'}') - print(f'[DIAG] tilePlikeFP32: {tilePlikeFP32}') - print(f'[DIAG] tP and tStS_P same iterator? {tP.iterator == tStS_P.iterator if hasattr(tP, 'iterator') else 'cant compare'}') - - print(f'[DIAG] nblk_pv = {nblk_pv}, nblk_qk = {nblk_qk}') - print(f'[DIAG] tCrV.size = {cute.size(tCrV)}') - print(f'[DIAG] tOrP0.size = {cute.size(tOrP0)}') - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ── TMA WARP ── - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_a, tAgA[(None,h.count)], tAsA[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_b, tBgB[(None,h.count)], tBsB[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_stage_b_v2.py b/tests/archive/test_stage_b_v2.py deleted file mode 100644 index 2d2284f7..00000000 --- a/tests/archive/test_stage_b_v2.py +++ /dev/null @@ -1,407 +0,0 @@ -""" -Stage B v2: Two MMAs (Q@K^T then Scores@V) — no softmax, no identity P. - -Tests: -- MMA1: Q @ K^T → tmem_scores (accumulate=False) -- MMA2: tmem_scores @ V → tmem_output (a_source=TMEM, accumulate=True) -- Two TMEM regions with pointer arithmetic -- The epilogue warps allocate TMEM and store output, but skip softmax for now - -Reference: output = Q @ K^T @ V (with K=V for simplicity) -""" -import torch -import cutlass -import cutlass.cute as cute -import cutlass.utils as utils -import cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -import cuda.bindings.driver as cuda - - -class StageBKernel: - def __init__(self, mma_tiler_mn, use_2cta_instrs=False): - self.acc_dtype = Float32 - self.use_2cta_instrs = use_2cta_instrs - self.mma_tiler_mn = mma_tiler_mn - self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE - # Warp layout: 4 epilogue + 1 MMA + 1 TMA = 6 warps = 192 threads - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4 - self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.epilog_sync_bar_id = 1 - self.tmem_alloc_sync_bar_id = 2 - - def _setup_attributes(self, tiled_mma1, tiled_mma2, a_dtype, b_dtype, c_dtype, - a_major, b_major, c_layout): - mma_inst_shape_k = cute.size(tiled_mma1.shape_mnk, mode=[2]) - self.mma_tiler = (*self.mma_tiler_mn, mma_inst_shape_k * 4) - self.cta_tile_shape_mnk = ( - self.mma_tiler[0] // cute.size(tiled_mma1.thr_id.shape), - self.mma_tiler[1], - self.mma_tiler[2], - ) - self.cluster_layout_vmnk = cute.tiled_divide( - cute.make_layout((1, 1, 1)), (tiled_mma1.thr_id.shape,)) - - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, self.use_2cta_instrs, c_layout, c_dtype) - - self.num_ab_stage = 1 - self.num_acc_stage = 1 - self.num_c_stage = 2 - - self.a_smem_layout_staged = utils.sm100.make_smem_layout_a( - tiled_mma1, self.mma_tiler, a_dtype, self.num_ab_stage) - self.b_smem_layout_staged = utils.sm100.make_smem_layout_b( - tiled_mma1, self.mma_tiler, b_dtype, self.num_ab_stage) - self.c_smem_layout_staged = utils.sm100.make_smem_layout_epi( - c_dtype, c_layout, self.epi_tile, self.num_c_stage) - - # TMEM: two regions (scores + output), each partition_shape_C columns - acc_shape = tiled_mma1.partition_shape_C(self.mma_tiler_mn) - tCtAcc_fake = tiled_mma1.make_fragment_C(cute.append(acc_shape, self.num_acc_stage)) - self.num_tmem_cols_per_region = utils.get_num_tmem_alloc_cols(tCtAcc_fake, arch="sm_100") - self.total_tmem_cols = max(self.num_tmem_cols_per_region * 2, 256) - - a_smem = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(a_dtype, a_smem) + - cute.size_in_bytes(b_dtype, b_smem) - ) * cute.size(tiled_mma1.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, - stream: cuda.CUstream): - a_dtype = a.element_type - b_dtype = b.element_type - c_dtype = c.element_type - a_major = LayoutEnum.from_tensor(a).mma_major_mode() - b_major = LayoutEnum.from_tensor(b).mma_major_mode() - c_layout = LayoutEnum.from_tensor(c) - - tiled_mma1 = utils.sm100.make_trivial_tiled_mma( - a_dtype, b_dtype, a_major, b_major, - self.acc_dtype, self.cta_group, self.mma_tiler_mn, - tcgen05.OperandSource.SMEM, - ) - tiled_mma2 = utils.sm100.make_trivial_tiled_mma( - a_dtype, b_dtype, a_major, b_major, - self.acc_dtype, self.cta_group, self.mma_tiler_mn, - tcgen05.OperandSource.SMEM, - ) - - self._setup_attributes(tiled_mma1, tiled_mma2, a_dtype, b_dtype, c_dtype, - a_major, b_major, c_layout) - # These are needed by epilogue_tma_store which accesses them via self - self.a_dtype = a_dtype - self.b_dtype = b_dtype - self.c_dtype = c_dtype - self.a_major_mode = a_major - self.b_major_mode = b_major - self.c_layout = c_layout - - a_smem = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) - - tma_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A( - self.cluster_shape_mn, tiled_mma1.thr_id), - a, a_smem, self.mma_tiler, tiled_mma1, - self.cluster_layout_vmnk.shape, - ) - tma_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B( - self.cluster_shape_mn, tiled_mma1.thr_id), - b, b_smem, self.mma_tiler, tiled_mma1, - self.cluster_layout_vmnk.shape, - ) - - epi_smem = cute.select(self.c_smem_layout_staged, mode=[0, 1]) - tma_c, tma_tensor_c = cpasync.make_tiled_tma_atom( - cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel( - tiled_mma1, tiled_mma2, - tma_a, tma_tensor_a, tma_b, tma_tensor_b, - tma_c, tma_tensor_c, self.cluster_layout_vmnk, - self.a_smem_layout_staged, self.b_smem_layout_staged, - self.c_smem_layout_staged, self.epi_tile, - ).launch(grid=(1, 1, 1), block=[self.threads_per_cta, 1, 1], stream=stream) - - @cute.kernel - def _kernel(self, tiled_mma1, tiled_mma2, - tma_atom_a, mA_mkl, tma_atom_b, mB_nkl, - tma_atom_c, mC_mnl, cluster_layout_vmnk, - a_smem_layout_staged, b_smem_layout_staged, - c_smem_layout_staged, epi_tile): - warp_idx = cute.arch.warp_idx() - warp_idx = cute.arch.make_warp_uniform(warp_idx) - tidx, _, _ = cute.arch.thread_idx() - use_2cta_instrs = cute.size(tiled_mma1.thr_id.shape) == 2 - is_leader_cta = True - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_atom_a) - cpasync.prefetch_descriptor(tma_atom_b) - cpasync.prefetch_descriptor(tma_atom_c) - - @cute.struct - class SharedStorage: - ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc_mbar: cutlass.Int64 - tmem_holding_buf: cutlass.Int32 - - smem = utils.SmemAllocator() - storage = smem.allocate(SharedStorage) - - ab_producer, ab_consumer = pipeline.PipelineTmaUmma.create( - barrier_storage=storage.ab_full_mbar_ptr.data_ptr(), - num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, - cta_layout_vmnk=cluster_layout_vmnk, - defer_sync=True, - ).make_participants() - - acc_pipeline = pipeline.PipelineUmmaAsync.create( - barrier_storage=storage.acc_full_mbar_ptr.data_ptr(), - num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta_instrs else 1)), - cta_layout_vmnk=cluster_layout_vmnk, - defer_sync=True, - ) - - tmem_alloc_barrier = pipeline.NamedBarrier( - barrier_id=self.tmem_alloc_sync_bar_id, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id)), - ) - tmem = utils.TmemAllocator( - storage.tmem_holding_buf.ptr, - barrier_for_retrieve=tmem_alloc_barrier, - allocator_warp_id=self.epilogue_warp_id[0], - is_two_cta=use_2cta_instrs, - two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr, - ) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cluster_layout_vmnk, is_relaxed=True) - - sA = smem.allocate_tensor( - element_type=BFloat16, - layout=a_smem_layout_staged.outer, - byte_alignment=128, swizzle=a_smem_layout_staged.inner) - sB = smem.allocate_tensor( - element_type=BFloat16, - layout=b_smem_layout_staged.outer, - byte_alignment=128, swizzle=b_smem_layout_staged.inner) - sC = smem.allocate_tensor( - element_type=BFloat16, - layout=c_smem_layout_staged.outer, - byte_alignment=128, swizzle=c_smem_layout_staged.inner) - - gA_mkl = cute.local_tile(mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)) - gB_nkl = cute.local_tile(mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None)) - gC_mnl = cute.local_tile(mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None)) - k_tile_cnt = cute.size(gA_mkl, mode=[3]) - - thr_mma1 = tiled_mma1.get_slice(0) - tCgA = thr_mma1.partition_A(gA_mkl) - tCgB = thr_mma1.partition_B(gB_nkl) - tCgC = thr_mma1.partition_C(gC_mnl) - - a_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape) - tAsA, tAgA = cpasync.tma_partition( - tma_atom_a, 0, a_cta_layout, - cute.group_modes(sA, 0, 3), cute.group_modes(tCgA, 0, 3)) - b_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape) - tBsB, tBgB = cpasync.tma_partition( - tma_atom_b, 0, b_cta_layout, - cute.group_modes(sB, 0, 3), cute.group_modes(tCgB, 0, 3)) - - tAgA_slice = tAgA[(None, 0, None, 0)] - tBgB_slice = tBgB[(None, 0, None, 0)] - - # MMA1 fragments - tCrA = tiled_mma1.make_fragment_A(sA) - tCrB = tiled_mma1.make_fragment_B(sB) - # MMA2 fragment for B (V) - same SMEM as K - tCrB_mma2 = tiled_mma2.make_fragment_B(sB) - # MMA2 fragment for A (a_source=SMEM, same as MMA1) - tCrA_mma2 = tiled_mma2.make_fragment_A(sA) - - # TMEM accumulator layout (same for both MMAs) - acc_shape = tiled_mma1.partition_shape_C(self.mma_tiler_mn) - tCtAcc_fake = tiled_mma1.make_fragment_C(cute.append(acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cluster_layout_vmnk) - - # ── TMA LOAD WARP (warp 5) ── - if warp_idx == self.tma_warp_id: - ab_producer.reset() - peek = ab_producer.try_acquire() - for k_tile in cutlass.range(k_tile_cnt, unroll=1): - handle = ab_producer.acquire_and_advance(peek) - cute.copy(tma_atom_a, tAgA_slice[(None, handle.count)], tAsA[(None, handle.index)], - tma_bar_ptr=handle.barrier) - cute.copy(tma_atom_b, tBgB_slice[(None, handle.count)], tBsB[(None, handle.index)], - tma_bar_ptr=handle.barrier) - peek = cutlass.Boolean(1) - if handle.count + 1 < k_tile_cnt: - peek = ab_producer.try_acquire() - ab_producer.tail() - - # ── MMA WARP (warp 4) ── - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) - - # TMEM region 0: scores (Q @ K^T) - tCtScores_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) - tCtScores = tCtScores_base[(None, None, None, 0)] - # MMA2 A fragment from SMEM (a_source=SMEM) - tCrA_mma2 = tiled_mma2.make_fragment_A(sA) - - # TMEM region 1: output (Scores @ V) - output_ptr = cute.recast_ptr( - tmem_ptr + self.num_tmem_cols_per_region, dtype=self.acc_dtype) - tCtOutput_base = cute.make_tensor(output_ptr, tCtAcc_fake.layout) - tCtOutput = tCtOutput_base[(None, None, None, 0)] - - ab_consumer.reset() - peek = cutlass.Boolean(1) - if is_leader_cta: - peek = ab_consumer.try_wait() - - # ── MMA1: Q @ K^T → tmem_scores ── - tiled_mma1.set(tcgen05.Field.ACCUMULATE, False) - for k_tile in range(k_tile_cnt): - if is_leader_cta: - handle = ab_consumer.wait_and_advance(peek) - nblk = cute.size(tCrA, mode=[2]) - for kblk in cutlass.range(nblk, unroll_full=True): - crd = (None, None, kblk, handle.index) - cute.gemm(tiled_mma1, tCtScores, tCrA[crd], tCrB[crd], tCtScores) - handle.release() - peek = cutlass.Boolean(1) - if handle.count + 1 < k_tile_cnt: - peek = ab_consumer.try_wait() - - # MMA1 done, scores in tmem_scores - # Fence to ensure TMEM writes are visible - cute.arch.fence_view_async_tmem_store() - - # ── MMA2: Scores @ V → tmem_output ── - # a_source=TMEM: the MMA instruction reads A from TMEM - # The A operand in TMEM is the scores from MMA1 - # We need to set the TMEM pointer for A using tiled_mma2.set() - # The C operand's TMEM pointer tells MMA2 where to write the output - # The A operand's TMEM pointer needs to be set separately - tiled_mma2.set(tcgen05.Field.ACCUMULATE, True) - # Set the A TMEM pointer to the scores region - # When a_source=TMEM, the MMA reads A from the C operand's base - # So we need tCtOutput to have the same base as tCtScores for MMA2 - # This means we can't use two separate TMEM regions for this approach - # Instead, let's try passing the TMEM A operand via the auxiliary list - # cute.gemm(tiled_mma2, D, [A_main, A_tmem], B, C) - # Or use tiled_mma2.set(Field.SFA, tCtScores.iterator) to set A ptr - nblk2 = cute.size(tCrB_mma2, mode=[2]) - for kblk in cutlass.range(nblk2, unroll_full=True): - crd = (None, None, kblk, 0) - cute.gemm(tiled_mma2, tCtOutput, tCrA_mma2[crd], tCrB_mma2[crd], tCtOutput) - - # Signal output ready - acc_producer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Producer, self.num_acc_stage) - if is_leader_cta: - acc_pipeline.producer_acquire(acc_producer_state) - acc_pipeline.producer_commit(acc_producer_state) - acc_producer_state.advance() - acc_pipeline.producer_tail(acc_producer_state) - - # ── EPILOGUE WARPS (0..3) ── - if warp_idx < self.mma_warp_id: - tmem.allocate(self.total_tmem_cols) - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) - - # TMEM region 1: output - output_ptr = cute.recast_ptr( - tmem_ptr + self.num_tmem_cols_per_region, dtype=self.acc_dtype) - tCtOutput_base = cute.make_tensor(output_ptr, tCtAcc_fake.layout) - - # Wait for MMA2 to finish - acc_consumer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Consumer, self.num_acc_stage) - - c_producer_group = pipeline.CooperativeGroup( - pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) - c_pipeline = pipeline.PipelineTmaStore.create( - num_stages=self.num_c_stage, producer_group=c_producer_group) - - mma_tile_coord_mnl = (0, 0, 0) - epilogue_op = const_expr(lambda x: x) - num_tiles_executed = 0 - - acc_consumer_state = utils.gemm.sm100.epilogue_tma_store( - self, tidx, warp_idx, tma_atom_c, tCtOutput_base, sC, tCgC, - epi_tile, num_tiles_executed, epilogue_op, - mma_tile_coord_mnl, acc_consumer_state, acc_pipeline, c_pipeline) - - c_pipeline.producer_tail() - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def test_stage_b(): - torch.manual_seed(42) - m, n, k = 128, 128, 128 - - a = torch.randn(m, k, 1, dtype=torch.bfloat16, device="cuda") - b = torch.randn(n, k, 1, dtype=torch.bfloat16, device="cuda") - c = torch.zeros(m, n, 1, dtype=torch.bfloat16, device="cuda") - - # Reference: Q @ K^T @ V (no softmax, K=V=b) - q = a[:, :, 0].float() - kt = b[:, :, 0].float().T # K^T - v = b[:, :, 0].float() - ref = q @ kt @ v # (128, 128) - - import cutlass.torch as cutlass_torch - mA = cutlass_torch.from_dlpack(a).mark_layout_dynamic( - leading_dim=cutlass_torch.get_leading_dim(a)) - mB = cutlass_torch.from_dlpack(b).mark_layout_dynamic( - leading_dim=cutlass_torch.get_leading_dim(b)) - mC = cutlass_torch.from_dlpack(c).mark_layout_dynamic( - leading_dim=cutlass_torch.get_leading_dim(c)) - - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - kernel = StageBKernel(mma_tiler_mn=(128, 128)) - print("Compiling...", flush=True) - compiled = cute.compile(kernel, mA, mB, mC, stream) - print("Running...", flush=True) - compiled(mA, mB, mC, stream) - torch.cuda.synchronize() - - output = c[:, :, 0].float() - cos = torch.nn.functional.cosine_similarity( - output.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - max_err = (output - ref).abs().max().item() - - print("Stage B v2: Q @ K^T @ V (no softmax)") - print(" Cosine: {:.6f}, Max error: {:.6f}".format(cos, max_err)) - print(" {}".format("PASS" if cos >= 0.99 else "FAIL")) - return cos - - -if __name__ == "__main__": - test_stage_b() diff --git a/tests/archive/test_stage_b_v20.py b/tests/archive/test_stage_b_v20.py deleted file mode 100644 index dbd704e3..00000000 --- a/tests/archive/test_stage_b_v20.py +++ /dev/null @@ -1,362 +0,0 @@ -""" -Stage B v20: FMHA-matching test with head_dim=64, proper V layout. - -KEY INSIGHT: The A-fragment ((128,16),1,(4,2)):((65536,1),0,(16,64)) is SEQUENTIAL -when flattened in CuTe order: addr = m*65536 + k0 + 16*k1 + 64*k2 = m*65536 + k. -So the C-fragment composition store aliases the SAME TMEM as the A-fragment read. - -Previous -0.02 cosine was caused by V dimension mismatch: -pv_mma_tiler=(128,64,128) expects V with N=64 (head_dim), -but the square 128x128 test had V=K (N=128). - -This test: Q=(128,64), K=(128,64), V=(64,128), O=(128,64) - -FOOTGUN: St32x32bOp MUST use Float32, NOT BFloat16! -The 16-bit values are packed via recast_ptr view. -St32x32bOp(BFloat16) causes ILLEGAL MEMORY ACCESS. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - - -class StageBIdentitySoftmax: - def __init__(self, mma_tiler_mn, use_2cta_instrs=False, use_tma_store=True): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16 - self.use_2cta_instrs = use_2cta_instrs; self.use_tma_store = use_tma_store - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.epilog_sync_bar_id = 1; self.tmem_alloc_sync_bar_id = 2; self.tmem_dealloc_sync_bar_id = 3 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - self.pv_mma_tiler = (self.qk_mma_tiler[0], self.qk_mma_tiler[2], self.qk_mma_tiler[1]) - self.mma_tiler = self.qk_mma_tiler - print(f"[v20] qk_mma_tiler = {self.qk_mma_tiler}") - print(f"[v20] pv_mma_tiler = {self.pv_mma_tiler}") - - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], - self.qk_mma_tiler[2], - ) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, self.use_2cta_instrs, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.a_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.b_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.b_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - self.tmem_alloc_cols = s_cols + o_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.a_dtype, a_smem) + cute.size_in_bytes(self.b_dtype, b_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, v: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - self.a_dtype = a.element_type; self.b_dtype = b.element_type; self.c_dtype = c.element_type - self.a_major = LayoutEnum.from_tensor(a).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(b).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.b_dtype, cute.nvgpu.OperandMajorMode.K, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_a, tma_ta = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - a, a_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_b, tma_tb = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - b, b_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - # Separate TMA for V (uses pv_mma_tiler and v_smem layout) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_a, tma_ta, tma_b, tma_tb, tma_v, tma_tv, tma_c, tma_tc, - self.cluster_layout_vmnk, self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_a, mA, tma_b, mB, tma_v, mV, tma_c, mC, cl_vmnk, - a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_a); cpasync.prefetch_descriptor(tma_b); cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=self.tmem_alloc_sync_bar_id, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sA = smem.allocate_tensor(element_type=self.a_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sB = smem.allocate_tensor(element_type=self.b_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.b_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gA = cute.local_tile(mA, cute.slice_(self.mma_tiler, (None,0,None)), (None,None,None)) - gB = cute.local_tile(mB, cute.slice_(self.mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gA, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgA = qk_thr.partition_A(gA); tCgB = qk_thr.partition_B(gB); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsA, tAgA = cpasync.tma_partition(tma_a, 0, a_lay, cute.group_modes(sA,0,3), cute.group_modes(tCgA,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsB, tBgB = cpasync.tma_partition(tma_b, 0, b_lay, cute.group_modes(sB,0,3), cute.group_modes(tCgB,0,3)) - tAgA = tAgA[(None,0,None,0)]; tBgB = tBgB[(None,0,None,0)] - - # V partition (from mV with pv_mma_tiler and tma_v) - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrA = qk_mma.make_fragment_A(sA); tCrB = qk_mma.make_fragment_B(sB) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_a, tAgA[(None,h.count)], tAsA[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_b, tBgB[(None,h.count)], tBsB[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_b, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1BF16 packing (EXACT FMHA pattern) - tTMEM_STORErS_x4 = cute.make_rmem_tensor(tTMEM_STOREcS.shape, self.qk_acc_dtype) - tTMEM_STORErS_x4_e = cute.make_tensor( - cute.recast_ptr(tTMEM_STORErS_x4.iterator, dtype=self.q_dtype), - tTMEM_LOADrS.layout) - - frg_cnt = 4 - frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt - tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile)) - tTMEM_STORErS_x4_e_frg = cute.logical_divide( - tTMEM_STORErS_x4_e, cute.make_layout(frg_tile)) - for j in range(frg_cnt): - s_vec = tTMEM_LOADrS_frg[None, j].load() - tTMEM_STORErS_x4_e_frg[None, j].store(s_vec.to(self.q_dtype)) - - cute.copy(tiled_tmem_store, tTMEM_STORErS_x4, tTMEM_STOREtS_x4) - cute.arch.fence_view_async_tmem_store() - si_handle.release() - - tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout) - acc_cons_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage) - c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) - c_pipe = pipeline.TmaStorePipeline.create(num_stages=self.num_c_stage, producer_group=c_grp) - acc_cons_st = utils.gemm.sm100.epilogue_tma_store( - self, tidx, warp_idx, tma_c, tCtO_base, sC, tCgC, - epi_tile, 0, const_expr(lambda x: x), (0,0,0), acc_cons_st, acc_pipe, c_pipe) - c_pipe.producer_tail() - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def test(): - torch.manual_seed(42) - # FMHA-matching dimensions: head_dim=64, seq=128 - # Q: (128, 64) K-major, K: (128, 64) K-major - # V: (64, 128) MN-major (transposed!) — FMHA requires v_major_mode=OperandMajorMode.MN - m, n, head_dim = 128, 128, 64 - q = torch.randn(m, head_dim, 1, dtype=torch.bfloat16, device='cuda') - kv = torch.randn(n, head_dim, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(head_dim, n, 1, dtype=torch.bfloat16, device='cuda') # Transposed! - c = torch.zeros(m, head_dim, 1, dtype=torch.bfloat16, device='cuda') - qf = q[:,:,0].float(); kvf = kv[:,:,0].float(); vf = v[:,:,0].float() - # Q@K^T = (128,128), P@V = (128,64) - ref = qf @ kvf.T @ vf - import cutlass.torch as ct - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(kv).mark_layout_dynamic(leading_dim=ct.get_leading_dim(kv)) - mV = ct.from_dlpack(v).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - kernel = StageBIdentitySoftmax(mma_tiler_mn=(128, 128), use_2cta_instrs=False, use_tma_store=True) - print('Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - print('Running...', flush=True) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - out = c[:,:,0].float() - cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - max_err = (out - ref).abs().max().item() - print('Stage B v20: FMHA-matching head_dim=64, proper V layout') - print(' Cosine: {:.6f}, Max error: {:.6f}'.format(cos, max_err)) - print(' {}'.format('PASS' if cos >= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_stage_b_v22.py b/tests/archive/test_stage_b_v22.py deleted file mode 100644 index bc711746..00000000 --- a/tests/archive/test_stage_b_v22.py +++ /dev/null @@ -1,314 +0,0 @@ -""" -Stage B v22: Q@K^T → S in TMEM, P@V → O in TMEM (identity softmax = S used as P) - -Based on working Stage A v2 pattern. Two MMAs on the MMA warp, no softmax pipeline. -P stays in TMEM between QK and PV — no copy, no packing, just use S directly as P. - -Bug 1 fix: V is MN-major, PV MMA uses v_major (OperandMajorMode.MN). -The PV MMA's A-operand (P) comes from TMEM (same as S accumulator). -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - - -class StageBIdentityKernel: - def __init__(self, mma_tiler_mn, use_2cta_instrs=False, use_tma_store=True): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16 - self.c_dtype = self.o_dtype # alias for epilogue_tma_store - self.use_2cta_instrs = use_2cta_instrs; self.use_tma_store = use_tma_store - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.epilog_sync_bar_id = 1; self.tmem_alloc_sync_bar_id = 2; self.tmem_dealloc_sync_bar_id = 3 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - self.pv_mma_tiler = (self.qk_mma_tiler[0], self.qk_mma_tiler[2], self.qk_mma_tiler[1]) - self.mma_tiler = self.qk_mma_tiler - print(f"[v22] qk_mma_tiler = {self.qk_mma_tiler}") - print(f"[v22] pv_mma_tiler = {self.pv_mma_tiler}") - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], - self.qk_mma_tiler[2], - ) - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, self.use_2cta_instrs, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.b_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.b_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tmem_s0_offset = 0 - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.b_dtype, b_smem) + cute.size_in_bytes(self.b_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q: cute.Tensor, k: cute.Tensor, v: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - self.q_dtype = q.element_type; self.b_dtype = k.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - print(f"[v22] a_major (Q) = {self.a_major}") - print(f"[v22] b_major (K) = {self.b_major}") - print(f"[v22] v_major (V) = {self.v_major}") - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.b_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - # BUG 1 FIX: PV MMA b_leading_mode = v_major (MN), NOT b_major (K) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.b_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, # Bug 1 fix - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=self.tmem_alloc_sync_bar_id, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.b_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.b_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - # V partition — pv_mma with pv_mma_tiler - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - # QK accumulator (S) in TMEM - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - # PV accumulator (O) in TMEM - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - # P tensor for PV MMA A-operand (from TMEM, same as S) - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - # P is at the same TMEM location as S (identity softmax: no copy/packing) - tOrP0 = tOrP - - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_stage_b_v22_bug1fix.py b/tests/archive/test_stage_b_v22_bug1fix.py deleted file mode 100644 index e6fdb8c3..00000000 --- a/tests/archive/test_stage_b_v22_bug1fix.py +++ /dev/null @@ -1,364 +0,0 @@ -""" -Stage B v22: Bug 1 Fix — V B-Operand Must Be MN-Major - -Fix over v20: PV MMA uses V's major mode (MN) instead of K's major mode (K). -V is shaped (head_dim, seq) = (64, 128) with strides (1, 64) → OperandMajorMode.MN. -K is shaped (seq, head_dim) = (128, 64) with strides (64, 1) → OperandMajorMode.K. - -These are DIFFERENT. The PV MMA's b_leading_mode MUST come from V, not from K. - -Also: separate TMA descriptor for V (already in v20), separate SMEM layout for V. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - - -class StageBIdentitySoftmax: - def __init__(self, mma_tiler_mn, use_2cta_instrs=False, use_tma_store=True): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16 - self.use_2cta_instrs = use_2cta_instrs; self.use_tma_store = use_tma_store - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.epilog_sync_bar_id = 1; self.tmem_alloc_sync_bar_id = 2; self.tmem_dealloc_sync_bar_id = 3 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - self.pv_mma_tiler = (self.qk_mma_tiler[0], self.qk_mma_tiler[2], self.qk_mma_tiler[1]) - self.mma_tiler = self.qk_mma_tiler - print(f"[v22] qk_mma_tiler = {self.qk_mma_tiler}") - print(f"[v22] pv_mma_tiler = {self.pv_mma_tiler}") - - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], - self.qk_mma_tiler[2], - ) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, self.use_2cta_instrs, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.a_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.b_dtype, 1) - # V uses pv_mma with MN-major B — its own SMEM layout - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.b_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - self.tmem_alloc_cols = s_cols + o_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.a_dtype, a_smem) + cute.size_in_bytes(self.b_dtype, b_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, v: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - self.a_dtype = a.element_type; self.b_dtype = b.element_type; self.c_dtype = c.element_type - self.a_major = LayoutEnum.from_tensor(a).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(b).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() # V's major mode — MN! - self.c_layout = LayoutEnum.from_tensor(c) - - print(f"[v22] a_major (Q) = {self.a_major}") - print(f"[v22] b_major (K) = {self.b_major}") - print(f"[v22] v_major (V) = {self.v_major}") - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - # BUG 1 FIX: PV MMA uses V's major mode (MN), NOT K's major mode (K) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.b_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_a, tma_ta = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - a, a_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_b, tma_tb = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - b, b_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - # Separate TMA for V — uses pv_mma and pv_mma_tiler - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_a, tma_ta, tma_b, tma_tb, tma_v, tma_tv, tma_c, tma_tc, - self.cluster_layout_vmnk, self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_a, mA, tma_b, mB, tma_v, mV, tma_c, mC, cl_vmnk, - a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_a); cpasync.prefetch_descriptor(tma_b); cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=self.tmem_alloc_sync_bar_id, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sA = smem.allocate_tensor(element_type=self.a_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sB = smem.allocate_tensor(element_type=self.b_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.b_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gA = cute.local_tile(mA, cute.slice_(self.mma_tiler, (None,0,None)), (None,None,None)) - gB = cute.local_tile(mB, cute.slice_(self.mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gA, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgA = qk_thr.partition_A(gA); tCgB = qk_thr.partition_B(gB); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsA, tAgA = cpasync.tma_partition(tma_a, 0, a_lay, cute.group_modes(sA,0,3), cute.group_modes(tCgA,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsB, tBgB = cpasync.tma_partition(tma_b, 0, b_lay, cute.group_modes(sB,0,3), cute.group_modes(tCgB,0,3)) - tAgA = tAgA[(None,0,None,0)]; tBgB = tBgB[(None,0,None,0)] - - # V partition — uses pv_mma, pv_mma_tiler, and tma_v - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrA = qk_mma.make_fragment_A(sA); tCrB = qk_mma.make_fragment_B(sB) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_a, tAgA[(None,h.count)], tAsA[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_b, tBgB[(None,h.count)], tBsB[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1BF16 packing (EXACT FMHA pattern) - tTMEM_STORErS_x4 = cute.make_rmem_tensor(tTMEM_STOREcS.shape, self.qk_acc_dtype) - tTMEM_STORErS_x4_e = cute.make_tensor( - cute.recast_ptr(tTMEM_STORErS_x4.iterator, dtype=self.q_dtype), - tTMEM_LOADrS.layout) - - frg_cnt = 4 - frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt - tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile)) - tTMEM_STORErS_x4_e_frg = cute.logical_divide( - tTMEM_STORErS_x4_e, cute.make_layout(frg_tile)) - for j in range(frg_cnt): - s_vec = tTMEM_LOADrS_frg[None, j].load() - tTMEM_STORErS_x4_e_frg[None, j].store(s_vec.to(self.q_dtype)) - - cute.copy(tiled_tmem_store, tTMEM_STORErS_x4, tTMEM_STOREtS_x4) - cute.arch.fence_view_async_tmem_store() - si_handle.release() - - tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout) - acc_cons_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage) - c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) - c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp) - acc_cons_st = utils.gemm.sm100.epilogue_tma_store( - self, tidx, warp_idx, tma_c, tCtO_base, sC, tCgC, - epi_tile, 0, const_expr(lambda x: x), (0,0,0), acc_cons_st, acc_pipe, c_pipe) - c_pipe.producer_tail() - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def test(): - torch.manual_seed(42) - # FMHA-matching dimensions: head_dim=64, seq=128 - # Q: (128, 64) K-major, K: (128, 64) K-major - # V: (64, 128) MN-major — FMHA requires v_major_mode=OperandMajorMode.MN - m, n, head_dim = 128, 128, 64 - q = torch.randn(m, head_dim, 1, dtype=torch.bfloat16, device='cuda') - kv = torch.randn(n, head_dim, 1, dtype=torch.bfloat16, device='cuda') - v_base = torch.randn(head_dim, n, dtype=torch.bfloat16, device='cuda') - v = v_base.as_strided((head_dim, n), (1, head_dim)).unsqueeze(-1) # MN-major: strides (1, 64) - c = torch.zeros(m, head_dim, 1, dtype=torch.bfloat16, device='cuda') - qf = q[:,:,0].float(); kvf = kv[:,:,0].float(); vf = v[:,:,0].float() - # Q@K^T = (128,128), P@V = (128,64) - ref = qf @ kvf.T @ vf.T - import cutlass.torch as ct - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(kv).mark_layout_dynamic(leading_dim=ct.get_leading_dim(kv)) - mV = ct.from_dlpack(v).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - kernel = StageBIdentitySoftmax(mma_tiler_mn=(128, 128), use_2cta_instrs=False, use_tma_store=True) - print('Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - print('Running...', flush=True) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - out = c[:,:,0].float() - cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - max_err = (out - ref).abs().max().item() - print('Stage B v22: Bug 1 fix — V MN-major') - print(' Cosine: {:.6f}, Max error: {:.6f}'.format(cos, max_err)) - print(' {}'.format('PASS' if cos >= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_stage_b_v23.py b/tests/archive/test_stage_b_v23.py deleted file mode 100644 index 6a9264ef..00000000 --- a/tests/archive/test_stage_b_v23.py +++ /dev/null @@ -1,364 +0,0 @@ -""" -Stage B v23: FMHA-matching test with head_dim=64, proper V layout. - -KEY INSIGHT: The A-fragment ((128,16),1,(4,2)):((65536,1),0,(16,64)) is SEQUENTIAL -when flattened in CuTe order: addr = m*65536 + k0 + 16*k1 + 64*k2 = m*65536 + k. -So the C-fragment composition store aliases the SAME TMEM as the A-fragment read. - -Previous -0.02 cosine was caused by V dimension mismatch: -pv_mma_tiler=(128,64,128) expects V with N=64 (head_dim), -but the square 128x128 test had V=K (N=128). - -This test: Q=(128,64), K=(128,64), V=(64,128), O=(128,64) - -FOOTGUN: St32x32bOp MUST use Float32, NOT BFloat16! -The 16-bit values are packed via recast_ptr view. -St32x32bOp(BFloat16) causes ILLEGAL MEMORY ACCESS. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - - -class StageBIdentitySoftmax: - def __init__(self, mma_tiler_mn, use_2cta_instrs=False, use_tma_store=True): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16 - self.use_2cta_instrs = use_2cta_instrs; self.use_tma_store = use_tma_store - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.epilog_sync_bar_id = 1; self.tmem_alloc_sync_bar_id = 2; self.tmem_dealloc_sync_bar_id = 3 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - self.pv_mma_tiler = (self.qk_mma_tiler[0], self.qk_mma_tiler[2], self.qk_mma_tiler[1]) - self.mma_tiler = self.qk_mma_tiler - print(f"[v23] qk_mma_tiler = {self.qk_mma_tiler}") - print(f"[v23] pv_mma_tiler = {self.pv_mma_tiler}") - - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], - self.qk_mma_tiler[2], - ) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, self.use_2cta_instrs, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.a_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.b_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.b_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - self.tmem_alloc_cols = s_cols + o_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.a_dtype, a_smem) + cute.size_in_bytes(self.b_dtype, b_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, v: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - self.a_dtype = a.element_type; self.b_dtype = b.element_type; self.c_dtype = c.element_type - self.a_major = LayoutEnum.from_tensor(a).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(b).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() # Bug 1: V MN-major - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.b_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, # Bug 1 fix: V MN-major - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_a, tma_ta = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - a, a_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_b, tma_tb = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - b, b_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - # Separate TMA for V (uses pv_mma_tiler and v_smem layout) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_a, tma_ta, tma_b, tma_tb, tma_v, tma_tv, tma_c, tma_tc, - self.cluster_layout_vmnk, self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_a, mA, tma_b, mB, tma_v, mV, tma_c, mC, cl_vmnk, - a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_a); cpasync.prefetch_descriptor(tma_b); cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=self.tmem_alloc_sync_bar_id, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sA = smem.allocate_tensor(element_type=self.a_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sB = smem.allocate_tensor(element_type=self.b_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.b_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gA = cute.local_tile(mA, cute.slice_(self.mma_tiler, (None,0,None)), (None,None,None)) - gB = cute.local_tile(mB, cute.slice_(self.mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gA, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgA = qk_thr.partition_A(gA); tCgB = qk_thr.partition_B(gB); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsA, tAgA = cpasync.tma_partition(tma_a, 0, a_lay, cute.group_modes(sA,0,3), cute.group_modes(tCgA,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsB, tBgB = cpasync.tma_partition(tma_b, 0, b_lay, cute.group_modes(sB,0,3), cute.group_modes(tCgB,0,3)) - tAgA = tAgA[(None,0,None,0)]; tBgB = tBgB[(None,0,None,0)] - - # V partition (from mV with pv_mma_tiler and tma_v) - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrA = qk_mma.make_fragment_A(sA); tCrB = qk_mma.make_fragment_B(sB) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_a, tAgA[(None,h.count)], tAsA[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_b, tBgB[(None,h.count)], tBsB[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_b, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1BF16 packing (EXACT FMHA pattern) - tTMEM_STORErS_x4 = cute.make_rmem_tensor(tTMEM_STOREcS.shape, self.qk_acc_dtype) - tTMEM_STORErS_x4_e = cute.make_tensor( - cute.recast_ptr(tTMEM_STORErS_x4.iterator, dtype=self.q_dtype), - tTMEM_LOADrS.layout) - - frg_cnt = 4 - frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt - tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile)) - tTMEM_STORErS_x4_e_frg = cute.logical_divide( - tTMEM_STORErS_x4_e, cute.make_layout(frg_tile)) - for j in range(frg_cnt): - s_vec = tTMEM_LOADrS_frg[None, j].load() - tTMEM_STORErS_x4_e_frg[None, j].store(s_vec.to(self.q_dtype)) - - cute.copy(tiled_tmem_store, tTMEM_STORErS_x4, tTMEM_STOREtS_x4) - cute.arch.fence_view_async_tmem_store() - si_handle.release() - - tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout) - acc_cons_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage) - c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) - c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp) - acc_cons_st = utils.gemm.sm100.epilogue_tma_store( - self, tidx, warp_idx, tma_c, tCtO_base, sC, tCgC, - epi_tile, 0, const_expr(lambda x: x), (0,0,0), acc_cons_st, acc_pipe, c_pipe) - c_pipe.producer_tail() - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def test(): - torch.manual_seed(42) - # FMHA-matching dimensions: head_dim=64, seq=128 - # Q: (128, 64) K-major, K: (128, 64) K-major - # V: (64, 128) MN-major (transposed!) — FMHA requires v_major_mode=OperandMajorMode.MN - m, n, head_dim = 128, 128, 64 - q = torch.randn(m, head_dim, 1, dtype=torch.bfloat16, device='cuda') - kv = torch.randn(n, head_dim, 1, dtype=torch.bfloat16, device='cuda') - v_base = torch.randn(head_dim, n, dtype=torch.bfloat16, device='cuda') - v = v_base.as_strided((head_dim, n), (1, head_dim)).unsqueeze(-1) # MN-major: strides (1, 64) - c = torch.zeros(m, head_dim, 1, dtype=torch.bfloat16, device='cuda') - qf = q[:,:,0].float(); kvf = kv[:,:,0].float(); vf = v_base.float() - # Q@K^T = (128,128), P@V = (128,64) - ref = qf @ kvf.T @ vf.T # P@V with V MN-major = Q@K^T @ V^T - import cutlass.torch as ct - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(kv).mark_layout_dynamic(leading_dim=ct.get_leading_dim(kv)) - mV = ct.from_dlpack(v).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - kernel = StageBIdentitySoftmax(mma_tiler_mn=(128, 128), use_2cta_instrs=False, use_tma_store=True) - print('Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - print('Running...', flush=True) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - out = c[:,:,0].float() - cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - max_err = (out - ref).abs().max().item() - print('Stage B v23: FMHA-matching head_dim=64, proper V layout') - print(' Cosine: {:.6f}, Max error: {:.6f}'.format(cos, max_err)) - print(' {}'.format('PASS' if cos >= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_stage_b_v24.py b/tests/archive/test_stage_b_v24.py deleted file mode 100644 index ee2a94f6..00000000 --- a/tests/archive/test_stage_b_v24.py +++ /dev/null @@ -1,377 +0,0 @@ -""" -Stage B v24: Q@K^T + identity softmax P packing + P@V -Bug 1 fix: V MN-major -Bug 2 fix: FP32->BF16 P packing (C-fragment composition store) -Pipeline fix: Use NamedBarrier instead of PipelineUmmaAsync for mma_si sync -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - - -class StageBIdentityKernel: - def __init__(self, mma_tiler_mn, use_2cta_instrs=False, use_tma_store=True): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = use_2cta_instrs; self.use_tma_store = use_tma_store - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.epilog_sync_bar_id = 1; self.tmem_alloc_sync_bar_id = 2; self.tmem_dealloc_sync_bar_id = 3 - self.num_c_stage = 2 - # Named barrier IDs for mma_si coordination - self.scores_ready_bar_id = 4 # MMA warp signals S ready - self.softmax_done_bar_id = 5 # Epilogue warps signal softmax done - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - self.pv_mma_tiler = (self.qk_mma_tiler[0], self.qk_mma_tiler[2], self.qk_mma_tiler[1]) - self.mma_tiler = self.qk_mma_tiler - print(f"[v24] qk_mma_tiler = {self.qk_mma_tiler}") - print(f"[v24] pv_mma_tiler = {self.pv_mma_tiler}") - - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], - self.qk_mma_tiler[2], - ) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, self.use_2cta_instrs, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.b_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.b_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 # P region in TMEM (for BF16 packing) - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.b_dtype, b_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q: cute.Tensor, k: cute.Tensor, v: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - self.q_dtype = q.element_type; self.b_dtype = k.element_type - self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - print(f"[v24] a_major (Q) = {self.a_major}") - print(f"[v24] b_major (K) = {self.b_major}") - print(f"[v24] v_major (V) = {self.v_major}") - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.b_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.b_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=self.tmem_alloc_sync_bar_id, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - # Named barriers for MMA ↔ softmax coordination - scores_ready_bar = pipeline.NamedBarrier( - barrier_id=self.scores_ready_bar_id, - num_threads=32 * (1 + len(self.epilogue_warp_id))) # MMA warp + epilogue warps - softmax_done_bar = pipeline.NamedBarrier( - barrier_id=self.softmax_done_bar_id, - num_threads=32 * (1 + len(self.epilogue_warp_id))) # MMA warp + epilogue warps - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.b_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.b_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - # P from TMEM — at tmem_p0_offset, read by PV MMA as A-operand - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - tmem.wait_for_alloc() - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1 MMA computes P @ V^T - ref = qf @ kf.T @ vf.T # (128,64) - - import cutlass.torch as ct - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - kernel = StageBIdentityKernel(mma_tiler_mn=(128, 128), use_2cta_instrs=False, use_tma_store=True) - print('Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - print('Running...', flush=True) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - out = c[:,:,0].float() - cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - max_err = (out - ref).abs().max().item() - print('Stage B v24: Q@K^T + identity softmax (NamedBarrier) + P@V (V MN-major)') - print(' Cosine: {:.6f}, Max error: {:.6f}'.format(cos, max_err)) - print(' {}'.format('PASS' if cos >= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_stage_b_v25.py b/tests/archive/test_stage_b_v25.py deleted file mode 100644 index da97566f..00000000 --- a/tests/archive/test_stage_b_v25.py +++ /dev/null @@ -1,380 +0,0 @@ -""" -Stage B v25: Q@K^T + identity softmax P packing + P@V -Uses PipelineAsync (not Umma) for mma_si sync — two separate one-shot pipelines. -Bug 1 fix: V MN-major. -Bug 2 fix: FP32→BF16 P packing (C-fragment composition store). -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - - -class StageBIdentityKernel: - def __init__(self, mma_tiler_mn, use_2cta_instrs=False, use_tma_store=True): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = use_2cta_instrs; self.use_tma_store = use_tma_store - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.epilog_sync_bar_id = 1; self.tmem_alloc_sync_bar_id = 2; self.tmem_dealloc_sync_bar_id = 3 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - self.pv_mma_tiler = (self.qk_mma_tiler[0], self.qk_mma_tiler[2], self.qk_mma_tiler[1]) - self.mma_tiler = self.qk_mma_tiler - print(f"[v25] qk_mma_tiler = {self.qk_mma_tiler}") - print(f"[v25] pv_mma_tiler = {self.pv_mma_tiler}") - - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], - self.qk_mma_tiler[2], - ) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, self.use_2cta_instrs, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.b_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.b_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.b_dtype, b_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q: cute.Tensor, k: cute.Tensor, v: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - self.q_dtype = q.element_type; self.b_dtype = k.element_type - self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - print(f"[v25] a_major (Q) = {self.a_major}") - print(f"[v25] b_major (K) = {self.b_major}") - print(f"[v25] v_major (V) = {self.v_major}") - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.b_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.b_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - # Two separate mbarriers for mma→softmax and softmax→mma signaling - scores_ready_bar: cute.struct.MemRange[cutlass.Int64, 2] # MMA→softmax - softmax_done_bar: cute.struct.MemRange[cutlass.Int64, 2] # softmax→MMA - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=self.tmem_alloc_sync_bar_id, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - # PipelineAsync for scores_ready (MMA → softmax) - scores_ready_prod, scores_ready_cons = pipeline.PipelineAsync.create( - barrier_storage=st.scores_ready_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - # PipelineAsync for softmax_done (softmax → MMA) - softmax_done_prod, softmax_done_cons = pipeline.PipelineAsync.create( - barrier_storage=st.softmax_done_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - ).make_participants() - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.b_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.b_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_stage_b_v26.py b/tests/archive/test_stage_b_v26.py deleted file mode 100644 index 9d933015..00000000 --- a/tests/archive/test_stage_b_v26.py +++ /dev/null @@ -1,367 +0,0 @@ -""" -Stage B v26: Q@K^T + identity softmax P packing + P@V -Uses SMEM spin-wait flags for mma↔softmax sync (no PipelineUmmaAsync). -Bug 1 fix: V MN-major. -Bug 2 fix: FP32→BF16 P packing (C-fragment composition store). -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - - -class StageBIdentityKernel: - def __init__(self, mma_tiler_mn, use_2cta_instrs=False, use_tma_store=True): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = use_2cta_instrs; self.use_tma_store = use_tma_store - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.epilog_sync_bar_id = 1; self.tmem_alloc_sync_bar_id = 2; self.tmem_dealloc_sync_bar_id = 3 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - self.pv_mma_tiler = (self.qk_mma_tiler[0], self.qk_mma_tiler[2], self.qk_mma_tiler[1]) - self.mma_tiler = self.qk_mma_tiler - print(f"[v26] qk_mma_tiler = {self.qk_mma_tiler}") - print(f"[v26] pv_mma_tiler = {self.pv_mma_tiler}") - - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], - self.qk_mma_tiler[2], - ) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, self.use_2cta_instrs, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.b_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.b_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.b_dtype, b_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q: cute.Tensor, k: cute.Tensor, v: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - self.q_dtype = q.element_type; self.b_dtype = k.element_type - self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.b_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.b_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - # Use acc_pipe's barrier for BOTH QK→softmax and PV→epilogue signaling - # The trick: re-acquire acc_pipe on the producer side after softmax - # This works because PipelineUmmaAsync with 1 stage blocks producer - # until consumer releases - - tmem_bar = pipeline.NamedBarrier(barrier_id=self.tmem_alloc_sync_bar_id, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.b_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.b_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP (no tmem.wait_for_alloc!) ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_stage_b_v27.py b/tests/archive/test_stage_b_v27.py deleted file mode 100644 index c544d363..00000000 --- a/tests/archive/test_stage_b_v27.py +++ /dev/null @@ -1,375 +0,0 @@ -""" -Stage B v27: Based on v20 with Bug 1 fix (V MN-major). -Fixes over v20: - 1. V MN-major + pv_mma uses v_major - 2. PipelineTmaStore (not TmaStorePipeline) - 3. TMA warp does NOT call tmem.wait_for_alloc - 4. mma_si PipelineUmmaAsync without cta_layout_vmnk (like FMHA) -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - - -class StageBIdentitySoftmax: - def __init__(self, mma_tiler_mn, use_2cta_instrs=False, use_tma_store=True): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = use_2cta_instrs; self.use_tma_store = use_tma_store - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.epilog_sync_bar_id = 1; self.tmem_alloc_sync_bar_id = 2; self.tmem_dealloc_sync_bar_id = 3 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - self.pv_mma_tiler = (self.qk_mma_tiler[0], self.qk_mma_tiler[2], self.qk_mma_tiler[1]) - self.mma_tiler = self.qk_mma_tiler - print(f"[v27] qk_mma_tiler = {self.qk_mma_tiler}") - print(f"[v27] pv_mma_tiler = {self.pv_mma_tiler}") - - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], - self.qk_mma_tiler[2], - ) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, self.use_2cta_instrs, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.b_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.b_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - # ⛔⛔⛔ CRITICAL: num_tma_load_bytes MUST include ALL TMA-loaded tensors (Q + K + V). - # Missing V bytes causes TMA barrier tx-count underflow → DEADLOCK. - # See FOOTGUN #0 in README. - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) - + cute.size_in_bytes(self.b_dtype, b_smem) - + cute.size_in_bytes(self.b_dtype, v_smem) # ← DO NOT FORGET V - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q: cute.Tensor, k: cute.Tensor, v: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - self.q_dtype = q.element_type; self.b_dtype = k.element_type - self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - print(f"[v27] a_major (Q) = {self.a_major}") - print(f"[v27] b_major (K) = {self.b_major}") - print(f"[v27] v_major (V) = {self.v_major}") - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.b_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.b_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - # mma_si pipeline — NO cta_layout_vmnk (matches FMHA) - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=self.tmem_alloc_sync_bar_id, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.b_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.b_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP (no tmem.wait_for_alloc) ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_stage_b_v28.py b/tests/archive/test_stage_b_v28.py deleted file mode 100644 index 3e88a0c5..00000000 --- a/tests/archive/test_stage_b_v28.py +++ /dev/null @@ -1,377 +0,0 @@ -""" -Stage B v28: Bug 4 fix — PV MMA with correct (128,64) tiler. - -Key fixes over v27: - 1. PV MMA created with pv_mma_tiler[:2] = (128, 64) instead of mma_tiler_mn = (128, 128) - -> A-fragment N_MMA=64 matches packed P layout (64 FP32 columns) - 2. epi_tile = pv_mma_tiler[:2] like FMHA (not from QK cta_tile) - 3. PV ACCUMULATE=False on first tile (FMHA: kphase_idx != 0) - 4. gC (output) tiling uses pv_mma_tiler for (M, head_dim) not qk_mma_tiler - 5. cta_tile_shape_mnk set to PV cta tile before epilogue -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - - -class StageBIdentitySoftmax: - def __init__(self, mma_tiler_mn, use_2cta_instrs=False, use_tma_store=True): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = use_2cta_instrs; self.use_tma_store = use_tma_store - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.epilog_sync_bar_id = 1; self.tmem_alloc_sync_bar_id = 2; self.tmem_dealloc_sync_bar_id = 3 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - self.pv_mma_tiler = (self.qk_mma_tiler[0], self.qk_mma_tiler[2], self.qk_mma_tiler[1]) - self.mma_tiler = self.qk_mma_tiler - print(f"[v28] qk_mma_tiler = {self.qk_mma_tiler}") - print(f"[v28] pv_mma_tiler = {self.pv_mma_tiler}") - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - - # cta_tile_shape_mnk from PV for epilogue - self.cta_tile_shape_mnk = ( - self.pv_mma_tiler[0] // cute.size(pv_mma.thr_id.shape), - self.pv_mma_tiler[1], - self.pv_mma_tiler[2], - ) - # FMHA: epi_tile from PV cta_tile, not QK - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, self.use_2cta_instrs, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.b_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.b_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.b_dtype, b_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q: cute.Tensor, k: cute.Tensor, v: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - self.q_dtype = q.element_type; self.b_dtype = k.element_type - self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.b_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - - # Compute pv_mma_tiler[:2] BEFORE creating pv_mma - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - pv_mma_tiler = (qk_mma_tiler[0], qk_mma_tiler[2], qk_mma_tiler[1]) - - # BUG 4 FIX: pv_mma_tiler[:2] = (128, 64) not (128, 128) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.b_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, pv_mma_tiler[:2], tcgen05.OperandSource.TMEM) - - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=self.tmem_alloc_sync_bar_id, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.b_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.b_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - # gC uses pv_mma_tiler output shape (M, head_dim) - gC = cute.local_tile(mC, cute.slice_(self.pv_mma_tiler, (None, None, 0)), (None, None, None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - # C partitioned by QK thr (for TMA store, FMHA uses the same pattern) - tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ===== TMA LOAD WARP ===== - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_stage_b_v29.py b/tests/archive/test_stage_b_v29.py deleted file mode 100644 index 502db5af..00000000 --- a/tests/archive/test_stage_b_v29.py +++ /dev/null @@ -1,365 +0,0 @@ -""" -Stage B v29: Pad V to (128,128), use (128,128) PV MMA. -V_padded = [V_real; zeros(64,128)] -O = P @ V_padded → (128, 128), but O[:, :64] = P @ V_real = correct -We only check the first 64 columns. - -This avoids the (128,64) PV MMA deadlock and TMEM alias issues. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - - -class StageBIdentitySoftmax: - def __init__(self, mma_tiler_mn, use_2cta_instrs=False, use_tma_store=True): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = use_2cta_instrs; self.use_tma_store = use_tma_store - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.epilog_sync_bar_id = 1; self.tmem_alloc_sync_bar_id = 2; self.tmem_dealloc_sync_bar_id = 3 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - self.pv_mma_tiler = (self.qk_mma_tiler[0], self.qk_mma_tiler[1], self.qk_mma_tiler[1]) - self.mma_tiler = self.qk_mma_tiler - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], - self.qk_mma_tiler[2], - ) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, self.use_2cta_instrs, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.b_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.b_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - # ⛔⛔⛔ CRITICAL: num_tma_load_bytes MUST include ALL TMA-loaded tensors (Q + K + V). - # Missing V bytes causes TMA barrier tx-count underflow → DEADLOCK. - # See FOOTGUN #0 in README. - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) - + cute.size_in_bytes(self.b_dtype, b_smem) - + cute.size_in_bytes(self.b_dtype, v_smem) # ← DO NOT FORGET V - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q: cute.Tensor, k: cute.Tensor, v: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - self.q_dtype = q.element_type; self.b_dtype = k.element_type - self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.b_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.b_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=self.tmem_alloc_sync_bar_id, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.b_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.b_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP (no tmem.wait_for_alloc) ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_stage_b_v3.py b/tests/archive/test_stage_b_v3.py deleted file mode 100644 index 0884a7a5..00000000 --- a/tests/archive/test_stage_b_v3.py +++ /dev/null @@ -1,375 +0,0 @@ -""" -Stage B: Two MMAs (Q@K^T then P@V) with a_source=TMEM for MMA2. - -Architecture (following NVIDIA's fmha.py reference): - MMA1: Q @ K^T → tmem_scores (a_source=SMEM, accumulate=False) - MMA2: P @ V → tmem_output (a_source=TMEM, accumulate=True) - P fragment: constructed with P TMEM layout from make_smem_layout_A(pv_tiled_mma, ...) - P TMEM address: fragment.iterator + (acc_width/a_width) * tmem_p_offset - -Reference: output = Q @ K^T @ V (no softmax, P = raw scores) -""" -import torch -import cutlass -import cutlass.cute as cute -import cutlass.utils as utils -import cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -import cuda.bindings.driver as cuda - - -class StageBKernel: - def __init__(self, mma_tiler_mn, use_2cta_instrs=False): - self.acc_dtype = Float32 - self.use_2cta_instrs = use_2cta_instrs - self.mma_tiler_mn = mma_tiler_mn - self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4 - self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.epilog_sync_bar_id = 1 - - def _setup(self, qk_mma, pv_mma, a_dtype, b_dtype, c_dtype, a_major, b_major, c_layout): - self.a_dtype = a_dtype - self.b_dtype = b_dtype - self.c_dtype = c_dtype - self.c_layout = c_layout - self.use_2cta_instrs = False - - # QK MMA tiler - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - # PV MMA tiler (same M,N but potentially different K) - pv_inst_k = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (*self.mma_tiler_mn, pv_inst_k * 4) - # Use QK tiler for the overall tiler (A/B come from QK layout) - self.mma_tiler = self.qk_mma_tiler - - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0], self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - - # Epilogue tile from PV MMA (output is O, same shape as PV MMA's C) - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, False, c_layout, c_dtype) - - self.num_ab_stage = 1 - self.num_acc_stage = 1 - self.num_c_stage = 2 - - # SMEM layouts for QK MMA (Q and K) - self.q_smem_layout = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, a_dtype, 1) - self.k_smem_layout = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, b_dtype, 1) - # SMEM layout for V (from PV MMA) - self.v_smem_layout = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, b_dtype, 1) - # TMEM layout for P (from PV MMA's A operand — this is the KEY) - self.p_tmem_layout = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, a_dtype, 1) - # C/Output SMEM layout - self.c_smem_layout = utils.sm100.make_smem_layout_epi(c_dtype, c_layout, self.epi_tile, 2) - - # TMEM allocation: two regions - # Region 0: scores (Q@K^T accumulator, QK MMA's C layout) - acc_shape = qk_mma.partition_shape_C(self.mma_tiler_mn) - tCtAcc_fake = qk_mma.make_fragment_C(cute.append(acc_shape, 1)) - self.num_tmem_cols_scores = utils.get_num_tmem_alloc_cols(tCtAcc_fake, arch="sm_100") - # Region 1: output (P@V accumulator, PV MMA's C layout) - acc_shape_pv = pv_mma.partition_shape_C(self.mma_tiler_mn) - tCtO_fake = pv_mma.make_fragment_C(cute.append(acc_shape_pv, 1)) - self.num_tmem_cols_output = utils.get_num_tmem_alloc_cols(tCtO_fake, arch="sm_100") - # Total - self.total_tmem_cols = max(self.num_tmem_cols_scores + self.num_tmem_cols_output, 256) - # tmem_p_offset: P TMEM starts at 0 (same as scores, P replaces scores in TMEM) - self.tmem_p_offset = 0 - - # TMA load bytes - q_smem = cute.slice_(self.q_smem_layout, (None, None, None, 0)) - k_smem = cute.slice_(self.k_smem_layout, (None, None, None, 0)) - self.num_tma_bytes = ( - cute.size_in_bytes(a_dtype, q_smem) + - cute.size_in_bytes(b_dtype, k_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, - stream: cuda.CUstream): - a_dtype = a.element_type - b_dtype = b.element_type - c_dtype = c.element_type - a_major = LayoutEnum.from_tensor(a).mma_major_mode() - b_major = LayoutEnum.from_tensor(b).mma_major_mode() - c_layout = LayoutEnum.from_tensor(c) - - # QK MMA: Q @ K^T, A from SMEM - qk_mma = utils.sm100.make_trivial_tiled_mma( - a_dtype, b_dtype, a_major, b_major, - self.acc_dtype, self.cta_group, self.mma_tiler_mn, - tcgen05.OperandSource.SMEM, - ) - - # PV MMA: P @ V, A from TMEM, P is K-major - # Following NVIDIA fmha.py: P (intermediate) is K-major and from TMEM - p_major = cute.nvgpu.OperandMajorMode.K - pv_mma = utils.sm100.make_trivial_tiled_mma( - a_dtype, b_dtype, p_major, b_major, - self.acc_dtype, self.cta_group, self.mma_tiler_mn, - tcgen05.OperandSource.TMEM, - ) - - self._setup(qk_mma, pv_mma, a_dtype, b_dtype, c_dtype, a_major, b_major, c_layout) - - q_smem = cute.slice_(self.q_smem_layout, (None, None, None, 0)) - k_smem = cute.slice_(self.k_smem_layout, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_layout, (None, None, None, 0)) - - tma_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - a, q_smem, self.qk_mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - b, k_smem, self.qk_mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tensor_v = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - b, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_layout, mode=[0, 1]) - tma_c, tma_tensor_c = cpasync.make_tiled_tma_atom( - cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel( - qk_mma, pv_mma, - tma_a, tma_tensor_a, tma_b, tma_tensor_b, - tma_v, tma_tensor_v, - tma_c, tma_tensor_c, self.cluster_layout_vmnk, - self.q_smem_layout, self.k_smem_layout, - self.v_smem_layout, self.p_tmem_layout, - self.c_smem_layout, self.epi_tile, - ).launch(grid=(1,1,1), block=[self.threads_per_cta, 1, 1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, - tma_q, mQ, tma_k, mK, - tma_v, mV, - tma_c, mC, cl_vmnk, - q_smem_layout, k_smem_layout, - v_smem_layout, p_tmem_layout, - c_smem_layout, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q) - cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v) - cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, 2] - dealloc_bar: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator() - storage = smem.allocate(SS) - - ab_prod, ab_cons = pipeline.PipelineTmaUmma.create( - barrier_storage=storage.ab_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_bytes, cta_layout_vmnk=cl_vmnk, - defer_sync=True).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=storage.acc_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 128), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, num_threads=160) - tmem = utils.TmemAllocator(storage.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=0, is_two_cta=False, - two_cta_tmem_dealloc_mbar_ptr=storage.dealloc_bar.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=BFloat16, layout=q_smem_layout.outer, - byte_alignment=128, swizzle=q_smem_layout.inner) - sK = smem.allocate_tensor(element_type=BFloat16, layout=k_smem_layout.outer, - byte_alignment=128, swizzle=k_smem_layout.inner) - sV = smem.allocate_tensor(element_type=BFloat16, layout=v_smem_layout.outer, - byte_alignment=128, swizzle=v_smem_layout.inner) - sC = smem.allocate_tensor(element_type=BFloat16, layout=c_smem_layout.outer, - byte_alignment=128, swizzle=c_smem_layout.inner) - - # Q and K TMA partition - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ) - tCgK = qk_thr.partition_B(gK) - tCgC = qk_thr.partition_C(gC) - - a_cta = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_cta, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_cta = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tAsK, tAgK = cpasync.tma_partition(tma_k, 0, b_cta, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)] - tAgK = tAgK[(None,0,None,0)] - - # QK MMA fragments - tCrQ = qk_mma.make_fragment_A(sQ) - tCrK = qk_mma.make_fragment_B(sK) - - # PV MMA fragments - tCrV = pv_mma.make_fragment_B(sV) - - # TMEM accumulator fake for QK (scores) - acc_shape_qk = qk_mma.partition_shape_C(self.mma_tiler_mn) - tCtS_fake = qk_mma.make_fragment_C(cute.append(acc_shape_qk, 1)) - - # TMEM accumulator fake for PV (output) - acc_shape_pv = pv_mma.partition_shape_C(self.mma_tiler_mn) - tCtO_fake = pv_mma.make_fragment_C(cute.append(acc_shape_pv, 1)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ── TMA LOAD WARP ── - if warp_idx == self.tma_warp_id: - ab_prod.reset() - peek = ab_prod.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_prod.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tAgK[(None,h.count)], tAsK[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else "FAIL")) - return cos - -if __name__ == "__main__": - test_stage_b() diff --git a/tests/archive/test_stage_b_v30.py b/tests/archive/test_stage_b_v30.py deleted file mode 100644 index 80184627..00000000 --- a/tests/archive/test_stage_b_v30.py +++ /dev/null @@ -1,382 +0,0 @@ -""" -Minimal PV-only test: Load P from GMEM to TMEM via QK-style MMA, then PV from TMEM. -Step 1: QK MMA writes FP32 S to TMEM (we know this works) -Step 2: Softmax packing writes BF16 P to TMEM (test this) -Step 3: PV MMA reads BF16 P from TMEM and V from SMEM, produces O - -But to isolate the bug, let me test just the PV MMA in isolation. -I'll write known BF16 values to TMEM using the softmax packing path, -then immediately read them back using the PV A-fragment path, -and compare. - -Actually, the simplest isolation test: -1. Do QK MMA to get S in TMEM (cosine 0.999999 verified) -2. Do softmax packing: S → P in TMEM (at offset 32) -3. Skip PV entirely — read P from TMEM using the C-fragment composition LOAD path -4. Output P to GMEM and compare against S.to(BF16) - -This tests whether the softmax packing writes P correctly to the same TMEM -that the PV would read from. - -But we can't easily read P from TMEM using the standard epilogue path -because the epilogue expects FP32 accumulator data. - -Alternative: Use the PV MMA with V=I (identity). If P is correct, -then P @ I = P. But V needs to be MN-major and (128, 128), not (128, 64). -The output would be (128, 128) which doesn't match our (128, 64) c tensor. - -Let me use V that selects the first 64 columns: V[k, n] = delta(k, n) for k in [0,63]. -This gives P @ V = P[:, :64], and the output is (128, 64). -But V is (128, 128) in the MMA K,N dims. V[k, n] for k in [0,127], n in [0,63]. -Hmm, this is getting complicated. Let me just do the identity approach with a (128, 128) output. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class PvDiagKernel: - """QK + softmax packing + PV with V=I to isolate PV MMA correctness. - Output should be P = S.to(BF16), i.e. (Q@K^T).bfloat16() - With V=I, O = P @ I = P. - But V is (K=128, N=128) in the MMA. We need a 128x128 identity in MN-major. - Output tensor is (128, 128). - """ - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False # needed by epilogue_tma_store - self.epilog_sync_bar_id = 1 # needed by epilogue_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - # PV with V=I: output is (128, 128), same as QK - self.pv_mma_tiler = (self.qk_mma_tiler[0], self.qk_mma_tiler[1], self.qk_mma_tiler[1]) - # pv_mma_tiler = (128, 128, 128) since V is 128x128 - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - # PV with 128x128 output (V=I) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_stage_b_v4.py b/tests/archive/test_stage_b_v4.py deleted file mode 100644 index a7990222..00000000 --- a/tests/archive/test_stage_b_v4.py +++ /dev/null @@ -1,259 +0,0 @@ -"""Stage B v4: Two MMAs with a_source=TMEM, minimal pipeline. -No softmax. P = raw scores. Reference: Q @ K^T @ V.""" -import torch -import cutlass -import cutlass.cute as cute -import cutlass.utils as utils -import cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -import cuda.bindings.driver as cuda - -class StageBKernel: - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32 - self.mma_tiler_mn = mma_tiler_mn - self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.use_2cta_instrs = False - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4 - self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.epilog_sync_bar_id = 1 - self.num_c_stage = 2 - - @cute.jit - def __call__(self, a, b, c, stream): - a_dtype = a.element_type - b_dtype = b.element_type - c_dtype = c.element_type - a_major = LayoutEnum.from_tensor(a).mma_major_mode() - b_major = LayoutEnum.from_tensor(b).mma_major_mode() - c_layout = LayoutEnum.from_tensor(c) - self.a_dtype = a_dtype - self.b_dtype = b_dtype - self.c_dtype = c_dtype - self.c_layout = c_layout - - qk_mma = utils.sm100.make_trivial_tiled_mma( - a_dtype, b_dtype, a_major, b_major, - self.acc_dtype, self.cta_group, self.mma_tiler_mn, - tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma( - a_dtype, b_dtype, cute.nvgpu.OperandMajorMode.K, b_major, - self.acc_dtype, self.cta_group, self.mma_tiler_mn, - tcgen05.OperandSource.TMEM) - - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - pv_inst_k = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (*self.mma_tiler_mn, pv_inst_k * 4) - self.mma_tiler = self.qk_mma_tiler - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0], self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, c_layout, c_dtype) - - q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, a_dtype, 1) - k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, b_dtype, 1) - p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, a_dtype, 1) - c_smem_s = utils.sm100.make_smem_layout_epi(c_dtype, c_layout, self.epi_tile, 2) - - acc_shape = qk_mma.partition_shape_C(self.mma_tiler_mn) - tCtS_fake = qk_mma.make_fragment_C(cute.append(acc_shape, 1)) - self.num_tmem_cols_scores = utils.get_num_tmem_alloc_cols(tCtS_fake, arch="sm_100") - acc_shape_pv = pv_mma.partition_shape_C(self.mma_tiler_mn) - tCtO_fake = pv_mma.make_fragment_C(cute.append(acc_shape_pv, 1)) - self.num_tmem_cols_output = utils.get_num_tmem_alloc_cols(tCtO_fake, arch="sm_100") - self.total_tmem_cols = max(self.num_tmem_cols_scores + self.num_tmem_cols_output, 256) - - q_smem = cute.slice_(q_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(k_smem_s, (None, None, None, 0)) - self.num_tma_bytes = (cute.size_in_bytes(a_dtype, q_smem) + cute.size_in_bytes(b_dtype, k_smem)) * cute.size(qk_mma.thr_id.shape) - - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - a, q_smem, self.qk_mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - b, k_smem, self.qk_mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom( - cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_c, tma_tc, - self.cluster_layout_vmnk, q_smem_s, k_smem_s, p_tmem_s, c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[192,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_c, mC, cl_vmnk, - q_smem_s, k_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q) - cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, 2] - dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator() - st = smem.allocate(SS) - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 128), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2, num_threads=160) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=0, is_two_cta=False, - two_cta_tmem_dealloc_mbar_ptr=st.dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=BFloat16, layout=q_smem_s.outer, byte_alignment=128, swizzle=q_smem_s.inner) - sK = smem.allocate_tensor(element_type=BFloat16, layout=k_smem_s.outer, byte_alignment=128, swizzle=k_smem_s.inner) - sC = smem.allocate_tensor(element_type=BFloat16, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ) - tCgK = qk_thr.partition_B(gK) - tCgC = qk_thr.partition_C(gC) - - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tAsK, tAgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)] - tAgK = tAgK[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ) - tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sK) - - acc_shape = qk_mma.partition_shape_C(self.mma_tiler_mn) - tCtS_fake = qk_mma.make_fragment_C(cute.append(acc_shape, 1)) - acc_shape_pv = pv_mma.partition_shape_C(self.mma_tiler_mn) - tCtO_fake = pv_mma.make_fragment_C(cute.append(acc_shape_pv, 1)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # TMA warp - if warp_idx == self.tma_warp_id: - ab_p.reset() - peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tAgK[(None,h.count)], tAsK[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_stage_b_v5.py b/tests/archive/test_stage_b_v5.py deleted file mode 100644 index 19527382..00000000 --- a/tests/archive/test_stage_b_v5.py +++ /dev/null @@ -1,341 +0,0 @@ -""" -Stage B: Two MMAs + Identity Softmax with Layout Transform - -Following fmha.py's synchronization pattern: - - MMA↔softmax sync via PipelineUmmaAsync (mma_si pipeline) - - MMA produces scores (after QK), softmax consumes - - Softmax produces P, MMA re-acquires (before PV) - - Identity softmax: tcgen05.ld from C-layout → F32→BF16 → tcgen05.st to A-layout - -Reference: output = (Q @ K^T) @ V -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -import cuda.bindings.driver as cuda - - -class StageBIdentitySoftmax: - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32 - self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16 - self.o_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn - self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.use_2cta_instrs = False - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4 - self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - # TMEM offsets (fmha.py for 128x128) - self.tmem_s0_offset = 0 - self.tmem_o0_offset = 256 - self.tmem_p0_offset = 32 - self.tmem_alloc_cols = 512 - self.epilog_sync_bar_id = 1 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - pv_inst_k = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (*self.mma_tiler_mn, pv_inst_k * 4) - self.mma_tiler = self.qk_mma_tiler - self.cta_tile_shape_mnk = tuple(self.qk_mma_tiler) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, self.a_dtype, 1) - self.k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, self.b_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - acc_shape = qk_mma.partition_shape_C(self.mma_tiler_mn) - tCtS_fake = qk_mma.make_fragment_C(cute.append(acc_shape, 1)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(tCtS_fake, arch="sm_100") - - q_smem = cute.slice_(self.q_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.k_smem_s, (None, None, None, 0)) - self.num_tma_bytes = (cute.size_in_bytes(self.a_dtype, q_smem) + cute.size_in_bytes(self.b_dtype, k_smem)) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, a, b, c, stream): - self.a_dtype = a.element_type; self.b_dtype = b.element_type; self.c_dtype = c.element_type - self.a_major = LayoutEnum.from_tensor(a).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(b).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, self.a_major, self.b_major, self.acc_dtype, self.cta_group, self.mma_tiler_mn, - tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, cute.nvgpu.OperandMajorMode.K, self.b_major, self.acc_dtype, self.cta_group, self.mma_tiler_mn, - tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.q_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.k_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - a, q_smem, self.qk_mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - b, k_smem, self.qk_mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_c, tma_tc, - self.cluster_layout_vmnk, self.q_smem_s, self.k_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[192,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_c, mC, cl_vmnk, - q_smem_s, k_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q) - cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, 2] # AB pipeline (1 stage) - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] # MMA↔softmax pipeline (1 stage) - acc_bar: cute.struct.MemRange[cutlass.Int64, 2] # ACC pipeline (1 stage) - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator() - st = smem.allocate(SS) - - # AB pipeline (TMA load) - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - # MMA↔softmax pipeline (following fmha.py's mma_s0 pattern) - # Producer = MMA warp (after QK: commit scores; before PV: re-acquire P) - # Consumer = softmax warps (wait for scores, process, release P) - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 128), - cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - # ACC pipeline (PV output → epilogue) - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 128), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - # TMEM allocator - tmem_bar = pipeline.NamedBarrier(barrier_id=2, num_threads=160) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=0, is_two_cta=False, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=BFloat16, layout=q_smem_s.outer, byte_alignment=128, swizzle=q_smem_s.inner) - sK = smem.allocate_tensor(element_type=BFloat16, layout=k_smem_s.outer, byte_alignment=128, swizzle=k_smem_s.inner) - sC = smem.allocate_tensor(element_type=BFloat16, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tAsK, tAgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tAgK = tAgK[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ) - tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sK) - - # TMEM tensors - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler_mn) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_mma.partition_shape_C(self.mma_tiler_mn) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - # P fragment for PV MMA (A-layout from TMEM) - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_mma.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - # Fake acc for epilogue - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, 1)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, 1)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ── TMA WARP ── - if warp_idx == self.tma_warp_id: - ab_p.reset() - peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tAgK[(None,h.count)], tAsK[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1=0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_stage_b_v6.py b/tests/archive/test_stage_b_v6.py deleted file mode 100644 index cd8a32ea..00000000 --- a/tests/archive/test_stage_b_v6.py +++ /dev/null @@ -1,342 +0,0 @@ -""" -Stage B: Two MMAs + Identity Softmax with Layout Transform - -Following fmha.py's softmax_step pattern exactly. - -Architecture: - MMA1: Q @ K^T → tmem_scores (a_source=SMEM, accumulate=False) - Identity softmax: tcgen05.ld from C-layout → convert F32→BF16 → tcgen05.st to A-layout - MMA2: P @ V → tmem_output (a_source=TMEM, accumulate=True) - -Reference: output = (Q @ K^T) @ V (no softmax, P = raw scores) -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -import cuda.bindings.driver as cuda - - -class StageBIdentitySoftmax: - def __init__(self, mma_tiler_mn, use_2cta_instrs=False, use_tma_store=True): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16 - self.use_2cta_instrs = use_2cta_instrs; self.use_tma_store = use_tma_store - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.epilog_sync_bar_id = 1; self.tmem_alloc_sync_bar_id = 2; self.tmem_dealloc_sync_bar_id = 3 - self.num_c_stage = 2 - - # TMEM offsets (fmha.py for 128x128) - self.tmem_s0_offset = 0; self.tmem_o0_offset = 256; self.tmem_p0_offset = 32 - self.tmem_alloc_cols = 512 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - pv_inst_k = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (*self.mma_tiler_mn, pv_inst_k * 4) - self.mma_tiler = self.qk_mma_tiler - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], - self.qk_mma_tiler[2], - ) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, self.use_2cta_instrs, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.a_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.b_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - # TMEM alloc cols — use the LARGER of QK and PV fragment sizes - qk_acc_shape = qk_mma.partition_shape_C(self.mma_tiler[:2]) - qk_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, 1)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(qk_fake, arch="sm_100") - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.a_dtype, a_smem) + cute.size_in_bytes(self.b_dtype, b_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - self.a_dtype = a.element_type; self.b_dtype = b.element_type; self.c_dtype = c.element_type - self.a_major = LayoutEnum.from_tensor(a).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(b).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, cute.nvgpu.OperandMajorMode.K, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - tma_a, tma_ta = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - a, a_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_b, tma_tb = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - b, b_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_a, tma_ta, tma_b, tma_tb, tma_c, tma_tc, - self.cluster_layout_vmnk, self.a_smem_s, self.b_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_a, mA, tma_b, mB, tma_c, mC, cl_vmnk, - a_smem_s, b_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_a); cpasync.prefetch_descriptor(tma_b); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] # MMA↔softmax pipeline (1 stage) - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - # AB pipeline - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - # MMA↔softmax pipeline - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - # ACC pipeline (PV output → epilogue) — KEY FIX: use warp count, NOT thread count - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - # TMEM allocator - tmem_bar = pipeline.NamedBarrier(barrier_id=self.tmem_alloc_sync_bar_id, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sA = smem.allocate_tensor(element_type=self.a_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sB = smem.allocate_tensor(element_type=self.b_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gA = cute.local_tile(mA, cute.slice_(self.mma_tiler, (None,0,None)), (None,None,None)) - gB = cute.local_tile(mB, cute.slice_(self.mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gA, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - tCgA = qk_thr.partition_A(gA); tCgB = qk_thr.partition_B(gB); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsA, tAgA = cpasync.tma_partition(tma_a, 0, a_lay, cute.group_modes(sA,0,3), cute.group_modes(tCgA,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsB, tBgB = cpasync.tma_partition(tma_b, 0, b_lay, cute.group_modes(sB,0,3), cute.group_modes(tCgB,0,3)) - tAgA = tAgA[(None,0,None,0)]; tBgB = tBgB[(None,0,None,0)] - - tCrA = qk_mma.make_fragment_A(sA); tCrB = qk_mma.make_fragment_B(sB) - tCrV = pv_mma.make_fragment_B(sB) - - # TMEM tensors - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - # P fragment for PV MMA (A-layout from TMEM) - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_mma.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ── TMA WARP ── - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_a, tAgA[(None,h.count)], tAsA[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_b, tBgB[(None,h.count)], tBsB[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_stage_b_v7.py b/tests/archive/test_stage_b_v7.py deleted file mode 100644 index a523ee6e..00000000 --- a/tests/archive/test_stage_b_v7.py +++ /dev/null @@ -1,450 +0,0 @@ -""" -Stage B v7: Two MMAs + Identity Softmax with COMPUTED TMEM offsets. - -Key fixes over v6: - - TMEM offsets computed via find_tmem_tensor_col_offset (same API as get_num_tmem_alloc_cols) - - P tensor constructed from p_tmem_s.outer (matching fmha.py pattern exactly) - - tilePlikeFP32 computed from qk_mma_tiler and dtype widths - - tmem_alloc_cols from get_num_tmem_alloc_cols with all fragments - - JIT-time diagnostic prints of all TMEM sizes - -Architecture (matches fmha.py exactly): - MMA1: Q @ K^T → tmem_scores (a_source=SMEM, accumulate=False) - Identity softmax: tcgen05.ld C-layout → F32→BF16 → tcgen05.st A-layout - MMA2: P @ V → tmem_output (a_source=TMEM, accumulate=True) -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - - -class StageBIdentitySoftmax: - def __init__(self, mma_tiler_mn, use_2cta_instrs=False, use_tma_store=True): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16 - self.use_2cta_instrs = use_2cta_instrs; self.use_tma_store = use_tma_store - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.epilog_sync_bar_id = 1; self.tmem_alloc_sync_bar_id = 2; self.tmem_dealloc_sync_bar_id = 3 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - pv_inst_k = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (*self.mma_tiler_mn, pv_inst_k * 4) - self.mma_tiler = self.qk_mma_tiler - print(f"[StageB] qk_mma.shape_mnk = {qk_mma.shape_mnk}") - print(f"[StageB] pv_mma.shape_mnk = {pv_mma.shape_mnk}") - print(f"[StageB] qk_mma_tiler = {self.qk_mma_tiler}") - print(f"[StageB] pv_mma_tiler = {self.pv_mma_tiler}") - print(f"[StageB] qk_inst_k = {qk_inst_k}, pv_inst_k = {pv_inst_k}") - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], - self.qk_mma_tiler[2], - ) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, self.use_2cta_instrs, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.a_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.b_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.b_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - # ── COMPUTE TMEM OFFSETS USING find_tmem_tensor_col_offset ── - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - # tilePlikeFP32 for the store-side composition - self.tilePlikeFP32 = self.qk_mma_tiler[1] * self.q_dtype.width // 32 - - # ── TMEM LAYOUT (matching fmha.py) ── - # P OVERLAPS S — after softmax, P (A-layout) is written on top of scores (C-layout) - # in the same TMEM region. The A-layout view starts partway through the S region. - # fmha.py: S0=0, P0=32, O0=256 (with S1=128, P1=160 double-buffered) - # The P offset of 32 means the A-layout starts at column 32 within the S region. - # This is because the C-layout and A-layout partition TMEM differently per-thread; - # the first 32 C-layout columns are "dead space" in the A-layout mapping. - # - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 # Original - self.tmem_o0_offset = s_cols # 128 - self.tmem_alloc_cols = s_cols + o_cols # 256 - - # Also compute via get_num_tmem_alloc_cols for the full allocation - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, 1)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, 1)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - print(f"[StageB] s_cols (QK accumulator) = {s_cols}") - print(f"[StageB] o_cols (PV accumulator) = {o_cols}") - print(f"[StageB] tilePlikeFP32 = {self.tilePlikeFP32}") - print(f"[StageB] tmem_s0_offset = {self.tmem_s0_offset}") - print(f"[StageB] tmem_p0_offset = {self.tmem_p0_offset}") - print(f"[StageB] tmem_o0_offset = {self.tmem_o0_offset}") - print(f"[StageB] tmem_alloc_cols (computed) = {self.tmem_alloc_cols}") - print(f"[StageB] num_tmem_alloc_cols (via utils) = {self.num_tmem_alloc_cols}") - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.a_dtype, a_smem) + cute.size_in_bytes(self.b_dtype, b_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - self.a_dtype = a.element_type; self.b_dtype = b.element_type; self.c_dtype = c.element_type - self.a_major = LayoutEnum.from_tensor(a).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(b).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, cute.nvgpu.OperandMajorMode.K, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - # Introspect PV MMA atom - print(f"[ATOM] PV MMA type: {type(pv_mma)}") - print(f"[ATOM] PV MMA op: {pv_mma.op if hasattr(pv_mma, "op") else "no op"}") - print(f"[ATOM] PV MMA trait: {pv_mma._trait if hasattr(pv_mma, "_trait") else "no trait"}") - print(f"[ATOM] PV MMA shape_mnk: {pv_mma.shape_mnk}") - print(f"[ATOM] QK MMA shape_mnk: {qk_mma.shape_mnk}") - # Check a_src - print(f"[ATOM] PV MMA op.a_src: {pv_mma.op.a_src}") - print(f"[ATOM] QK MMA op.a_src: {qk_mma.op.a_src}") - print(f"[ATOM] PV MMA op: {pv_mma.op}") - self._setup(qk_mma, pv_mma) - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - tma_a, tma_ta = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - a, a_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_b, tma_tb = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - b, b_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_a, tma_ta, tma_b, tma_tb, tma_c, tma_tc, - self.cluster_layout_vmnk, self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_a, mA, tma_b, mB, tma_c, mC, cl_vmnk, - a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_a); cpasync.prefetch_descriptor(tma_b); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=self.tmem_alloc_sync_bar_id, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sA = smem.allocate_tensor(element_type=self.a_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sB = smem.allocate_tensor(element_type=self.b_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - # V shares the same SMEM as B (same data, different layout for PV MMA) - sV_ptr = cute.recast_ptr(sB.iterator, v_smem_s.inner) - sV = cute.make_tensor(sV_ptr, v_smem_s.outer) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gA = cute.local_tile(mA, cute.slice_(self.mma_tiler, (None,0,None)), (None,None,None)) - gB = cute.local_tile(mB, cute.slice_(self.mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gA, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - tCgA = qk_thr.partition_A(gA); tCgB = qk_thr.partition_B(gB); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsA, tAgA = cpasync.tma_partition(tma_a, 0, a_lay, cute.group_modes(sA,0,3), cute.group_modes(tCgA,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsB, tBgB = cpasync.tma_partition(tma_b, 0, b_lay, cute.group_modes(sB,0,3), cute.group_modes(tCgB,0,3)) - tAgA = tAgA[(None,0,None,0)]; tBgB = tBgB[(None,0,None,0)] - - tCrA = qk_mma.make_fragment_A(sA); tCrB = qk_mma.make_fragment_B(sB) - tCrV = pv_mma.make_fragment_B(sV) # V fragment from V SMEM layout - print(f"[DIAG] tCrV.size = {cute.size(tCrV)}") - print(f"[DIAG] tCrA.size = {cute.size(tCrA)}") - print(f"[DIAG] tCrB.size = {cute.size(tCrB)}") - print(f"[DIAG] nblk_qk (tCrA mode 2) = {cute.size(tCrA, mode=[2])}") - - # ── TMEM tensors with computed offsets (matching fmha.py pattern) ── - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - # P fragment: construct from p_tmem_s layout (matching fmha.py exactly) - # fmha.py: tP = cute.make_tensor(tStS.iterator, p_tmem_layout_staged.outer) - # tOrP = pv_thr_mma.make_fragment_A(tP)[None, None, None, 0] - # tOrP0 = cute.make_tensor(tOrP.iterator + dtype_width_ratio * tmem_p0_offset, tOrP.layout) - print(f'[TMEM] p_tmem_s: {p_tmem_s}') - print(f'[TMEM] p_tmem_s.outer: {p_tmem_s.outer}') - print(f'[TMEM] p_tmem_s.inner: {p_tmem_s.inner}') - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - print(f'[DIAG] tStS.layout: {tStS.layout}') - print(f'[DIAG] tStS.size: {cute.size(tStS)}') - print(f'[DIAG] p_tmem_s.outer: {p_tmem_s.outer}') - print(f'[DIAG] p_tmem_s.inner: {p_tmem_s.inner}') - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - # Compute nblk_pv for diagnostics - nblk_pv = cute.size(tOrP0, mode=[2]) - nblk_qk = cute.size(tCrA, mode=[2]) - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - # COMPREHENSIVE LAYOUT DUMP - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tilePlikeFP32 = self.qk_mma_tiler[1] * self.q_dtype.width // 32 - tStS_P_layout = cute.composition(tStS.layout, cute.make_layout((128, tilePlikeFP32))) - tStS_P = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStS_P_layout) - tScS_P_layout = cute.composition(tScS.layout, cute.make_layout((128, tilePlikeFP32))) - tScS_P = cute.make_tensor(tScS.iterator, tScS_P_layout) - - print(f'[LAYOUT] QK C-fragment tStS.layout: {tStS.layout}') - print(f'[LAYOUT] QK C-fragment tStS cosize: {cute.cosize(tStS.layout)}') - print(f'[LAYOUT] QK C-fragment tStS.size: {cute.size(tStS)}') - print(f'[LAYOUT] QK C-fragment tScS.layout: {tScS.layout}') - print(f'[LAYOUT] QK C-fragment tScS cosize: {cute.cosize(tScS.layout)}') - print(f'[LAYOUT] PV A-fragment tOrP.layout: {tOrP.layout}') - print(f'[LAYOUT] PV A-fragment tOrP cosize: {cute.cosize(tOrP.layout)}') - print(f'[LAYOUT] PV A-fragment tOrP.size: {cute.size(tOrP)}') - print(f'[LAYOUT] PV A-fragment tOrP0.layout: {tOrP0.layout}') - print(f'[LAYOUT] PV A-fragment tOrP0 cosize: {cute.cosize(tOrP0.layout)}') - print(f'[LAYOUT] tP.layout: {tP.layout}') - print(f'[LAYOUT] tP cosize: {cute.cosize(tP.layout)}') - print(f'[LAYOUT] tStS_P (composed) layout: {tStS_P.layout}') - print(f'[LAYOUT] tStS_P (composed) cosize: {cute.cosize(tStS_P.layout)}') - print(f'[LAYOUT] tScS_P (composed) layout: {tScS_P.layout}') - print(f'[LAYOUT] tScS_P (composed) cosize: {cute.cosize(tScS_P.layout)}') - print(f'[LAYOUT] tOtO.layout: {tOtO.layout}') - print(f'[LAYOUT] tOtO cosize: {cute.cosize(tOtO.layout)}') - print(f'[LAYOUT] pv_mma_tiler: {self.pv_mma_tiler}') - print(f'[LAYOUT] qk_mma_tiler: {self.qk_mma_tiler}') - print(f'[LAYOUT] tilePlikeFP32: {tilePlikeFP32}') - - # DIAGNOSTIC: Compare tP (A-layout) vs tStS_P (composition) - tilePlikeFP32 = self.qk_mma_tiler[1] * self.q_dtype.width // 32 - tStS_P_layout = cute.composition(tStS.layout, cute.make_layout((128, tilePlikeFP32))) - tStS_P = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStS_P_layout) - print(f'[DIAG] tP.layout: {tP.layout}') - print(f'[DIAG] tP.size: {cute.size(tP)}') - print(f'[DIAG] tP.element_type: {tP.element_type if hasattr(tP, 'element_type') else 'N/A'}') - print(f'[DIAG] tStS_P.layout: {tStS_P.layout}') - print(f'[DIAG] tStS_P.size: {cute.size(tStS_P)}') - print(f'[DIAG] tStS_P.element_type: {tStS_P.element_type if hasattr(tStS_P, 'element_type') else 'N/A'}') - print(f'[DIAG] tilePlikeFP32: {tilePlikeFP32}') - print(f'[DIAG] tP and tStS_P same iterator? {tP.iterator == tStS_P.iterator if hasattr(tP, 'iterator') else 'cant compare'}') - - print(f'[DIAG] nblk_pv = {nblk_pv}, nblk_qk = {nblk_qk}') - print(f'[DIAG] tCrV.size = {cute.size(tCrV)}') - print(f'[DIAG] tOrP0.size = {cute.size(tOrP0)}') - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ── TMA WARP ── - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_a, tAgA[(None,h.count)], tAsA[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_b, tBgB[(None,h.count)], tBsB[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_stage_b_v7_rep128.py b/tests/archive/test_stage_b_v7_rep128.py deleted file mode 100644 index 8f0d3e15..00000000 --- a/tests/archive/test_stage_b_v7_rep128.py +++ /dev/null @@ -1,445 +0,0 @@ -""" -Stage B v7: Two MMAs + Identity Softmax with COMPUTED TMEM offsets. - -Key fixes over v6: - - TMEM offsets computed via find_tmem_tensor_col_offset (same API as get_num_tmem_alloc_cols) - - P tensor constructed from p_tmem_s.outer (matching fmha.py pattern exactly) - - tilePlikeFP32 computed from qk_mma_tiler and dtype widths - - tmem_alloc_cols from get_num_tmem_alloc_cols with all fragments - - JIT-time diagnostic prints of all TMEM sizes - -Architecture (matches fmha.py exactly): - MMA1: Q @ K^T → tmem_scores (a_source=SMEM, accumulate=False) - Identity softmax: tcgen05.ld C-layout → F32→BF16 → tcgen05.st A-layout - MMA2: P @ V → tmem_output (a_source=TMEM, accumulate=True) -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - - -class StageBIdentitySoftmax: - def __init__(self, mma_tiler_mn, use_2cta_instrs=False, use_tma_store=True): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16 - self.use_2cta_instrs = use_2cta_instrs; self.use_tma_store = use_tma_store - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.epilog_sync_bar_id = 1; self.tmem_alloc_sync_bar_id = 2; self.tmem_dealloc_sync_bar_id = 3 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - pv_inst_k = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (*self.mma_tiler_mn, pv_inst_k * 4) - self.mma_tiler = self.qk_mma_tiler - print(f"[StageB] qk_mma.shape_mnk = {qk_mma.shape_mnk}") - print(f"[StageB] pv_mma.shape_mnk = {pv_mma.shape_mnk}") - print(f"[StageB] qk_mma_tiler = {self.qk_mma_tiler}") - print(f"[StageB] pv_mma_tiler = {self.pv_mma_tiler}") - print(f"[StageB] qk_inst_k = {qk_inst_k}, pv_inst_k = {pv_inst_k}") - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], - self.qk_mma_tiler[2], - ) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, self.use_2cta_instrs, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.a_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.b_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.b_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - # ── COMPUTE TMEM OFFSETS USING find_tmem_tensor_col_offset ── - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - # tilePlikeFP32 for the store-side composition - self.tilePlikeFP32 = self.qk_mma_tiler[1] * self.q_dtype.width // 32 - - # ── TMEM LAYOUT (matching fmha.py) ── - # P OVERLAPS S — after softmax, P (A-layout) is written on top of scores (C-layout) - # in the same TMEM region. The A-layout view starts partway through the S region. - # fmha.py: S0=0, P0=32, O0=256 (with S1=128, P1=160 double-buffered) - # The P offset of 32 means the A-layout starts at column 32 within the S region. - # This is because the C-layout and A-layout partition TMEM differently per-thread; - # the first 32 C-layout columns are "dead space" in the A-layout mapping. - # - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 # Same as fmha.py - self.tmem_o0_offset = s_cols # 128 - self.tmem_alloc_cols = s_cols + o_cols # 256 - - # Also compute via get_num_tmem_alloc_cols for the full allocation - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, 1)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, 1)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - print(f"[StageB] s_cols (QK accumulator) = {s_cols}") - print(f"[StageB] o_cols (PV accumulator) = {o_cols}") - print(f"[StageB] tilePlikeFP32 = {self.tilePlikeFP32}") - print(f"[StageB] tmem_s0_offset = {self.tmem_s0_offset}") - print(f"[StageB] tmem_p0_offset = {self.tmem_p0_offset}") - print(f"[StageB] tmem_o0_offset = {self.tmem_o0_offset}") - print(f"[StageB] tmem_alloc_cols (computed) = {self.tmem_alloc_cols}") - print(f"[StageB] num_tmem_alloc_cols (via utils) = {self.num_tmem_alloc_cols}") - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.a_dtype, a_smem) + cute.size_in_bytes(self.b_dtype, b_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - self.a_dtype = a.element_type; self.b_dtype = b.element_type; self.c_dtype = c.element_type - self.a_major = LayoutEnum.from_tensor(a).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(b).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, cute.nvgpu.OperandMajorMode.K, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - tma_a, tma_ta = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - a, a_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_b, tma_tb = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - b, b_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_a, tma_ta, tma_b, tma_tb, tma_c, tma_tc, - self.cluster_layout_vmnk, self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_a, mA, tma_b, mB, tma_c, mC, cl_vmnk, - a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_a); cpasync.prefetch_descriptor(tma_b); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=self.tmem_alloc_sync_bar_id, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sA = smem.allocate_tensor(element_type=self.a_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sB = smem.allocate_tensor(element_type=self.b_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - # V shares the same SMEM as B (same data, different layout for PV MMA) - sV_ptr = cute.recast_ptr(sB.iterator, v_smem_s.inner) - sV = cute.make_tensor(sV_ptr, v_smem_s.outer) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gA = cute.local_tile(mA, cute.slice_(self.mma_tiler, (None,0,None)), (None,None,None)) - gB = cute.local_tile(mB, cute.slice_(self.mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gA, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - tCgA = qk_thr.partition_A(gA); tCgB = qk_thr.partition_B(gB); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsA, tAgA = cpasync.tma_partition(tma_a, 0, a_lay, cute.group_modes(sA,0,3), cute.group_modes(tCgA,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsB, tBgB = cpasync.tma_partition(tma_b, 0, b_lay, cute.group_modes(sB,0,3), cute.group_modes(tCgB,0,3)) - tAgA = tAgA[(None,0,None,0)]; tBgB = tBgB[(None,0,None,0)] - - tCrA = qk_mma.make_fragment_A(sA); tCrB = qk_mma.make_fragment_B(sB) - tCrV = pv_mma.make_fragment_B(sV) # V fragment from V SMEM layout - print(f"[DIAG] tCrV.size = {cute.size(tCrV)}") - print(f"[DIAG] tCrA.size = {cute.size(tCrA)}") - print(f"[DIAG] tCrB.size = {cute.size(tCrB)}") - print(f"[DIAG] nblk_qk (tCrA mode 2) = {cute.size(tCrA, mode=[2])}") - - # ── TMEM tensors with computed offsets (matching fmha.py pattern) ── - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - # P fragment: construct from p_tmem_s layout (matching fmha.py exactly) - # fmha.py: tP = cute.make_tensor(tStS.iterator, p_tmem_layout_staged.outer) - # tOrP = pv_thr_mma.make_fragment_A(tP)[None, None, None, 0] - # tOrP0 = cute.make_tensor(tOrP.iterator + dtype_width_ratio * tmem_p0_offset, tOrP.layout) - print(f'[TMEM] p_tmem_s: {p_tmem_s}') - print(f'[TMEM] p_tmem_s.outer: {p_tmem_s.outer}') - print(f'[TMEM] p_tmem_s.inner: {p_tmem_s.inner}') - # Check SMEM layout compatibility: K (b_smem_s) vs V (v_smem_s) - print(f'[SMEM] b_smem_s.outer: {self.b_smem_s.outer}') - print(f'[SMEM] v_smem_s.outer: {self.v_smem_s.outer}') - print(f'[SMEM] b_smem_s.inner: {self.b_smem_s.inner}') - print(f'[SMEM] v_smem_s.inner: {self.v_smem_s.inner}') - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - print(f'[DIAG] tStS.layout: {tStS.layout}') - print(f'[DIAG] tStS.size: {cute.size(tStS)}') - print(f'[DIAG] p_tmem_s.outer: {p_tmem_s.outer}') - print(f'[DIAG] p_tmem_s.inner: {p_tmem_s.inner}') - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - # Compute nblk_pv for diagnostics - nblk_pv = cute.size(tOrP0, mode=[2]) - nblk_qk = cute.size(tCrA, mode=[2]) - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - # COMPREHENSIVE LAYOUT DUMP - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tilePlikeFP32 = self.qk_mma_tiler[1] * self.q_dtype.width // 32 - tStS_P_layout = cute.composition(tStS.layout, cute.make_layout((128, tilePlikeFP32))) - tStS_P = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStS_P_layout) - tScS_P_layout = cute.composition(tScS.layout, cute.make_layout((128, tilePlikeFP32))) - tScS_P = cute.make_tensor(tScS.iterator, tScS_P_layout) - - print(f'[LAYOUT] QK C-fragment tStS.layout: {tStS.layout}') - print(f'[LAYOUT] QK C-fragment tStS cosize: {cute.cosize(tStS.layout)}') - print(f'[LAYOUT] QK C-fragment tStS.size: {cute.size(tStS)}') - print(f'[LAYOUT] QK C-fragment tScS.layout: {tScS.layout}') - print(f'[LAYOUT] QK C-fragment tScS cosize: {cute.cosize(tScS.layout)}') - print(f'[LAYOUT] PV A-fragment tOrP.layout: {tOrP.layout}') - print(f'[LAYOUT] PV A-fragment tOrP cosize: {cute.cosize(tOrP.layout)}') - print(f'[LAYOUT] PV A-fragment tOrP.size: {cute.size(tOrP)}') - print(f'[LAYOUT] PV A-fragment tOrP0.layout: {tOrP0.layout}') - print(f'[LAYOUT] PV A-fragment tOrP0 cosize: {cute.cosize(tOrP0.layout)}') - print(f'[LAYOUT] tP.layout: {tP.layout}') - print(f'[LAYOUT] tP cosize: {cute.cosize(tP.layout)}') - print(f'[LAYOUT] tStS_P (composed) layout: {tStS_P.layout}') - print(f'[LAYOUT] tStS_P (composed) cosize: {cute.cosize(tStS_P.layout)}') - print(f'[LAYOUT] tScS_P (composed) layout: {tScS_P.layout}') - print(f'[LAYOUT] tScS_P (composed) cosize: {cute.cosize(tScS_P.layout)}') - print(f'[LAYOUT] tOtO.layout: {tOtO.layout}') - print(f'[LAYOUT] tOtO cosize: {cute.cosize(tOtO.layout)}') - print(f'[LAYOUT] pv_mma_tiler: {self.pv_mma_tiler}') - print(f'[LAYOUT] qk_mma_tiler: {self.qk_mma_tiler}') - print(f'[LAYOUT] tilePlikeFP32: {tilePlikeFP32}') - - # DIAGNOSTIC: Compare tP (A-layout) vs tStS_P (composition) - tilePlikeFP32 = self.qk_mma_tiler[1] * self.q_dtype.width // 32 - tStS_P_layout = cute.composition(tStS.layout, cute.make_layout((128, tilePlikeFP32))) - tStS_P = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStS_P_layout) - print(f'[DIAG] tP.layout: {tP.layout}') - print(f'[DIAG] tP.size: {cute.size(tP)}') - print(f'[DIAG] tP.element_type: {tP.element_type if hasattr(tP, 'element_type') else 'N/A'}') - print(f'[DIAG] tStS_P.layout: {tStS_P.layout}') - print(f'[DIAG] tStS_P.size: {cute.size(tStS_P)}') - print(f'[DIAG] tStS_P.element_type: {tStS_P.element_type if hasattr(tStS_P, 'element_type') else 'N/A'}') - print(f'[DIAG] tilePlikeFP32: {tilePlikeFP32}') - print(f'[DIAG] tP and tStS_P same iterator? {tP.iterator == tStS_P.iterator if hasattr(tP, 'iterator') else 'cant compare'}') - - print(f'[DIAG] nblk_pv = {nblk_pv}, nblk_qk = {nblk_qk}') - print(f'[DIAG] tCrV.size = {cute.size(tCrV)}') - print(f'[DIAG] tOrP0.size = {cute.size(tOrP0)}') - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ── TMA WARP ── - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_a, tAgA[(None,h.count)], tAsA[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_b, tBgB[(None,h.count)], tBsB[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_stage_b_v7_rep16.py b/tests/archive/test_stage_b_v7_rep16.py deleted file mode 100644 index 8f0d3e15..00000000 --- a/tests/archive/test_stage_b_v7_rep16.py +++ /dev/null @@ -1,445 +0,0 @@ -""" -Stage B v7: Two MMAs + Identity Softmax with COMPUTED TMEM offsets. - -Key fixes over v6: - - TMEM offsets computed via find_tmem_tensor_col_offset (same API as get_num_tmem_alloc_cols) - - P tensor constructed from p_tmem_s.outer (matching fmha.py pattern exactly) - - tilePlikeFP32 computed from qk_mma_tiler and dtype widths - - tmem_alloc_cols from get_num_tmem_alloc_cols with all fragments - - JIT-time diagnostic prints of all TMEM sizes - -Architecture (matches fmha.py exactly): - MMA1: Q @ K^T → tmem_scores (a_source=SMEM, accumulate=False) - Identity softmax: tcgen05.ld C-layout → F32→BF16 → tcgen05.st A-layout - MMA2: P @ V → tmem_output (a_source=TMEM, accumulate=True) -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - - -class StageBIdentitySoftmax: - def __init__(self, mma_tiler_mn, use_2cta_instrs=False, use_tma_store=True): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16 - self.use_2cta_instrs = use_2cta_instrs; self.use_tma_store = use_tma_store - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.epilog_sync_bar_id = 1; self.tmem_alloc_sync_bar_id = 2; self.tmem_dealloc_sync_bar_id = 3 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - pv_inst_k = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (*self.mma_tiler_mn, pv_inst_k * 4) - self.mma_tiler = self.qk_mma_tiler - print(f"[StageB] qk_mma.shape_mnk = {qk_mma.shape_mnk}") - print(f"[StageB] pv_mma.shape_mnk = {pv_mma.shape_mnk}") - print(f"[StageB] qk_mma_tiler = {self.qk_mma_tiler}") - print(f"[StageB] pv_mma_tiler = {self.pv_mma_tiler}") - print(f"[StageB] qk_inst_k = {qk_inst_k}, pv_inst_k = {pv_inst_k}") - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], - self.qk_mma_tiler[2], - ) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, self.use_2cta_instrs, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.a_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.b_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.b_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - # ── COMPUTE TMEM OFFSETS USING find_tmem_tensor_col_offset ── - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - # tilePlikeFP32 for the store-side composition - self.tilePlikeFP32 = self.qk_mma_tiler[1] * self.q_dtype.width // 32 - - # ── TMEM LAYOUT (matching fmha.py) ── - # P OVERLAPS S — after softmax, P (A-layout) is written on top of scores (C-layout) - # in the same TMEM region. The A-layout view starts partway through the S region. - # fmha.py: S0=0, P0=32, O0=256 (with S1=128, P1=160 double-buffered) - # The P offset of 32 means the A-layout starts at column 32 within the S region. - # This is because the C-layout and A-layout partition TMEM differently per-thread; - # the first 32 C-layout columns are "dead space" in the A-layout mapping. - # - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 # Same as fmha.py - self.tmem_o0_offset = s_cols # 128 - self.tmem_alloc_cols = s_cols + o_cols # 256 - - # Also compute via get_num_tmem_alloc_cols for the full allocation - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, 1)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, 1)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - print(f"[StageB] s_cols (QK accumulator) = {s_cols}") - print(f"[StageB] o_cols (PV accumulator) = {o_cols}") - print(f"[StageB] tilePlikeFP32 = {self.tilePlikeFP32}") - print(f"[StageB] tmem_s0_offset = {self.tmem_s0_offset}") - print(f"[StageB] tmem_p0_offset = {self.tmem_p0_offset}") - print(f"[StageB] tmem_o0_offset = {self.tmem_o0_offset}") - print(f"[StageB] tmem_alloc_cols (computed) = {self.tmem_alloc_cols}") - print(f"[StageB] num_tmem_alloc_cols (via utils) = {self.num_tmem_alloc_cols}") - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.a_dtype, a_smem) + cute.size_in_bytes(self.b_dtype, b_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - self.a_dtype = a.element_type; self.b_dtype = b.element_type; self.c_dtype = c.element_type - self.a_major = LayoutEnum.from_tensor(a).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(b).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, cute.nvgpu.OperandMajorMode.K, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - tma_a, tma_ta = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - a, a_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_b, tma_tb = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - b, b_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_a, tma_ta, tma_b, tma_tb, tma_c, tma_tc, - self.cluster_layout_vmnk, self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_a, mA, tma_b, mB, tma_c, mC, cl_vmnk, - a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_a); cpasync.prefetch_descriptor(tma_b); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=self.tmem_alloc_sync_bar_id, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sA = smem.allocate_tensor(element_type=self.a_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sB = smem.allocate_tensor(element_type=self.b_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - # V shares the same SMEM as B (same data, different layout for PV MMA) - sV_ptr = cute.recast_ptr(sB.iterator, v_smem_s.inner) - sV = cute.make_tensor(sV_ptr, v_smem_s.outer) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gA = cute.local_tile(mA, cute.slice_(self.mma_tiler, (None,0,None)), (None,None,None)) - gB = cute.local_tile(mB, cute.slice_(self.mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gA, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - tCgA = qk_thr.partition_A(gA); tCgB = qk_thr.partition_B(gB); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsA, tAgA = cpasync.tma_partition(tma_a, 0, a_lay, cute.group_modes(sA,0,3), cute.group_modes(tCgA,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsB, tBgB = cpasync.tma_partition(tma_b, 0, b_lay, cute.group_modes(sB,0,3), cute.group_modes(tCgB,0,3)) - tAgA = tAgA[(None,0,None,0)]; tBgB = tBgB[(None,0,None,0)] - - tCrA = qk_mma.make_fragment_A(sA); tCrB = qk_mma.make_fragment_B(sB) - tCrV = pv_mma.make_fragment_B(sV) # V fragment from V SMEM layout - print(f"[DIAG] tCrV.size = {cute.size(tCrV)}") - print(f"[DIAG] tCrA.size = {cute.size(tCrA)}") - print(f"[DIAG] tCrB.size = {cute.size(tCrB)}") - print(f"[DIAG] nblk_qk (tCrA mode 2) = {cute.size(tCrA, mode=[2])}") - - # ── TMEM tensors with computed offsets (matching fmha.py pattern) ── - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - # P fragment: construct from p_tmem_s layout (matching fmha.py exactly) - # fmha.py: tP = cute.make_tensor(tStS.iterator, p_tmem_layout_staged.outer) - # tOrP = pv_thr_mma.make_fragment_A(tP)[None, None, None, 0] - # tOrP0 = cute.make_tensor(tOrP.iterator + dtype_width_ratio * tmem_p0_offset, tOrP.layout) - print(f'[TMEM] p_tmem_s: {p_tmem_s}') - print(f'[TMEM] p_tmem_s.outer: {p_tmem_s.outer}') - print(f'[TMEM] p_tmem_s.inner: {p_tmem_s.inner}') - # Check SMEM layout compatibility: K (b_smem_s) vs V (v_smem_s) - print(f'[SMEM] b_smem_s.outer: {self.b_smem_s.outer}') - print(f'[SMEM] v_smem_s.outer: {self.v_smem_s.outer}') - print(f'[SMEM] b_smem_s.inner: {self.b_smem_s.inner}') - print(f'[SMEM] v_smem_s.inner: {self.v_smem_s.inner}') - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - print(f'[DIAG] tStS.layout: {tStS.layout}') - print(f'[DIAG] tStS.size: {cute.size(tStS)}') - print(f'[DIAG] p_tmem_s.outer: {p_tmem_s.outer}') - print(f'[DIAG] p_tmem_s.inner: {p_tmem_s.inner}') - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - # Compute nblk_pv for diagnostics - nblk_pv = cute.size(tOrP0, mode=[2]) - nblk_qk = cute.size(tCrA, mode=[2]) - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - # COMPREHENSIVE LAYOUT DUMP - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tilePlikeFP32 = self.qk_mma_tiler[1] * self.q_dtype.width // 32 - tStS_P_layout = cute.composition(tStS.layout, cute.make_layout((128, tilePlikeFP32))) - tStS_P = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStS_P_layout) - tScS_P_layout = cute.composition(tScS.layout, cute.make_layout((128, tilePlikeFP32))) - tScS_P = cute.make_tensor(tScS.iterator, tScS_P_layout) - - print(f'[LAYOUT] QK C-fragment tStS.layout: {tStS.layout}') - print(f'[LAYOUT] QK C-fragment tStS cosize: {cute.cosize(tStS.layout)}') - print(f'[LAYOUT] QK C-fragment tStS.size: {cute.size(tStS)}') - print(f'[LAYOUT] QK C-fragment tScS.layout: {tScS.layout}') - print(f'[LAYOUT] QK C-fragment tScS cosize: {cute.cosize(tScS.layout)}') - print(f'[LAYOUT] PV A-fragment tOrP.layout: {tOrP.layout}') - print(f'[LAYOUT] PV A-fragment tOrP cosize: {cute.cosize(tOrP.layout)}') - print(f'[LAYOUT] PV A-fragment tOrP.size: {cute.size(tOrP)}') - print(f'[LAYOUT] PV A-fragment tOrP0.layout: {tOrP0.layout}') - print(f'[LAYOUT] PV A-fragment tOrP0 cosize: {cute.cosize(tOrP0.layout)}') - print(f'[LAYOUT] tP.layout: {tP.layout}') - print(f'[LAYOUT] tP cosize: {cute.cosize(tP.layout)}') - print(f'[LAYOUT] tStS_P (composed) layout: {tStS_P.layout}') - print(f'[LAYOUT] tStS_P (composed) cosize: {cute.cosize(tStS_P.layout)}') - print(f'[LAYOUT] tScS_P (composed) layout: {tScS_P.layout}') - print(f'[LAYOUT] tScS_P (composed) cosize: {cute.cosize(tScS_P.layout)}') - print(f'[LAYOUT] tOtO.layout: {tOtO.layout}') - print(f'[LAYOUT] tOtO cosize: {cute.cosize(tOtO.layout)}') - print(f'[LAYOUT] pv_mma_tiler: {self.pv_mma_tiler}') - print(f'[LAYOUT] qk_mma_tiler: {self.qk_mma_tiler}') - print(f'[LAYOUT] tilePlikeFP32: {tilePlikeFP32}') - - # DIAGNOSTIC: Compare tP (A-layout) vs tStS_P (composition) - tilePlikeFP32 = self.qk_mma_tiler[1] * self.q_dtype.width // 32 - tStS_P_layout = cute.composition(tStS.layout, cute.make_layout((128, tilePlikeFP32))) - tStS_P = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStS_P_layout) - print(f'[DIAG] tP.layout: {tP.layout}') - print(f'[DIAG] tP.size: {cute.size(tP)}') - print(f'[DIAG] tP.element_type: {tP.element_type if hasattr(tP, 'element_type') else 'N/A'}') - print(f'[DIAG] tStS_P.layout: {tStS_P.layout}') - print(f'[DIAG] tStS_P.size: {cute.size(tStS_P)}') - print(f'[DIAG] tStS_P.element_type: {tStS_P.element_type if hasattr(tStS_P, 'element_type') else 'N/A'}') - print(f'[DIAG] tilePlikeFP32: {tilePlikeFP32}') - print(f'[DIAG] tP and tStS_P same iterator? {tP.iterator == tStS_P.iterator if hasattr(tP, 'iterator') else 'cant compare'}') - - print(f'[DIAG] nblk_pv = {nblk_pv}, nblk_qk = {nblk_qk}') - print(f'[DIAG] tCrV.size = {cute.size(tCrV)}') - print(f'[DIAG] tOrP0.size = {cute.size(tOrP0)}') - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ── TMA WARP ── - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_a, tAgA[(None,h.count)], tAsA[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_b, tBgB[(None,h.count)], tBsB[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_stage_b_v7_rep64.py b/tests/archive/test_stage_b_v7_rep64.py deleted file mode 100644 index 8f0d3e15..00000000 --- a/tests/archive/test_stage_b_v7_rep64.py +++ /dev/null @@ -1,445 +0,0 @@ -""" -Stage B v7: Two MMAs + Identity Softmax with COMPUTED TMEM offsets. - -Key fixes over v6: - - TMEM offsets computed via find_tmem_tensor_col_offset (same API as get_num_tmem_alloc_cols) - - P tensor constructed from p_tmem_s.outer (matching fmha.py pattern exactly) - - tilePlikeFP32 computed from qk_mma_tiler and dtype widths - - tmem_alloc_cols from get_num_tmem_alloc_cols with all fragments - - JIT-time diagnostic prints of all TMEM sizes - -Architecture (matches fmha.py exactly): - MMA1: Q @ K^T → tmem_scores (a_source=SMEM, accumulate=False) - Identity softmax: tcgen05.ld C-layout → F32→BF16 → tcgen05.st A-layout - MMA2: P @ V → tmem_output (a_source=TMEM, accumulate=True) -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - - -class StageBIdentitySoftmax: - def __init__(self, mma_tiler_mn, use_2cta_instrs=False, use_tma_store=True): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16 - self.use_2cta_instrs = use_2cta_instrs; self.use_tma_store = use_tma_store - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.epilog_sync_bar_id = 1; self.tmem_alloc_sync_bar_id = 2; self.tmem_dealloc_sync_bar_id = 3 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - pv_inst_k = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (*self.mma_tiler_mn, pv_inst_k * 4) - self.mma_tiler = self.qk_mma_tiler - print(f"[StageB] qk_mma.shape_mnk = {qk_mma.shape_mnk}") - print(f"[StageB] pv_mma.shape_mnk = {pv_mma.shape_mnk}") - print(f"[StageB] qk_mma_tiler = {self.qk_mma_tiler}") - print(f"[StageB] pv_mma_tiler = {self.pv_mma_tiler}") - print(f"[StageB] qk_inst_k = {qk_inst_k}, pv_inst_k = {pv_inst_k}") - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], - self.qk_mma_tiler[2], - ) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, self.use_2cta_instrs, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.a_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.b_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.b_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - # ── COMPUTE TMEM OFFSETS USING find_tmem_tensor_col_offset ── - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - # tilePlikeFP32 for the store-side composition - self.tilePlikeFP32 = self.qk_mma_tiler[1] * self.q_dtype.width // 32 - - # ── TMEM LAYOUT (matching fmha.py) ── - # P OVERLAPS S — after softmax, P (A-layout) is written on top of scores (C-layout) - # in the same TMEM region. The A-layout view starts partway through the S region. - # fmha.py: S0=0, P0=32, O0=256 (with S1=128, P1=160 double-buffered) - # The P offset of 32 means the A-layout starts at column 32 within the S region. - # This is because the C-layout and A-layout partition TMEM differently per-thread; - # the first 32 C-layout columns are "dead space" in the A-layout mapping. - # - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 # Same as fmha.py - self.tmem_o0_offset = s_cols # 128 - self.tmem_alloc_cols = s_cols + o_cols # 256 - - # Also compute via get_num_tmem_alloc_cols for the full allocation - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, 1)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, 1)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - print(f"[StageB] s_cols (QK accumulator) = {s_cols}") - print(f"[StageB] o_cols (PV accumulator) = {o_cols}") - print(f"[StageB] tilePlikeFP32 = {self.tilePlikeFP32}") - print(f"[StageB] tmem_s0_offset = {self.tmem_s0_offset}") - print(f"[StageB] tmem_p0_offset = {self.tmem_p0_offset}") - print(f"[StageB] tmem_o0_offset = {self.tmem_o0_offset}") - print(f"[StageB] tmem_alloc_cols (computed) = {self.tmem_alloc_cols}") - print(f"[StageB] num_tmem_alloc_cols (via utils) = {self.num_tmem_alloc_cols}") - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.a_dtype, a_smem) + cute.size_in_bytes(self.b_dtype, b_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - self.a_dtype = a.element_type; self.b_dtype = b.element_type; self.c_dtype = c.element_type - self.a_major = LayoutEnum.from_tensor(a).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(b).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, cute.nvgpu.OperandMajorMode.K, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - tma_a, tma_ta = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - a, a_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_b, tma_tb = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - b, b_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_a, tma_ta, tma_b, tma_tb, tma_c, tma_tc, - self.cluster_layout_vmnk, self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_a, mA, tma_b, mB, tma_c, mC, cl_vmnk, - a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_a); cpasync.prefetch_descriptor(tma_b); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=self.tmem_alloc_sync_bar_id, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sA = smem.allocate_tensor(element_type=self.a_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sB = smem.allocate_tensor(element_type=self.b_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - # V shares the same SMEM as B (same data, different layout for PV MMA) - sV_ptr = cute.recast_ptr(sB.iterator, v_smem_s.inner) - sV = cute.make_tensor(sV_ptr, v_smem_s.outer) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gA = cute.local_tile(mA, cute.slice_(self.mma_tiler, (None,0,None)), (None,None,None)) - gB = cute.local_tile(mB, cute.slice_(self.mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gA, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - tCgA = qk_thr.partition_A(gA); tCgB = qk_thr.partition_B(gB); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsA, tAgA = cpasync.tma_partition(tma_a, 0, a_lay, cute.group_modes(sA,0,3), cute.group_modes(tCgA,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsB, tBgB = cpasync.tma_partition(tma_b, 0, b_lay, cute.group_modes(sB,0,3), cute.group_modes(tCgB,0,3)) - tAgA = tAgA[(None,0,None,0)]; tBgB = tBgB[(None,0,None,0)] - - tCrA = qk_mma.make_fragment_A(sA); tCrB = qk_mma.make_fragment_B(sB) - tCrV = pv_mma.make_fragment_B(sV) # V fragment from V SMEM layout - print(f"[DIAG] tCrV.size = {cute.size(tCrV)}") - print(f"[DIAG] tCrA.size = {cute.size(tCrA)}") - print(f"[DIAG] tCrB.size = {cute.size(tCrB)}") - print(f"[DIAG] nblk_qk (tCrA mode 2) = {cute.size(tCrA, mode=[2])}") - - # ── TMEM tensors with computed offsets (matching fmha.py pattern) ── - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - # P fragment: construct from p_tmem_s layout (matching fmha.py exactly) - # fmha.py: tP = cute.make_tensor(tStS.iterator, p_tmem_layout_staged.outer) - # tOrP = pv_thr_mma.make_fragment_A(tP)[None, None, None, 0] - # tOrP0 = cute.make_tensor(tOrP.iterator + dtype_width_ratio * tmem_p0_offset, tOrP.layout) - print(f'[TMEM] p_tmem_s: {p_tmem_s}') - print(f'[TMEM] p_tmem_s.outer: {p_tmem_s.outer}') - print(f'[TMEM] p_tmem_s.inner: {p_tmem_s.inner}') - # Check SMEM layout compatibility: K (b_smem_s) vs V (v_smem_s) - print(f'[SMEM] b_smem_s.outer: {self.b_smem_s.outer}') - print(f'[SMEM] v_smem_s.outer: {self.v_smem_s.outer}') - print(f'[SMEM] b_smem_s.inner: {self.b_smem_s.inner}') - print(f'[SMEM] v_smem_s.inner: {self.v_smem_s.inner}') - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - print(f'[DIAG] tStS.layout: {tStS.layout}') - print(f'[DIAG] tStS.size: {cute.size(tStS)}') - print(f'[DIAG] p_tmem_s.outer: {p_tmem_s.outer}') - print(f'[DIAG] p_tmem_s.inner: {p_tmem_s.inner}') - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - # Compute nblk_pv for diagnostics - nblk_pv = cute.size(tOrP0, mode=[2]) - nblk_qk = cute.size(tCrA, mode=[2]) - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - # COMPREHENSIVE LAYOUT DUMP - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tilePlikeFP32 = self.qk_mma_tiler[1] * self.q_dtype.width // 32 - tStS_P_layout = cute.composition(tStS.layout, cute.make_layout((128, tilePlikeFP32))) - tStS_P = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStS_P_layout) - tScS_P_layout = cute.composition(tScS.layout, cute.make_layout((128, tilePlikeFP32))) - tScS_P = cute.make_tensor(tScS.iterator, tScS_P_layout) - - print(f'[LAYOUT] QK C-fragment tStS.layout: {tStS.layout}') - print(f'[LAYOUT] QK C-fragment tStS cosize: {cute.cosize(tStS.layout)}') - print(f'[LAYOUT] QK C-fragment tStS.size: {cute.size(tStS)}') - print(f'[LAYOUT] QK C-fragment tScS.layout: {tScS.layout}') - print(f'[LAYOUT] QK C-fragment tScS cosize: {cute.cosize(tScS.layout)}') - print(f'[LAYOUT] PV A-fragment tOrP.layout: {tOrP.layout}') - print(f'[LAYOUT] PV A-fragment tOrP cosize: {cute.cosize(tOrP.layout)}') - print(f'[LAYOUT] PV A-fragment tOrP.size: {cute.size(tOrP)}') - print(f'[LAYOUT] PV A-fragment tOrP0.layout: {tOrP0.layout}') - print(f'[LAYOUT] PV A-fragment tOrP0 cosize: {cute.cosize(tOrP0.layout)}') - print(f'[LAYOUT] tP.layout: {tP.layout}') - print(f'[LAYOUT] tP cosize: {cute.cosize(tP.layout)}') - print(f'[LAYOUT] tStS_P (composed) layout: {tStS_P.layout}') - print(f'[LAYOUT] tStS_P (composed) cosize: {cute.cosize(tStS_P.layout)}') - print(f'[LAYOUT] tScS_P (composed) layout: {tScS_P.layout}') - print(f'[LAYOUT] tScS_P (composed) cosize: {cute.cosize(tScS_P.layout)}') - print(f'[LAYOUT] tOtO.layout: {tOtO.layout}') - print(f'[LAYOUT] tOtO cosize: {cute.cosize(tOtO.layout)}') - print(f'[LAYOUT] pv_mma_tiler: {self.pv_mma_tiler}') - print(f'[LAYOUT] qk_mma_tiler: {self.qk_mma_tiler}') - print(f'[LAYOUT] tilePlikeFP32: {tilePlikeFP32}') - - # DIAGNOSTIC: Compare tP (A-layout) vs tStS_P (composition) - tilePlikeFP32 = self.qk_mma_tiler[1] * self.q_dtype.width // 32 - tStS_P_layout = cute.composition(tStS.layout, cute.make_layout((128, tilePlikeFP32))) - tStS_P = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStS_P_layout) - print(f'[DIAG] tP.layout: {tP.layout}') - print(f'[DIAG] tP.size: {cute.size(tP)}') - print(f'[DIAG] tP.element_type: {tP.element_type if hasattr(tP, 'element_type') else 'N/A'}') - print(f'[DIAG] tStS_P.layout: {tStS_P.layout}') - print(f'[DIAG] tStS_P.size: {cute.size(tStS_P)}') - print(f'[DIAG] tStS_P.element_type: {tStS_P.element_type if hasattr(tStS_P, 'element_type') else 'N/A'}') - print(f'[DIAG] tilePlikeFP32: {tilePlikeFP32}') - print(f'[DIAG] tP and tStS_P same iterator? {tP.iterator == tStS_P.iterator if hasattr(tP, 'iterator') else 'cant compare'}') - - print(f'[DIAG] nblk_pv = {nblk_pv}, nblk_qk = {nblk_qk}') - print(f'[DIAG] tCrV.size = {cute.size(tCrV)}') - print(f'[DIAG] tOrP0.size = {cute.size(tOrP0)}') - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ── TMA WARP ── - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_a, tAgA[(None,h.count)], tAsA[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_b, tBgB[(None,h.count)], tBsB[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_stage_b_v7_rep8.py b/tests/archive/test_stage_b_v7_rep8.py deleted file mode 100644 index 8f0d3e15..00000000 --- a/tests/archive/test_stage_b_v7_rep8.py +++ /dev/null @@ -1,445 +0,0 @@ -""" -Stage B v7: Two MMAs + Identity Softmax with COMPUTED TMEM offsets. - -Key fixes over v6: - - TMEM offsets computed via find_tmem_tensor_col_offset (same API as get_num_tmem_alloc_cols) - - P tensor constructed from p_tmem_s.outer (matching fmha.py pattern exactly) - - tilePlikeFP32 computed from qk_mma_tiler and dtype widths - - tmem_alloc_cols from get_num_tmem_alloc_cols with all fragments - - JIT-time diagnostic prints of all TMEM sizes - -Architecture (matches fmha.py exactly): - MMA1: Q @ K^T → tmem_scores (a_source=SMEM, accumulate=False) - Identity softmax: tcgen05.ld C-layout → F32→BF16 → tcgen05.st A-layout - MMA2: P @ V → tmem_output (a_source=TMEM, accumulate=True) -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - - -class StageBIdentitySoftmax: - def __init__(self, mma_tiler_mn, use_2cta_instrs=False, use_tma_store=True): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16 - self.use_2cta_instrs = use_2cta_instrs; self.use_tma_store = use_tma_store - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.epilog_sync_bar_id = 1; self.tmem_alloc_sync_bar_id = 2; self.tmem_dealloc_sync_bar_id = 3 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - pv_inst_k = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (*self.mma_tiler_mn, pv_inst_k * 4) - self.mma_tiler = self.qk_mma_tiler - print(f"[StageB] qk_mma.shape_mnk = {qk_mma.shape_mnk}") - print(f"[StageB] pv_mma.shape_mnk = {pv_mma.shape_mnk}") - print(f"[StageB] qk_mma_tiler = {self.qk_mma_tiler}") - print(f"[StageB] pv_mma_tiler = {self.pv_mma_tiler}") - print(f"[StageB] qk_inst_k = {qk_inst_k}, pv_inst_k = {pv_inst_k}") - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], - self.qk_mma_tiler[2], - ) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, self.use_2cta_instrs, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.a_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.b_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.b_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - # ── COMPUTE TMEM OFFSETS USING find_tmem_tensor_col_offset ── - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - # tilePlikeFP32 for the store-side composition - self.tilePlikeFP32 = self.qk_mma_tiler[1] * self.q_dtype.width // 32 - - # ── TMEM LAYOUT (matching fmha.py) ── - # P OVERLAPS S — after softmax, P (A-layout) is written on top of scores (C-layout) - # in the same TMEM region. The A-layout view starts partway through the S region. - # fmha.py: S0=0, P0=32, O0=256 (with S1=128, P1=160 double-buffered) - # The P offset of 32 means the A-layout starts at column 32 within the S region. - # This is because the C-layout and A-layout partition TMEM differently per-thread; - # the first 32 C-layout columns are "dead space" in the A-layout mapping. - # - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 # Same as fmha.py - self.tmem_o0_offset = s_cols # 128 - self.tmem_alloc_cols = s_cols + o_cols # 256 - - # Also compute via get_num_tmem_alloc_cols for the full allocation - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, 1)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, 1)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - print(f"[StageB] s_cols (QK accumulator) = {s_cols}") - print(f"[StageB] o_cols (PV accumulator) = {o_cols}") - print(f"[StageB] tilePlikeFP32 = {self.tilePlikeFP32}") - print(f"[StageB] tmem_s0_offset = {self.tmem_s0_offset}") - print(f"[StageB] tmem_p0_offset = {self.tmem_p0_offset}") - print(f"[StageB] tmem_o0_offset = {self.tmem_o0_offset}") - print(f"[StageB] tmem_alloc_cols (computed) = {self.tmem_alloc_cols}") - print(f"[StageB] num_tmem_alloc_cols (via utils) = {self.num_tmem_alloc_cols}") - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.a_dtype, a_smem) + cute.size_in_bytes(self.b_dtype, b_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - self.a_dtype = a.element_type; self.b_dtype = b.element_type; self.c_dtype = c.element_type - self.a_major = LayoutEnum.from_tensor(a).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(b).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, cute.nvgpu.OperandMajorMode.K, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - tma_a, tma_ta = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - a, a_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_b, tma_tb = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - b, b_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_a, tma_ta, tma_b, tma_tb, tma_c, tma_tc, - self.cluster_layout_vmnk, self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_a, mA, tma_b, mB, tma_c, mC, cl_vmnk, - a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_a); cpasync.prefetch_descriptor(tma_b); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=self.tmem_alloc_sync_bar_id, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sA = smem.allocate_tensor(element_type=self.a_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sB = smem.allocate_tensor(element_type=self.b_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - # V shares the same SMEM as B (same data, different layout for PV MMA) - sV_ptr = cute.recast_ptr(sB.iterator, v_smem_s.inner) - sV = cute.make_tensor(sV_ptr, v_smem_s.outer) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gA = cute.local_tile(mA, cute.slice_(self.mma_tiler, (None,0,None)), (None,None,None)) - gB = cute.local_tile(mB, cute.slice_(self.mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gA, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - tCgA = qk_thr.partition_A(gA); tCgB = qk_thr.partition_B(gB); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsA, tAgA = cpasync.tma_partition(tma_a, 0, a_lay, cute.group_modes(sA,0,3), cute.group_modes(tCgA,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsB, tBgB = cpasync.tma_partition(tma_b, 0, b_lay, cute.group_modes(sB,0,3), cute.group_modes(tCgB,0,3)) - tAgA = tAgA[(None,0,None,0)]; tBgB = tBgB[(None,0,None,0)] - - tCrA = qk_mma.make_fragment_A(sA); tCrB = qk_mma.make_fragment_B(sB) - tCrV = pv_mma.make_fragment_B(sV) # V fragment from V SMEM layout - print(f"[DIAG] tCrV.size = {cute.size(tCrV)}") - print(f"[DIAG] tCrA.size = {cute.size(tCrA)}") - print(f"[DIAG] tCrB.size = {cute.size(tCrB)}") - print(f"[DIAG] nblk_qk (tCrA mode 2) = {cute.size(tCrA, mode=[2])}") - - # ── TMEM tensors with computed offsets (matching fmha.py pattern) ── - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - # P fragment: construct from p_tmem_s layout (matching fmha.py exactly) - # fmha.py: tP = cute.make_tensor(tStS.iterator, p_tmem_layout_staged.outer) - # tOrP = pv_thr_mma.make_fragment_A(tP)[None, None, None, 0] - # tOrP0 = cute.make_tensor(tOrP.iterator + dtype_width_ratio * tmem_p0_offset, tOrP.layout) - print(f'[TMEM] p_tmem_s: {p_tmem_s}') - print(f'[TMEM] p_tmem_s.outer: {p_tmem_s.outer}') - print(f'[TMEM] p_tmem_s.inner: {p_tmem_s.inner}') - # Check SMEM layout compatibility: K (b_smem_s) vs V (v_smem_s) - print(f'[SMEM] b_smem_s.outer: {self.b_smem_s.outer}') - print(f'[SMEM] v_smem_s.outer: {self.v_smem_s.outer}') - print(f'[SMEM] b_smem_s.inner: {self.b_smem_s.inner}') - print(f'[SMEM] v_smem_s.inner: {self.v_smem_s.inner}') - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - print(f'[DIAG] tStS.layout: {tStS.layout}') - print(f'[DIAG] tStS.size: {cute.size(tStS)}') - print(f'[DIAG] p_tmem_s.outer: {p_tmem_s.outer}') - print(f'[DIAG] p_tmem_s.inner: {p_tmem_s.inner}') - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - # Compute nblk_pv for diagnostics - nblk_pv = cute.size(tOrP0, mode=[2]) - nblk_qk = cute.size(tCrA, mode=[2]) - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - # COMPREHENSIVE LAYOUT DUMP - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tilePlikeFP32 = self.qk_mma_tiler[1] * self.q_dtype.width // 32 - tStS_P_layout = cute.composition(tStS.layout, cute.make_layout((128, tilePlikeFP32))) - tStS_P = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStS_P_layout) - tScS_P_layout = cute.composition(tScS.layout, cute.make_layout((128, tilePlikeFP32))) - tScS_P = cute.make_tensor(tScS.iterator, tScS_P_layout) - - print(f'[LAYOUT] QK C-fragment tStS.layout: {tStS.layout}') - print(f'[LAYOUT] QK C-fragment tStS cosize: {cute.cosize(tStS.layout)}') - print(f'[LAYOUT] QK C-fragment tStS.size: {cute.size(tStS)}') - print(f'[LAYOUT] QK C-fragment tScS.layout: {tScS.layout}') - print(f'[LAYOUT] QK C-fragment tScS cosize: {cute.cosize(tScS.layout)}') - print(f'[LAYOUT] PV A-fragment tOrP.layout: {tOrP.layout}') - print(f'[LAYOUT] PV A-fragment tOrP cosize: {cute.cosize(tOrP.layout)}') - print(f'[LAYOUT] PV A-fragment tOrP.size: {cute.size(tOrP)}') - print(f'[LAYOUT] PV A-fragment tOrP0.layout: {tOrP0.layout}') - print(f'[LAYOUT] PV A-fragment tOrP0 cosize: {cute.cosize(tOrP0.layout)}') - print(f'[LAYOUT] tP.layout: {tP.layout}') - print(f'[LAYOUT] tP cosize: {cute.cosize(tP.layout)}') - print(f'[LAYOUT] tStS_P (composed) layout: {tStS_P.layout}') - print(f'[LAYOUT] tStS_P (composed) cosize: {cute.cosize(tStS_P.layout)}') - print(f'[LAYOUT] tScS_P (composed) layout: {tScS_P.layout}') - print(f'[LAYOUT] tScS_P (composed) cosize: {cute.cosize(tScS_P.layout)}') - print(f'[LAYOUT] tOtO.layout: {tOtO.layout}') - print(f'[LAYOUT] tOtO cosize: {cute.cosize(tOtO.layout)}') - print(f'[LAYOUT] pv_mma_tiler: {self.pv_mma_tiler}') - print(f'[LAYOUT] qk_mma_tiler: {self.qk_mma_tiler}') - print(f'[LAYOUT] tilePlikeFP32: {tilePlikeFP32}') - - # DIAGNOSTIC: Compare tP (A-layout) vs tStS_P (composition) - tilePlikeFP32 = self.qk_mma_tiler[1] * self.q_dtype.width // 32 - tStS_P_layout = cute.composition(tStS.layout, cute.make_layout((128, tilePlikeFP32))) - tStS_P = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStS_P_layout) - print(f'[DIAG] tP.layout: {tP.layout}') - print(f'[DIAG] tP.size: {cute.size(tP)}') - print(f'[DIAG] tP.element_type: {tP.element_type if hasattr(tP, 'element_type') else 'N/A'}') - print(f'[DIAG] tStS_P.layout: {tStS_P.layout}') - print(f'[DIAG] tStS_P.size: {cute.size(tStS_P)}') - print(f'[DIAG] tStS_P.element_type: {tStS_P.element_type if hasattr(tStS_P, 'element_type') else 'N/A'}') - print(f'[DIAG] tilePlikeFP32: {tilePlikeFP32}') - print(f'[DIAG] tP and tStS_P same iterator? {tP.iterator == tStS_P.iterator if hasattr(tP, 'iterator') else 'cant compare'}') - - print(f'[DIAG] nblk_pv = {nblk_pv}, nblk_qk = {nblk_qk}') - print(f'[DIAG] tCrV.size = {cute.size(tCrV)}') - print(f'[DIAG] tOrP0.size = {cute.size(tOrP0)}') - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ── TMA WARP ── - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_a, tAgA[(None,h.count)], tAsA[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_b, tBgB[(None,h.count)], tBsB[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_stage_b_v8.py b/tests/archive/test_stage_b_v8.py deleted file mode 100644 index e3c7f372..00000000 --- a/tests/archive/test_stage_b_v8.py +++ /dev/null @@ -1,403 +0,0 @@ -""" -Stage B v8: Fix the identity softmax by using the A-fragment layout -for the TMEM store target instead of the C-fragment composition. - -The bug in v7: tStS_P uses composition(tStS.layout, (128, tilePlikeFP32)) -which gives layout (128,64):(65536,1) — C-fragment strides. -But the PV MMA reads from TMEM using the A-fragment layout -((128,16),1,4):((64,1),0,16) — physical TMEM strides. - -For the store to be read correctly by the PV MMA, the store target -must use the same physical TMEM addressing as the A-fragment. - -Key insight from CUTLASS source: - Physical TMEM for M=128, BK=64, BF16 A from TMEM (K-major): - tmem[dp=m, col=base_col + 16*mma_k + k_inner] for mma_k in 0..3, k_inner in 0..15 - - This means: A-fragment address = 64*m + k_inner + 16*mma_k - C-fragment address for (m, col) = ??? (virtual layout, not physical) - - The St32x32b copy atom with tStS_P (C-composition) writes to C-layout addresses. - The PV MMA reads from A-layout addresses. These are different physical locations. - -Fix: Use tP (from p_tmem_s, the A-fragment source layout) as the store target -instead of tStS_P (the C-fragment composition). This ensures the store writes -to physical TMEM addresses that the PV MMA's A-fragment will read correctly. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - - -class StageBIdentitySoftmaxV8: - def __init__(self, mma_tiler_mn, use_2cta_instrs=False, use_tma_store=True): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16 - self.use_2cta_instrs = use_2cta_instrs; self.use_tma_store = use_tma_store - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.epilog_sync_bar_id = 1; self.tmem_alloc_sync_bar_id = 2; self.tmem_dealloc_sync_bar_id = 3 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - pv_inst_k = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (*self.mma_tiler_mn, pv_inst_k * 4) - self.mma_tiler = self.qk_mma_tiler - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], - self.qk_mma_tiler[2], - ) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, self.use_2cta_instrs, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.a_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.b_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.b_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols # 128 - self.tmem_alloc_cols = s_cols + o_cols # 256 - - self.tilePlikeFP32 = self.qk_mma_tiler[1] * self.q_dtype.width // 32 - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, 1)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, 1)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.a_dtype, a_smem) + cute.size_in_bytes(self.b_dtype, b_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - self.a_dtype = a.element_type; self.b_dtype = b.element_type; self.c_dtype = c.element_type - self.a_major = LayoutEnum.from_tensor(a).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(b).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, cute.nvgpu.OperandMajorMode.K, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - tma_a, tma_ta = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - a, a_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_b, tma_tb = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - b, b_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_a, tma_ta, tma_b, tma_tb, tma_c, tma_tc, - self.cluster_layout_vmnk, self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_a, mA, tma_b, mB, tma_c, mC, cl_vmnk, - a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_a); cpasync.prefetch_descriptor(tma_b); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=self.tmem_alloc_sync_bar_id, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sA = smem.allocate_tensor(element_type=self.a_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sB = smem.allocate_tensor(element_type=self.b_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV_ptr = cute.recast_ptr(sB.iterator, v_smem_s.inner) - sV = cute.make_tensor(sV_ptr, v_smem_s.outer) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gA = cute.local_tile(mA, cute.slice_(self.mma_tiler, (None,0,None)), (None,None,None)) - gB = cute.local_tile(mB, cute.slice_(self.mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gA, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - tCgA = qk_thr.partition_A(gA); tCgB = qk_thr.partition_B(gB); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsA, tAgA = cpasync.tma_partition(tma_a, 0, a_lay, cute.group_modes(sA,0,3), cute.group_modes(tCgA,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsB, tBgB = cpasync.tma_partition(tma_b, 0, b_lay, cute.group_modes(sB,0,3), cute.group_modes(tCgB,0,3)) - tAgA = tAgA[(None,0,None,0)]; tBgB = tBgB[(None,0,None,0)] - - tCrA = qk_mma.make_fragment_A(sA); tCrB = qk_mma.make_fragment_B(sB) - tCrV = pv_mma.make_fragment_B(sV) - - # TMEM tensors - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - # P tensor for PV MMA A-fragment - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP = pv_thr.make_fragment_A(tP)[None, None, None, 0] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - # ── KEY FIX: Store target uses A-fragment layout (tP) not C-fragment composition ── - # The store must write to physical TMEM addresses that the PV MMA reads via A-fragment. - # tP has layout ((128,16),1,4,1):((64,1),0,16,0) — the A-fragment's physical TMEM layout. - # We need the store target at tmem_p0_offset = 32 columns into the S region. - # tP's iterator starts at tStS.iterator (base of S region). - # tOrP0 starts at tStS.iterator + 2*32 (scaled by F32/BF16 width ratio for the A-fragment). - # The store target should use the SAME layout as tP but with the p0 offset applied. - tP_store = cute.make_tensor( - tStS.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - p_tmem_s.outer) - - print(f'[v8] tP.layout: {tP.layout}') - print(f'[v8] tP_store.layout: {tP_store.layout}') - print(f'[v8] tOrP0.layout: {tOrP0.layout}') - print(f'[v8] tStS.layout: {tStS.layout}') - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ── TMA WARP ── - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_a, tAgA[(None,h.count)], tAsA[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_b, tBgB[(None,h.count)], tBsB[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1 BF16 - tTMEM_STORErP = cute.make_rmem_tensor(tScS_A.shape, self.qk_acc_dtype) - tTMEM_STORErP_e = cute.make_tensor( - cute.recast_ptr(tTMEM_STORErP.iterator, dtype=self.q_dtype), - tTMEM_LOADrS.layout) - s_vec = tTMEM_LOADrS.load() - tTMEM_STORErP_e.store(s_vec.to(self.q_dtype)) - - # Store to A-layout TMEM - cute.copy(tiled_tmem_store, tTMEM_STORErP, tTMEM_STOREtP) - cute.arch.fence_view_async_tmem_store() - - si_handle.release() - - # ── Epilogue ── - tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout) - acc_cons_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage) - c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) - c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp) - acc_cons_st = utils.gemm.sm100.epilogue_tma_store( - self, tidx, warp_idx, tma_c, tCtO_base, sC, tCgC, - epi_tile, 0, const_expr(lambda x: x), (0,0,0), acc_cons_st, acc_pipe, c_pipe) - c_pipe.producer_tail() - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def test(): - torch.manual_seed(42) - m, n, k = 128, 128, 128 - q = torch.randn(m, k, 1, dtype=torch.bfloat16, device='cuda') - kv = torch.randn(n, k, 1, dtype=torch.bfloat16, device='cuda') - c = torch.zeros(m, n, 1, dtype=torch.bfloat16, device='cuda') - qf = q[:,:,0].float(); kvf = kv[:,:,0].float() - ref = qf @ kvf.T @ kvf - import cutlass.torch as ct - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(kv).mark_layout_dynamic(leading_dim=ct.get_leading_dim(kv)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - kernel = StageBIdentitySoftmaxV8(mma_tiler_mn=(128, 128), use_2cta_instrs=False, use_tma_store=True) - print('Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mC, stream) - print('Running...', flush=True) - compiled(mQ, mK, mC, stream) - torch.cuda.synchronize() - out = c[:,:,0].float() - cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - max_err = (out - ref).abs().max().item() - print('Stage B v8: (Q @ K^T) @ V with identity softmax (A-layout store)') - print(' Cosine: {:.6f}, Max error: {:.6f}'.format(cos, max_err)) - print(' {}'.format('PASS' if cos >= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_stage_b_v8b.py b/tests/archive/test_stage_b_v8b.py deleted file mode 100644 index 67b8fd89..00000000 --- a/tests/archive/test_stage_b_v8b.py +++ /dev/null @@ -1,354 +0,0 @@ -""" -Stage B v8b: BF16 store directly to tOrP0 (diagnostic) - -Key fixes over v6: - - TMEM offsets computed via find_tmem_tensor_col_offset (same API as get_num_tmem_alloc_cols) - - P tensor constructed from p_tmem_s.outer (matching fmha.py pattern exactly) - - tilePlikeFP32 computed from qk_mma_tiler and dtype widths - - tmem_alloc_cols from get_num_tmem_alloc_cols with all fragments - - JIT-time diagnostic prints of all TMEM sizes - -Architecture (matches fmha.py exactly): - MMA1: Q @ K^T → tmem_scores (a_source=SMEM, accumulate=False) - Identity softmax: tcgen05.ld C-layout → F32→BF16 → tcgen05.st A-layout - MMA2: P @ V → tmem_output (a_source=TMEM, accumulate=True) -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - - -class StageBIdentitySoftmax: - def __init__(self, mma_tiler_mn, use_2cta_instrs=False, use_tma_store=True): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16 - self.use_2cta_instrs = use_2cta_instrs; self.use_tma_store = use_tma_store - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.epilog_sync_bar_id = 1; self.tmem_alloc_sync_bar_id = 2; self.tmem_dealloc_sync_bar_id = 3 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - pv_inst_k = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (*self.mma_tiler_mn, pv_inst_k * 4) - self.mma_tiler = self.qk_mma_tiler - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], - self.qk_mma_tiler[2], - ) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, self.use_2cta_instrs, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.a_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.b_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.b_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - # ── COMPUTE TMEM OFFSETS USING find_tmem_tensor_col_offset ── - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - # tilePlikeFP32 for the store-side composition - self.tilePlikeFP32 = self.qk_mma_tiler[1] * self.q_dtype.width // 32 - - # ── TMEM LAYOUT (matching fmha.py) ── - # P OVERLAPS S — after softmax, P (A-layout) is written on top of scores (C-layout) - # in the same TMEM region. The A-layout view starts partway through the S region. - # fmha.py: S0=0, P0=32, O0=256 (with S1=128, P1=160 double-buffered) - # The P offset of 32 means the A-layout starts at column 32 within the S region. - # This is because the C-layout and A-layout partition TMEM differently per-thread; - # the first 32 C-layout columns are "dead space" in the A-layout mapping. - # - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 # Same as fmha.py - self.tmem_o0_offset = s_cols # 128 - self.tmem_alloc_cols = s_cols + o_cols # 256 - - # Also compute via get_num_tmem_alloc_cols for the full allocation - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, 1)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, 1)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - print(f"[StageB] s_cols (QK accumulator) = {s_cols}") - print(f"[StageB] o_cols (PV accumulator) = {o_cols}") - print(f"[StageB] tilePlikeFP32 = {self.tilePlikeFP32}") - print(f"[StageB] tmem_s0_offset = {self.tmem_s0_offset}") - print(f"[StageB] tmem_p0_offset = {self.tmem_p0_offset}") - print(f"[StageB] tmem_o0_offset = {self.tmem_o0_offset}") - print(f"[StageB] tmem_alloc_cols (computed) = {self.tmem_alloc_cols}") - print(f"[StageB] num_tmem_alloc_cols (via utils) = {self.num_tmem_alloc_cols}") - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.a_dtype, a_smem) + cute.size_in_bytes(self.b_dtype, b_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - self.a_dtype = a.element_type; self.b_dtype = b.element_type; self.c_dtype = c.element_type - self.a_major = LayoutEnum.from_tensor(a).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(b).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, cute.nvgpu.OperandMajorMode.K, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - tma_a, tma_ta = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - a, a_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_b, tma_tb = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - b, b_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_a, tma_ta, tma_b, tma_tb, tma_c, tma_tc, - self.cluster_layout_vmnk, self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_a, mA, tma_b, mB, tma_c, mC, cl_vmnk, - a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_a); cpasync.prefetch_descriptor(tma_b); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=self.tmem_alloc_sync_bar_id, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sA = smem.allocate_tensor(element_type=self.a_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sB = smem.allocate_tensor(element_type=self.b_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - # V shares the same SMEM as B (same data, different layout for PV MMA) - sV_ptr = cute.recast_ptr(sB.iterator, v_smem_s.inner) - sV = cute.make_tensor(sV_ptr, v_smem_s.outer) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gA = cute.local_tile(mA, cute.slice_(self.mma_tiler, (None,0,None)), (None,None,None)) - gB = cute.local_tile(mB, cute.slice_(self.mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gA, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - tCgA = qk_thr.partition_A(gA); tCgB = qk_thr.partition_B(gB); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsA, tAgA = cpasync.tma_partition(tma_a, 0, a_lay, cute.group_modes(sA,0,3), cute.group_modes(tCgA,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsB, tBgB = cpasync.tma_partition(tma_b, 0, b_lay, cute.group_modes(sB,0,3), cute.group_modes(tCgB,0,3)) - tAgA = tAgA[(None,0,None,0)]; tBgB = tBgB[(None,0,None,0)] - - tCrA = qk_mma.make_fragment_A(sA); tCrB = qk_mma.make_fragment_B(sB) - tCrV = pv_mma.make_fragment_B(sV) # V fragment from V SMEM layout - - # ── TMEM tensors with computed offsets (matching fmha.py pattern) ── - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - # P fragment: construct from p_tmem_s layout (matching fmha.py exactly) - # fmha.py: tP = cute.make_tensor(tStS.iterator, p_tmem_layout_staged.outer) - # tOrP = pv_thr_mma.make_fragment_A(tP)[None, None, None, 0] - # tOrP0 = cute.make_tensor(tOrP.iterator + dtype_width_ratio * tmem_p0_offset, tOrP.layout) - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ── TMA WARP ── - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_a, tAgA[(None,h.count)], tAsA[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_b, tBgB[(None,h.count)], tBsB[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_stage_b_v9.py b/tests/archive/test_stage_b_v9.py deleted file mode 100644 index e8ea6e94..00000000 --- a/tests/archive/test_stage_b_v9.py +++ /dev/null @@ -1,348 +0,0 @@ -""" -Stage B v9: Identity softmax matching fmha.py softmax_step EXACTLY. - -Line by line match of the fmha softmax_step, but with identity softmax: - - No masking - - scale = 1 (no log2 scaling) - - row_max = 0 (skip max, just do exp(x) = 1 for identity) - - Actually for identity softmax, P = S (scores). So we just ld S, st as P. - - But we need to go through the C→A layout transform properly. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - -class StageBv9: - def __init__(self, mma_tiler_mn): - self.qk_acc_dtype = Float32; self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - self.acc_dtype = Float32; self.epilog_sync_bar_id = 1 - self.use_2cta_instrs = False - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = 128 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - pv_inst_k = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (*self.mma_tiler_mn, pv_inst_k * 4) - self.mma_tiler = self.qk_mma_tiler - self.tilePlikeFP32 = self.qk_mma_tiler[1] * self.q_dtype.width // 32 - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, False, c_layout, self.o_dtype) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, c_layout, self.epi_tile, 2) - self.c_layout = c_layout - self.num_ab_stage = 1; self.num_acc_stage = 1 - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols # 128 - self.tmem_alloc_cols = s_cols + o_cols # 256 - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, 1)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, 1)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, - LayoutEnum.from_tensor(a).mma_major_mode(), - LayoutEnum.from_tensor(b).mma_major_mode(), - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, - tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, - cute.nvgpu.OperandMajorMode.K, - LayoutEnum.from_tensor(b).mma_major_mode(), - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, - tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - tma_a, tma_ta = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - a, a_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_b, tma_tb = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - b, b_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_a, tma_ta, tma_b, tma_tb, tma_c, tma_tc, - self.cluster_layout_vmnk, self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_a, mA, tma_b, mB, tma_c, mC, cl_vmnk, - a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_a); cpasync.prefetch_descriptor(tma_b); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sA = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sB = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV_ptr = cute.recast_ptr(sB.iterator, v_smem_s.inner) - sV = cute.make_tensor(sV_ptr, v_smem_s.outer) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gA = cute.local_tile(mA, cute.slice_(self.mma_tiler, (None,0,None)), (None,None,None)) - gB = cute.local_tile(mB, cute.slice_(self.mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gA, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - tCgA = qk_thr.partition_A(gA); tCgB = qk_thr.partition_B(gB); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsA, tAgA = cpasync.tma_partition(tma_a, 0, a_lay, cute.group_modes(sA,0,3), cute.group_modes(tCgA,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsB, tBgB = cpasync.tma_partition(tma_b, 0, b_lay, cute.group_modes(sB,0,3), cute.group_modes(tCgB,0,3)) - tAgA = tAgA[(None,0,None,0)]; tBgB = tBgB[(None,0,None,0)] - - tCrA = qk_mma.make_fragment_A(sA); tCrB = qk_mma.make_fragment_B(sB) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP = pv_thr.make_fragment_A(tP)[None, None, None, 0] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - # ── TMEM copy atoms (matching fmha softmax_step exactly) ── - tmem_load_atom = cute.make_copy_atom( - tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS0) - thr_tmem_load = tiled_tmem_load.get_slice(tidx % (32 * len(self.epilogue_warp_id))) - tTMEM_LOADtS = thr_tmem_load.partition_S(tStS0) - - # Store target: composition of C-fragment (matching fmha) - tStS_P = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, - cute.composition(tStS.layout, cute.make_layout((128, self.tilePlikeFP32)))) - tmem_store_atom = cute.make_copy_atom( - tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStS_P) - thr_tmem_store = tiled_tmem_store.get_slice(tidx % (32 * len(self.epilogue_warp_id))) - tTMEM_STOREtS_x4 = thr_tmem_store.partition_D(tStS_P) - - # Identity tensors - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tScS_P = cute.make_tensor(tScS.iterator, - cute.composition(tScS.layout, cute.make_layout((128, self.tilePlikeFP32)))) - - tTMEM_LOADcS = thr_tmem_load.partition_D(tScS) - tTMEM_STOREcS = thr_tmem_store.partition_S(tScS_P) - - print(f'[v9] tTMEM_LOADcS.shape: {tTMEM_LOADcS.shape}') - print(f'[v9] tTMEM_STOREcS.shape: {tTMEM_STOREcS.shape}') - print(f'[v9] LOAD size: {cute.size(tTMEM_LOADcS)}') - print(f'[v9] STORE size: {cute.size(tTMEM_STOREcS)}') - print(f'[v9] tilePlikeFP32: {self.tilePlikeFP32}') - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ── TMA WARP ── - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_a, tAgA[(None,h.count)], tAsA[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_b, tBgB[(None,h.count)], tBsB[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_step2_subtile.py b/tests/archive/test_step2_subtile.py deleted file mode 100644 index 2c76e1f5..00000000 --- a/tests/archive/test_step2_subtile.py +++ /dev/null @@ -1,98 +0,0 @@ -"""Test: Validate gate/up subtile detection and SiLU on gate subtiles. - -This test runs the fused kernel with: -- Gate subtiles (0,1): SiLU applied, NOT written to GMEM -- Up subtiles (2,3): kept as-is, written to GMEM at positions (0,1) - -Expected output: (M, intermediate) BF16 with up values. -The output should match the up portion of the standard L1 GEMM output. -""" -import torch -import sys -sys.path.insert(0, '/root/dsv4-nvfp4-workspace/kernel') - -from dsv4.ops.quantize import ( - quantize_weight_to_nvfp4, - quantize_activation_nvfp4, -) -from dsv4.ops.layouts import ( - make_b_k_major, - assemble_scales_2d_side, - assemble_scales_3d_side, -) -from dsv4.ops.gemm_runner import ( - run_nvfp4_grouped_gemm, - run_fused_swiglu_grouped_gemm, - warmup_compilation, -) - - -def test_gate_up_subtile(): - device = "cuda" - num_experts = 4 - hidden = 512 - intermediate = 256 - num_tokens = 32 - - torch.manual_seed(42) - x = torch.randn(num_tokens, hidden, dtype=torch.bfloat16, device=device) - l1_w = torch.randn(num_experts, 2 * intermediate, hidden, dtype=torch.bfloat16, device=device) - - l1_fp4_list, l1_sf_list, l1_gs_list = [], [], [] - for e in range(num_experts): - w_fp4, w_sf, w_gs = quantize_weight_to_nvfp4(l1_w[e].T) - l1_fp4_list.append(w_fp4) - l1_sf_list.append(w_sf) - l1_gs_list.append(w_gs) - - l1_mat_b = make_b_k_major(torch.stack(l1_fp4_list)) - l1_scale_b = assemble_scales_3d_side(l1_sf_list) - l1_gs = torch.tensor(l1_gs_list, dtype=torch.float32, device=device) - - gs_val = x.abs().max().item() / (6.0 * 448.0) - x_fp4, x_sf = quantize_activation_nvfp4(x, gs_val) - tokens_per_expert = [num_tokens // num_experts] * num_experts - scale_a = assemble_scales_2d_side([x_sf[i*tpe:(i+1)*tpe] for i, tpe in enumerate(tokens_per_expert)]) - expert_offsets = torch.tensor( - [sum(tokens_per_expert[:e+1]) for e in range(num_experts)], - dtype=torch.int32, device=device, - ) - global_scale_a = torch.full((num_experts,), gs_val, dtype=torch.float32, device=device) - - warmup_compilation(num_experts, hidden // 2, (2 * intermediate) // 2, device) - - # 1. Standard L1 GEMM - out_bf16 = run_nvfp4_grouped_gemm( - mat_a=x_fp4, mat_b=l1_mat_b, - scale_a=scale_a, scale_b=l1_scale_b, - expert_offsets=expert_offsets, - global_scale_a=global_scale_a, global_scale_b=l1_gs, - ) - gate_ref = out_bf16[:, :intermediate] - up_ref = out_bf16[:, intermediate:] - print(f"Standard L1 output: shape={out_bf16.shape}") - print(f"Gate ref amax: {gate_ref.abs().amax().item():.4f}") - print(f"Up ref amax: {up_ref.abs().amax().item():.4f}") - - # 2. Fused kernel (gate: SiLU, up: as-is, only up written to GMEM) - print("\nRunning fused kernel...") - out_fused = run_fused_swiglu_grouped_gemm( - mat_a=x_fp4, mat_b=l1_mat_b, - scale_a=scale_a, scale_b=l1_scale_b, - expert_offsets=expert_offsets, - global_scale_a=global_scale_a, global_scale_b=l1_gs, - ) - print(f"Fused output: shape={out_fused.shape}, amax={out_fused.abs().amax().item():.4f}") - - # 3. Compare: fused output should match the up half of the standard output - diff = (out_fused - up_ref).float() - rel_err = diff.norm() / up_ref.float().norm() - max_err = diff.abs().max() - print(f"\n=== Results ===") - print(f"Rel error vs up_ref: {rel_err.item():.6f}") - print(f"Max abs error: {max_err.item():.6f}") - print(f"PASS" if rel_err.item() < 0.05 else "FAIL") - - -if __name__ == "__main__": - test_gate_up_subtile() diff --git a/tests/archive/test_step2_subtile_v2.py b/tests/archive/test_step2_subtile_v2.py deleted file mode 100644 index a91015a6..00000000 --- a/tests/archive/test_step2_subtile_v2.py +++ /dev/null @@ -1,114 +0,0 @@ -"""Test: Validate gate/up subtile detection (Step 2). - -The fused kernel writes: -- Gate subtiles (0,1): SiLU applied, stored to C tensor at positions 0,1 -- Up subtiles (2,3): raw values, stored to C tensor at positions 0,1 (overwriting gate) - (because TMA store uses gate_subtile_idx for up subtiles) - -For now, the output is still (M, 2*intermediate). We compare the -gate half of the output against SiLU(gate_ref) and the up half against up_ref. -""" -import torch -import sys -sys.path.insert(0, '/root/dsv4-nvfp4-workspace/kernel') - -from dsv4.ops.quantize import ( - quantize_weight_to_nvfp4, - quantize_activation_nvfp4, -) -from dsv4.ops.layouts import ( - make_b_k_major, - assemble_scales_2d_side, - assemble_scales_3d_side, -) -from dsv4.ops.gemm_runner import ( - run_nvfp4_grouped_gemm, - run_fused_swiglu_grouped_gemm, - warmup_compilation, -) - - -def test_gate_up_subtile(): - device = "cuda" - num_experts = 4 - hidden = 512 - intermediate = 256 - num_tokens = 32 - - torch.manual_seed(42) - x = torch.randn(num_tokens, hidden, dtype=torch.bfloat16, device=device) - l1_w = torch.randn(num_experts, 2 * intermediate, hidden, dtype=torch.bfloat16, device=device) - - l1_fp4_list, l1_sf_list, l1_gs_list = [], [], [] - for e in range(num_experts): - w_fp4, w_sf, w_gs = quantize_weight_to_nvfp4(l1_w[e].T) - l1_fp4_list.append(w_fp4) - l1_sf_list.append(w_sf) - l1_gs_list.append(w_gs) - - l1_mat_b = make_b_k_major(torch.stack(l1_fp4_list)) - l1_scale_b = assemble_scales_3d_side(l1_sf_list) - l1_gs = torch.tensor(l1_gs_list, dtype=torch.float32, device=device) - - gs_val = x.abs().max().item() / (6.0 * 448.0) - x_fp4, x_sf = quantize_activation_nvfp4(x, gs_val) - tokens_per_expert = [num_tokens // num_experts] * num_experts - scale_a = assemble_scales_2d_side([x_sf[i*tpe:(i+1)*tpe] for i, tpe in enumerate(tokens_per_expert)]) - expert_offsets = torch.tensor( - [sum(tokens_per_expert[:e+1]) for e in range(num_experts)], - dtype=torch.int32, device=device, - ) - global_scale_a = torch.full((num_experts,), gs_val, dtype=torch.float32, device=device) - - warmup_compilation(num_experts, hidden // 2, (2 * intermediate) // 2, device) - - # Standard L1 GEMM - out_bf16 = run_nvfp4_grouped_gemm( - mat_a=x_fp4, mat_b=l1_mat_b, - scale_a=scale_a, scale_b=l1_scale_b, - expert_offsets=expert_offsets, - global_scale_a=global_scale_a, global_scale_b=l1_gs, - ) - gate_ref = out_bf16[:, :intermediate] - up_ref = out_bf16[:, intermediate:] - silu_gate_ref = torch.nn.functional.silu(gate_ref) - - # Fused kernel - print("Running fused kernel...") - out_fused = run_fused_swiglu_grouped_gemm( - mat_a=x_fp4, mat_b=l1_mat_b, - scale_a=scale_a, scale_b=l1_scale_b, - expert_offsets=expert_offsets, - global_scale_a=global_scale_a, global_scale_b=l1_gs, - ) - - print(f"Fused output: shape={out_fused.shape}, amax={out_fused.abs().amax().item():.4f}") - - # The output has both gate (with SiLU) and up (raw) subtiles - # Gate is in the first half, up in the second half - fused_gate = out_fused[:, :intermediate] - fused_up = out_fused[:, intermediate:] - - # Compare gate: fused should have SiLU applied - gate_diff = (fused_gate - silu_gate_ref).float() - gate_rel_err = gate_diff.norm() / silu_gate_ref.float().norm() - gate_max_err = gate_diff.abs().max() - - # Compare up: fused should have raw values (no SiLU) - up_diff = (fused_up - up_ref).float() - up_rel_err = up_diff.norm() / up_ref.float().norm() - up_max_err = up_diff.abs().max() - - print(f"\n=== Gate Comparison (SiLU applied) ===") - print(f"Rel error: {gate_rel_err.item():.6f}") - print(f"Max abs error: {gate_max_err.item():.6f}") - print(f"Gate PASS" if gate_rel_err.item() < 0.05 else "Gate FAIL") - - print(f"\n=== Up Comparison (raw values) ===") - print(f"Rel error: {up_rel_err.item():.6f}") - print(f"Max abs error: {up_max_err.item():.6f}") - print(f"Up PASS" if up_rel_err.item() < 0.05 else "Up FAIL") - - -if __name__ == "__main__": - test_gate_up_subtile() diff --git a/tests/archive/test_store_verify.py b/tests/archive/test_store_verify.py deleted file mode 100644 index 85a0b41c..00000000 --- a/tests/archive/test_store_verify.py +++ /dev/null @@ -1,285 +0,0 @@ -""" -Test: Q@K^T → TMEM (scores), ld scores, st to P region, -then epilogue reads P region as C-fragment (not PV MMA). - -If the ld/st roundtrip preserves the data, the epilogue should -output the same as Stage A (Q@K^T result). -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - -class StoreVerify: - def __init__(self, mma_tiler_mn): - self.qk_acc_dtype = Float32; self.q_dtype = BFloat16; self.o_dtype = BFloat16 - self.c_dtype = BFloat16; self.acc_dtype = Float32 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2; self.use_2cta_instrs = False - self.epilog_sync_bar_id = 1 - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - - def _setup(self, qk_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - self.mma_tiler = self.qk_mma_tiler - self.tilePlikeFP32 = self.qk_mma_tiler[1] * self.q_dtype.width // 32 - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - c_layout = LayoutEnum.ROW_MAJOR - self.c_layout = c_layout - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, False, c_layout, self.o_dtype) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, c_layout, self.epi_tile, 2) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - self.tmem_alloc_cols = s_cols # Only need scores region, no O region - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, - LayoutEnum.from_tensor(a).mma_major_mode(), - LayoutEnum.from_tensor(b).mma_major_mode(), - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, - tcgen05.OperandSource.SMEM) - self._setup(qk_mma) - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - tma_a, tma_ta = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - a, a_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_b, tma_tb = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - b, b_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, tma_a, tma_ta, tma_b, tma_tb, tma_c, tma_tc, - self.cluster_layout_vmnk, self.a_smem_s, self.b_smem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, tma_a, mA, tma_b, mB, tma_c, mC, cl_vmnk, - a_smem_s, b_smem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_a); cpasync.prefetch_descriptor(tma_b); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, len(self.epilogue_warp_id)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=False, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sA = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sB = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gA = cute.local_tile(mA, cute.slice_(self.mma_tiler, (None,0,None)), (None,None,None)) - gB = cute.local_tile(mB, cute.slice_(self.mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gA, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - tCgA = qk_thr.partition_A(gA); tCgB = qk_thr.partition_B(gB); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsA, tAgA = cpasync.tma_partition(tma_a, 0, a_lay, cute.group_modes(sA,0,3), cute.group_modes(tCgA,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsB, tBgB = cpasync.tma_partition(tma_b, 0, b_lay, cute.group_modes(sB,0,3), cute.group_modes(tCgB,0,3)) - tAgA = tAgA[(None,0,None,0)]; tBgB = tBgB[(None,0,None,0)] - - tCrA = qk_mma.make_fragment_A(sA); tCrB = qk_mma.make_fragment_B(sB) - - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - # TMEM copy atoms for ld/st - tmem_load_atom = cute.make_copy_atom( - tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS0) - sfw_idx = tidx % (32 * len(self.epilogue_warp_id)) - thr_tmem_load = tiled_tmem_load.get_slice(sfw_idx) - tTMEM_LOADtS = thr_tmem_load.partition_S(tStS0) - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tTMEM_LOADcS = thr_tmem_load.partition_D(tScS) - - # Store target: P region (composition of C-fragment, offset 32) - tStS_P = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, - cute.composition(tStS.layout, cute.make_layout((128, self.tilePlikeFP32)))) - tmem_store_atom = cute.make_copy_atom( - tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStS_P) - thr_tmem_store = tiled_tmem_store.get_slice(sfw_idx) - tTMEM_STOREtP = thr_tmem_store.partition_D(tStS_P) - tScS_P = cute.make_tensor(tScS.iterator, - cute.composition(tScS.layout, cute.make_layout((128, self.tilePlikeFP32)))) - tTMEM_STOREcS = thr_tmem_store.partition_S(tScS_P) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, 1)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ── TMA WARP ── - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_a, tAgA[(None,h.count)], tAsA[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_b, tBgB[(None,h.count)], tBsB[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_store_verify2.py b/tests/archive/test_store_verify2.py deleted file mode 100644 index fe5a11a3..00000000 --- a/tests/archive/test_store_verify2.py +++ /dev/null @@ -1,272 +0,0 @@ -"""Store verify v2: ld S (full 128x128), st P (full 128x128 at offset 128), -then epilogue reads P region. No subview, no composition.""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - -class StoreVerify2: - def __init__(self, mma_tiler_mn): - self.qk_acc_dtype = Float32; self.q_dtype = BFloat16; self.o_dtype = BFloat16 - self.c_dtype = BFloat16; self.acc_dtype = Float32 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2; self.use_2cta_instrs = False - self.epilog_sync_bar_id = 1 - - def _setup(self, qk_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - self.mma_tiler = self.qk_mma_tiler - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - c_layout = LayoutEnum.ROW_MAJOR; self.c_layout = c_layout - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, False, c_layout, self.o_dtype) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, c_layout, self.epi_tile, 2) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - self.tmem_alloc_cols = s_cols * 2 # Two full regions - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, - LayoutEnum.from_tensor(a).mma_major_mode(), - LayoutEnum.from_tensor(b).mma_major_mode(), - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, - tcgen05.OperandSource.SMEM) - self._setup(qk_mma) - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - tma_a, tma_ta = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - a, a_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_b, tma_tb = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - b, b_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, tma_a, tma_ta, tma_b, tma_tb, tma_c, tma_tc, - self.cluster_layout_vmnk, self.a_smem_s, self.b_smem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, tma_a, mA, tma_b, mB, tma_c, mC, cl_vmnk, - a_smem_s, b_smem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_a); cpasync.prefetch_descriptor(tma_b); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, len(self.epilogue_warp_id)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=False, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sA = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sB = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gA = cute.local_tile(mA, cute.slice_(self.mma_tiler, (None,0,None)), (None,None,None)) - gB = cute.local_tile(mB, cute.slice_(self.mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gA, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - tCgA = qk_thr.partition_A(gA); tCgB = qk_thr.partition_B(gB); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsA, tAgA = cpasync.tma_partition(tma_a, 0, a_lay, cute.group_modes(sA,0,3), cute.group_modes(tCgA,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsB, tBgB = cpasync.tma_partition(tma_b, 0, b_lay, cute.group_modes(sB,0,3), cute.group_modes(tCgB,0,3)) - tAgA = tAgA[(None,0,None,0)]; tBgB = tBgB[(None,0,None,0)] - - tCrA = qk_mma.make_fragment_A(sA); tCrB = qk_mma.make_fragment_B(sB) - - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - tStS0 = cute.make_tensor(tStS.iterator, tStS.layout) # offset 0 - tStS1 = cute.make_tensor(tStS.iterator + s_cols, tStS.layout) # offset 128 (second region) - - # Load from S0, Store to S1 (same layout, just different offset) - tmem_load_atom = cute.make_copy_atom( - tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS0) - sfw_idx = tidx % (32 * len(self.epilogue_warp_id)) - thr_load = tiled_tmem_load.get_slice(sfw_idx) - tTMEM_LOADtS = thr_load.partition_S(tStS0) - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tTMEM_LOADcS = thr_load.partition_D(tScS) - - tmem_store_atom = cute.make_copy_atom( - tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStS1) - thr_store = tiled_tmem_store.get_slice(sfw_idx) - tTMEM_STOREtS1 = thr_store.partition_D(tStS1) - tTMEM_STOREcS = thr_store.partition_S(tScS) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, 1)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ── TMA WARP ── - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_a, tAgA[(None,h.count)], tAsA[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_b, tBgB[(None,h.count)], tBsB[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_tma_coord.py b/tests/archive/test_tma_coord.py deleted file mode 100644 index e07cc854..00000000 --- a/tests/archive/test_tma_coord.py +++ /dev/null @@ -1,102 +0,0 @@ -"""Minimal TMA multi-tile test: does the TMA coordinate propagate at runtime? -This is a tiny kernel that loads K from two different tiles and checks if the data is different. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean -from cutlass.utils import LayoutEnum -import cuda.bindings.driver as cuda -import cutlass.torch as ct -import math - -HEAD_DIM = 64 - -class TmaMultiTileTest: - def __init__(self, s_k=256): - self.s_k = s_k - self.q_dtype = BFloat16 - self.cta_group = tcgen05.CtaGroup.ONE - self.kv_stage = 2 - - def _setup(self, qk_mma): - qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (128, 128, qk_ik * 4) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, self.q_dtype, self.kv_stage) - cta = cute.size(qk_mma.thr_id.shape) - k_s = cute.slice_(self.k_smem_s,(None,None,None,0)) - self.kv_tx_bytes = cute.size_in_bytes(self.q_dtype, k_s) * cta - - @cute.jit - def __call__(self, k, out_buf, stream): - b_major = LayoutEnum.from_tensor(k).mma_major_mode() - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, - cute.nvgpu.OperandMajorMode.K, b_major, - Float32, self.cta_group, (128,128), tcgen05.OperandSource.SMEM) - self._setup(qk_mma) - k_s = cute.slice_(self.k_smem_s,(None,None,None,0)) - tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_layout_vmnk.shape, qk_mma.thr_id), - k, k_s, self.qk_mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - self._kernel(qk_mma, tma_k, mK, out_buf, self.cluster_layout_vmnk, self.k_smem_s).launch( - grid=(1,1,1), block=[32,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, tma_k, mK, out_buf, cl_vmnk, k_smem_s): - tidx,_,_ = cute.arch.thread_idx() - cpasync.prefetch_descriptor(tma_k) - - sK = cute.make_tensor(self.q_dtype, cute.slice_(k_smem_s,(None,None,None,0)), byte_alignment=128) - - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None)) - n_kv_tiles = cute.size(gK, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - tCgK = qk_thr.partition_B(gK) - b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape) - tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3)) - - # Test different slice patterns - tBgK_nn = tBgK[(None,None,0,0)] # Mode 1 = GMEM tiles (should work for multi-tile) - tBgK_n0 = tBgK[(None,0,None,0)] # Mode 1 fixed to 0 (broken for multi-tile) - - print(f"tBgK shape: {cute.shape(tBgK)}", end="") - print(f" tBgK_nn shape: {cute.shape(tBgK_nn)}", end="") - print(f" tBgK_n0 shape: {cute.shape(tBgK_n0)}", end="") - print(f" n_kv_tiles: {n_kv_tiles}") - - # Load K from each tile using the (None,None,0,0) slice - # and write the first 4 elements to out_buf for verification - for kt in range(n_kv_tiles): - cute.copy(tma_k, tBgK_nn[(None, Int32(kt))], tBsK[(None, 0)], tma_bar_ptr=None) - # Read first 4 BF16 values from SMEM and write to out_buf - for i in range(4): - out_buf[kt * 4 + i] = sK[0, i, 0, 0, 0] - - -def test(): - for n in [256]: - torch.manual_seed(42) - k = torch.randn(n, HEAD_DIM, 1, dtype=torch.bfloat16, device='cuda') - out_buf = torch.zeros(n // 128 * 4, dtype=torch.bfloat16, device='cuda') - - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mO = ct.from_dlpack(out_buf).mark_layout_dynamic(leading_dim=ct.get_leading_dim(out_buf)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - kernel = TmaMultiTileTest(s_k=n) - print(f'n={n}: Compiling...', flush=True) - compiled = cute.compile(kernel, mK, mO, stream) - compiled(mK, mO, stream) - torch.cuda.synchronize() - - print(f'out_buf: {out_buf.tolist()}') - # Check if tile 0 and tile 1 data are different - tile0 = k[:128, 0, 0][:4].tolist() - tile1 = k[128:, 0, 0][:4].tolist() - print(f'tile0 first 4: {tile0}') - print(f'tile1 first 4: {tile1}') - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_tmem_addressing.py b/tests/archive/test_tmem_addressing.py deleted file mode 100644 index f59faeb4..00000000 --- a/tests/archive/test_tmem_addressing.py +++ /dev/null @@ -1,263 +0,0 @@ -""" -TMEM Addressing Test: verify offset computation from layouts. - -Allocates TMEM, computes offsets from QK accumulator and PV fragment sizes, -writes known values via tcgen05.st at each offset region, reads them back -via tcgen05.ld, and verifies correctness. No MMA, no softmax, no V load. - -This validates that our offset arithmetic is correct before wiring it into Stage B. -""" -import torch -import cutlass -import cutlass.cute as cute -import cutlass.utils as utils -import cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -import cuda.bindings.driver as cuda - - -class TmemAddressingTest: - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32 - self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16 - self.o_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn - self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4 - self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.tmem_alloc_sync_bar_id = 2 - self.tmem_dealloc_sync_bar_id = 3 - self.num_c_stage = 2 - - @cute.jit - def __call__(self, debug_buf: cute.Tensor, stream: cuda.CUstream): - self.a_dtype = BFloat16 - self.b_dtype = BFloat16 - self.a_major = cute.nvgpu.OperandMajorMode.K - self.b_major = cute.nvgpu.OperandMajorMode.K - self.c_layout = LayoutEnum.RowMajor - - # Create the same MMAs as Stage B to get the same fragment layouts - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, - tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, self.b_dtype, cute.nvgpu.OperandMajorMode.K, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, - tcgen05.OperandSource.TMEM) - - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - pv_inst_k = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (*self.mma_tiler_mn, pv_inst_k * 4) - self.mma_tiler = self.qk_mma_tiler - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], - self.qk_mma_tiler[2], - ) - self.cluster_layout_vmnk = cute.tiled_divide( - cute.make_layout((1, 1, 1)), (qk_mma.thr_id.shape,)) - - # Compute TMEM fragment sizes from layouts - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - qk_acc_cols = cute.size(tStS.layout, mode=[1]) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - pv_acc_cols = cute.size(tOtO.layout, mode=[1]) - - # P operand size: tilePlikeFP32 = qk_mma_tiler[1] * q_dtype.width // 32 - tilePlikeFP32 = self.qk_mma_tiler[1] * self.q_dtype.width // 32 - - # Compute offsets - tmem_s_offset = 0 - tmem_p_offset = qk_acc_cols # P right after QK accumulator - tmem_o_offset = qk_acc_cols + tilePlikeFP32 # O right after P - - # Total allocation - tmem_alloc_cols = tmem_o_offset + pv_acc_cols - - # JIT-time prints — these appear during compilation - print(f"[TMEM] qk_acc_cols = {qk_acc_cols}") - print(f"[TMEM] tilePlikeFP32 = {tilePlikeFP32}") - print(f"[TMEM] pv_acc_cols = {pv_acc_cols}") - print(f"[TMEM] tmem_s_offset = {tmem_s_offset}") - print(f"[TMEM] tmem_p_offset = {tmem_p_offset}") - print(f"[TMEM] tmem_o_offset = {tmem_o_offset}") - print(f"[TMEM] tmem_alloc_cols = {tmem_alloc_cols}") - - self._kernel( - qk_mma, pv_mma, tStS, tOtO, tmem_alloc_cols, - tmem_s_offset, tmem_p_offset, tmem_o_offset, tilePlikeFP32, - debug_buf, self.cluster_layout_vmnk - ).launch(grid=(1, 1, 1), block=[self.threads_per_cta, 1, 1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tStS, tOtO, tmem_alloc_cols, - tmem_s_offset, tmem_p_offset, tmem_o_offset, tilePlikeFP32, - debug_buf, cl_vmnk): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - @cute.struct - class SS: - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator() - st = smem.allocate(SS) - - tmem_bar = pipeline.NamedBarrier( - barrier_id=self.tmem_alloc_sync_bar_id, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator( - st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], - is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - pipeline.pipeline_init_wait(cluster_shape_mvnk=cl_vmnk) - - # ── MMA WARP: allocate TMEM, write test values ── - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) - - # Create TMEM tensors at computed offsets - # Scores region: write 1.0 - tStS0 = cute.make_tensor(tStS.iterator + tmem_s_offset, tStS.layout) - # P region: write 2.0 - tStS_P_layout = cute.composition(tStS.layout, cute.make_layout((128, tilePlikeFP32))) - tStS_P = cute.make_tensor(tStS.iterator + tmem_p_offset, tStS_P_layout) - # Output region: write 3.0 - tOtO0 = cute.make_tensor(tOtO.iterator + tmem_o_offset, tOtO.layout) - - # Use tcgen05.st to write known values into each region - # We'll use the store copy atom - sfw_idx = tidx % (32 * len(self.epilogue_warp_id)) - - # Store to scores region (value = 1.0) - tmem_store_atom = cute.make_copy_atom( - tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.acc_dtype) - tiled_store = tcgen05.make_tmem_copy(tmem_store_atom, tStS0) - thr_store = tiled_store.get_slice(sfw_idx) - tTMEM_STOREtS = thr_store.partition_D(tStS0) - # We need a source tensor with the same shape - tTMEM_STOREcS = thr_store.partition_S( - cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1]))) - - # Load from scores region (verify readback) - tmem_load_atom = cute.make_copy_atom( - tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.acc_dtype) - tiled_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS0) - thr_load = tiled_load.get_slice(sfw_idx) - tTMEM_LOADtS = thr_load.partition_S(tStS0) - - # The MMA warp doesn't do the ld/st — the epilogue warps do. - # For this test, just signal that TMEM is ready, epilogue will verify. - # But actually, MMA warp CAN write to TMEM via cute.fill or direct MMA. - # The simplest test: MMA warp issues a QK MMA with accumulate=False (known result), - # then epilogue warps tcgen05.ld from the scores region and dump to debug_buf. - - # For now: the MMA warp just signals and the epilogue does the verification. - # We'll write test values using tcgen05.st from epilogue warps (they have the copy atoms). - pass - - # ── EPILOGUE WARPS: allocate TMEM, write test values, read back ── - if warp_idx < self.mma_warp_id: - tmem.allocate(tmem_alloc_cols) - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) - sfw_idx = tidx % (32 * len(self.epilogue_warp_id)) - - # ── Write 1.0 to scores region ── - tStS0 = cute.make_tensor(tStS.iterator + tmem_s_offset, tStS.layout) - - tmem_store_atom = cute.make_copy_atom( - tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.acc_dtype) - tiled_store_s = tcgen05.make_tmem_copy(tmem_store_atom, tStS0) - thr_store_s = tiled_store_s.get_slice(sfw_idx) - tTMEM_STOREtS = thr_store_s.partition_D(tStS0) - tScS_s = qk_mma.get_slice(0).partition_C( - cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1]))) - tTMEM_STOREcS = thr_store_s.partition_S(tScS_s) - - # Create register tensor filled with 1.0 - tTMEM_STORErS = cute.make_rmem_tensor(tTMEM_STOREcS.shape, self.acc_dtype) - # Fill with 1.0 - for i in cutlass.range(cute.size(tTMEM_STORErS), unroll_full=True): - tTMEM_STORErS.store(i, cutlass.Float32(1.0)) - - cute.copy(tiled_store_s, tTMEM_STORErS, tTMEM_STOREtS) - cute.arch.fence_view_async_tmem_store() - - # ── Read back from scores region ── - tmem_load_atom = cute.make_copy_atom( - tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.acc_dtype) - tiled_load_s = tcgen05.make_tmem_copy(tmem_load_atom, tStS0) - thr_load_s = tiled_load_s.get_slice(sfw_idx) - tTMEM_LOADtS = thr_load_s.partition_S(tStS0) - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_mma.get_slice(0).partition_C(cS) - tTMEM_LOADcS = thr_load_s.partition_D(tScS) - - tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.acc_dtype) - cute.copy(tiled_load_s, tTMEM_LOADtS, tTMEM_LOADrS) - cute.arch.fence_view_async_tmem_load() - - # Dump one value per thread to debug_buf for verification - # debug_buf shape: (threads_per_cta,) Float32 - # Only epilogue warps (0..3, 128 threads) write - if tidx < 128: - val = tTMEM_LOADrS.load() - # Store first element of the loaded vector - debug_buf[tidx] = val # type: ignore - - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def test_tmem_addressing(): - device = torch.device("cuda") - debug_buf = torch.zeros(128, dtype=torch.float32, device=device) - - import cutlass.torch as cutlass_torch - mD = cutlass_torch.from_dlpack(debug_buf).mark_layout_dynamic( - leading_dim=cutlass_torch.get_leading_dim(debug_buf)) - - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - kernel = TmemAddressingTest(mma_tiler_mn=(128, 128)) - print("Compiling TMEM addressing test...", flush=True) - compiled = cute.compile(kernel, mD, stream) - print("Running...", flush=True) - compiled(mD, stream) - torch.cuda.synchronize() - - print("Debug buffer (first 16 values):", debug_buf[:16].tolist()) - # All values should be 1.0 if addressing is correct - nonzero = (debug_buf[:128] != 0).sum().item() - ones = (debug_buf[:128] == 1.0).sum().item() - print(f"Non-zero: {nonzero}/128, Ones: {ones}/128") - if nonzero > 0: - print("PASS: TMEM addressing works — read back non-zero values") - else: - print("FAIL: All zeros — TMEM addressing broken") - - -if __name__ == "__main__": - test_tmem_addressing() diff --git a/tests/archive/test_tmem_col2.py b/tests/archive/test_tmem_col2.py deleted file mode 100644 index 2c32228b..00000000 --- a/tests/archive/test_tmem_col2.py +++ /dev/null @@ -1,383 +0,0 @@ -""" -Minimal PV-only test: Load P from GMEM to TMEM via QK-style MMA, then PV from TMEM. -Step 1: QK MMA writes FP32 S to TMEM (we know this works) -Step 2: Softmax packing writes BF16 P to TMEM (test this) -Step 3: PV MMA reads BF16 P from TMEM and V from SMEM, produces O - -But to isolate the bug, let me test just the PV MMA in isolation. -I'll write known BF16 values to TMEM using the softmax packing path, -then immediately read them back using the PV A-fragment path, -and compare. - -Actually, the simplest isolation test: -1. Do QK MMA to get S in TMEM (cosine 0.999999 verified) -2. Do softmax packing: S → P in TMEM (at offset 32) -3. Skip PV entirely — read P from TMEM using the C-fragment composition LOAD path -4. Output P to GMEM and compare against S.to(BF16) - -This tests whether the softmax packing writes P correctly to the same TMEM -that the PV would read from. - -But we can't easily read P from TMEM using the standard epilogue path -because the epilogue expects FP32 accumulator data. - -Alternative: Use the PV MMA with V=I (identity). If P is correct, -then P @ I = P. But V needs to be MN-major and (128, 128), not (128, 64). -The output would be (128, 128) which doesn't match our (128, 64) c tensor. - -Let me use V that selects the first 64 columns: V[k, n] = delta(k, n) for k in [0,63]. -This gives P @ V = P[:, :64], and the output is (128, 64). -But V is (128, 128) in the MMA K,N dims. V[k, n] for k in [0,127], n in [0,63]. -Hmm, this is getting complicated. Let me just do the identity approach with a (128, 128) output. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class PvDiagKernel: - """QK + softmax packing + PV with V=I to isolate PV MMA correctness. - Output should be P = S.to(BF16), i.e. (Q@K^T).bfloat16() - With V=I, O = P @ I = P. - But V is (K=128, N=128) in the MMA. We need a 128x128 identity in MN-major. - Output tensor is (128, 128). - """ - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False # needed by epilogue_tma_store - self.epilog_sync_bar_id = 1 # needed by epilogue_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - # PV with V=I: output is (128, 128), same as QK - self.pv_mma_tiler = (self.qk_mma_tiler[0], self.qk_mma_tiler[1], self.qk_mma_tiler[1]) - # pv_mma_tiler = (128, 128, 128) since V is 128x128 - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - # ⛔⛔⛔ CRITICAL: num_tma_load_bytes MUST include ALL TMA-loaded tensors (Q + K + V). Missing V → DEADLOCK. See FOOTGUN #0 in README. - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - # PV with 128x128 output (V=I) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_tmem_col3.py b/tests/archive/test_tmem_col3.py deleted file mode 100644 index b86c20aa..00000000 --- a/tests/archive/test_tmem_col3.py +++ /dev/null @@ -1,392 +0,0 @@ -""" -Minimal PV-only test: Load P from GMEM to TMEM via QK-style MMA, then PV from TMEM. -Step 1: QK MMA writes FP32 S to TMEM (we know this works) -Step 2: Softmax packing writes BF16 P to TMEM (test this) -Step 3: PV MMA reads BF16 P from TMEM and V from SMEM, produces O - -But to isolate the bug, let me test just the PV MMA in isolation. -I'll write known BF16 values to TMEM using the softmax packing path, -then immediately read them back using the PV A-fragment path, -and compare. - -Actually, the simplest isolation test: -1. Do QK MMA to get S in TMEM (cosine 0.999999 verified) -2. Do softmax packing: S → P in TMEM (at offset 32) -3. Skip PV entirely — read P from TMEM using the C-fragment composition LOAD path -4. Output P to GMEM and compare against S.to(BF16) - -This tests whether the softmax packing writes P correctly to the same TMEM -that the PV would read from. - -But we can't easily read P from TMEM using the standard epilogue path -because the epilogue expects FP32 accumulator data. - -Alternative: Use the PV MMA with V=I (identity). If P is correct, -then P @ I = P. But V needs to be MN-major and (128, 128), not (128, 64). -The output would be (128, 128) which doesn't match our (128, 64) c tensor. - -Let me use V that selects the first 64 columns: V[k, n] = delta(k, n) for k in [0,63]. -This gives P @ V = P[:, :64], and the output is (128, 64). -But V is (128, 128) in the MMA K,N dims. V[k, n] for k in [0,127], n in [0,63]. -Hmm, this is getting complicated. Let me just do the identity approach with a (128, 128) output. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class TmemCol3Kernel: - """QK + softmax packing + PV with V=I to isolate PV MMA correctness. - Output should be P = S.to(BF16), i.e. (Q@K^T).bfloat16() - With V=I, O = P @ I = P. - But V is (K=128, N=128) in the MMA. We need a 128x128 identity in MN-major. - Output tensor is (128, 128). - """ - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False # needed by epilogue_tma_store - self.epilog_sync_bar_id = 1 # needed by epilogue_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = int(cute.size(qk_mma.shape_mnk, mode=[2])) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - # PV with V=I: output is (128, 128), same as QK - self.pv_mma_tiler = (self.qk_mma_tiler[0], self.qk_mma_tiler[1], self.qk_mma_tiler[1]) - # pv_mma_tiler = (128, 128, 128) since V is 128x128 - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - # DIAGNOSTIC: compare layout sizes and shapes - tP = cute.make_tensor(tStS.iterator, self.p_tmem_s.outer) - tOrP2 = pv_mma.get_slice(0).make_fragment_A(tP) - tOrP2_s = tOrP2[(None, None, None, 0)] - print(f"LAYOUT: tStS_size={int(cute.size(tStS))} tP_size={int(cute.size(tP))} tOrP2_size={int(cute.size(tOrP2))} tOrP2_s_size={int(cute.size(tOrP2_s))}") - -dims = [int(cute.size(tOrP2_s, mode=[i])) for i in range(tOrP2_s.layout.ndim)] -print(f"LAYOUT: tOrP2_s_dims={dims}") - print(f"LAYOUT: pv_mma_tiler={self.pv_mma_tiler} tilePlikeFP32={self.tilePlikeFP32}") - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - # ⛔⛔⛔ CRITICAL: num_tma_load_bytes MUST include ALL TMA-loaded tensors (Q + K + V). Missing V → DEADLOCK. See FOOTGUN #0 in README. - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - # PV with 128x128 output (V=I) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_tmem_col4.py b/tests/archive/test_tmem_col4.py deleted file mode 100644 index 0cc87e62..00000000 --- a/tests/archive/test_tmem_col4.py +++ /dev/null @@ -1,388 +0,0 @@ -""" -Minimal PV-only test: Load P from GMEM to TMEM via QK-style MMA, then PV from TMEM. -Step 1: QK MMA writes FP32 S to TMEM (we know this works) -Step 2: Softmax packing writes BF16 P to TMEM (test this) -Step 3: PV MMA reads BF16 P from TMEM and V from SMEM, produces O - -But to isolate the bug, let me test just the PV MMA in isolation. -I'll write known BF16 values to TMEM using the softmax packing path, -then immediately read them back using the PV A-fragment path, -and compare. - -Actually, the simplest isolation test: -1. Do QK MMA to get S in TMEM (cosine 0.999999 verified) -2. Do softmax packing: S → P in TMEM (at offset 32) -3. Skip PV entirely — read P from TMEM using the C-fragment composition LOAD path -4. Output P to GMEM and compare against S.to(BF16) - -This tests whether the softmax packing writes P correctly to the same TMEM -that the PV would read from. - -But we can't easily read P from TMEM using the standard epilogue path -because the epilogue expects FP32 accumulator data. - -Alternative: Use the PV MMA with V=I (identity). If P is correct, -then P @ I = P. But V needs to be MN-major and (128, 128), not (128, 64). -The output would be (128, 128) which doesn't match our (128, 64) c tensor. - -Let me use V that selects the first 64 columns: V[k, n] = delta(k, n) for k in [0,63]. -This gives P @ V = P[:, :64], and the output is (128, 64). -But V is (128, 128) in the MMA K,N dims. V[k, n] for k in [0,127], n in [0,63]. -Hmm, this is getting complicated. Let me just do the identity approach with a (128, 128) output. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class TmemCol4: - """QK + softmax packing + PV with V=I to isolate PV MMA correctness. - Output should be P = S.to(BF16), i.e. (Q@K^T).bfloat16() - With V=I, O = P @ I = P. - But V is (K=128, N=128) in the MMA. We need a 128x128 identity in MN-major. - Output tensor is (128, 128). - """ - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False # needed by epilogue_tma_store - self.epilog_sync_bar_id = 1 # needed by epilogue_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = int(cute.size(qk_mma.shape_mnk, mode=[2])) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - # PV with V=I: output is (128, 128), same as QK - self.pv_mma_tiler = (self.qk_mma_tiler[0], self.qk_mma_tiler[1], self.qk_mma_tiler[1]) - # pv_mma_tiler = (128, 128, 128) since V is 128x128 - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - tP = cute.make_tensor(tStS.iterator, self.p_tmem_s.outer) - tOrP2 = pv_mma.get_slice(0).make_fragment_A(tP) - tOrP2_s = tOrP2[(None, None, None, 0)] - ndim = tOrP2_s.layout.ndim - print(int(cute.size(tStS)), int(cute.size(tP)), int(cute.size(tOrP2)), int(cute.size(tOrP2_s)), int(ndim)) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - # ⛔⛔⛔ CRITICAL: num_tma_load_bytes MUST include ALL TMA-loaded tensors (Q + K + V). Missing V → DEADLOCK. See FOOTGUN #0 in README. - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - # PV with 128x128 output (V=I) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_tmem_col5.py b/tests/archive/test_tmem_col5.py deleted file mode 100644 index ce5e3901..00000000 --- a/tests/archive/test_tmem_col5.py +++ /dev/null @@ -1,392 +0,0 @@ -""" -Minimal PV-only test: Load P from GMEM to TMEM via QK-style MMA, then PV from TMEM. -Step 1: QK MMA writes FP32 S to TMEM (we know this works) -Step 2: Softmax packing writes BF16 P to TMEM (test this) -Step 3: PV MMA reads BF16 P from TMEM and V from SMEM, produces O - -But to isolate the bug, let me test just the PV MMA in isolation. -I'll write known BF16 values to TMEM using the softmax packing path, -then immediately read them back using the PV A-fragment path, -and compare. - -Actually, the simplest isolation test: -1. Do QK MMA to get S in TMEM (cosine 0.999999 verified) -2. Do softmax packing: S → P in TMEM (at offset 32) -3. Skip PV entirely — read P from TMEM using the C-fragment composition LOAD path -4. Output P to GMEM and compare against S.to(BF16) - -This tests whether the softmax packing writes P correctly to the same TMEM -that the PV would read from. - -But we can't easily read P from TMEM using the standard epilogue path -because the epilogue expects FP32 accumulator data. - -Alternative: Use the PV MMA with V=I (identity). If P is correct, -then P @ I = P. But V needs to be MN-major and (128, 128), not (128, 64). -The output would be (128, 128) which doesn't match our (128, 64) c tensor. - -Let me use V that selects the first 64 columns: V[k, n] = delta(k, n) for k in [0,63]. -This gives P @ V = P[:, :64], and the output is (128, 64). -But V is (128, 128) in the MMA K,N dims. V[k, n] for k in [0,127], n in [0,63]. -Hmm, this is getting complicated. Let me just do the identity approach with a (128, 128) output. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class TmemCol5: - """QK + softmax packing + PV with V=I to isolate PV MMA correctness. - Output should be P = S.to(BF16), i.e. (Q@K^T).bfloat16() - With V=I, O = P @ I = P. - But V is (K=128, N=128) in the MMA. We need a 128x128 identity in MN-major. - Output tensor is (128, 128). - """ - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False # needed by epilogue_tma_store - self.epilog_sync_bar_id = 1 # needed by epilogue_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = int(cute.size(qk_mma.shape_mnk, mode=[2])) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - # PV with V=I: output is (128, 128), same as QK - self.pv_mma_tiler = (self.qk_mma_tiler[0], self.qk_mma_tiler[1], self.qk_mma_tiler[1]) - # pv_mma_tiler = (128, 128, 128) since V is 128x128 - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - print(f"V: smem_size={int(cute.size(v_smem))} O: tOtO={int(cute.size(tOtO))} o_cols={o_cols} cta_tile={self.cta_tile_shape_mnk} epi_tile={self.epi_tile}") - tP = cute.make_tensor(tStS.iterator, self.p_tmem_s.outer) - tOrP2 = pv_mma.get_slice(0).make_fragment_A(tP) - tOrP2_s = tOrP2[(None, None, None, 0)] - sz0 = int(cute.size(tOrP2_s, mode=[0])) - sz1 = int(cute.size(tOrP2_s, mode=[1])) - sz2 = int(cute.size(tOrP2_s, mode=[2])) if cute.rank(tOrP2_s.layout) >= 3 else 0 - print(sz0, sz1, sz2, int(cute.size(tStS)), int(cute.size(tP)), int(cute.size(tOrP2_s))) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - # ⛔⛔⛔ CRITICAL: num_tma_load_bytes MUST include ALL TMA-loaded tensors (Q + K + V). Missing V → DEADLOCK. See FOOTGUN #0 in README. - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - # PV with 128x128 output (V=I) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_tmem_col5_16.py b/tests/archive/test_tmem_col5_16.py deleted file mode 100644 index 37e7bb3b..00000000 --- a/tests/archive/test_tmem_col5_16.py +++ /dev/null @@ -1,397 +0,0 @@ -""" -Minimal PV-only test: Load P from GMEM to TMEM via QK-style MMA, then PV from TMEM. -Step 1: QK MMA writes FP32 S to TMEM (we know this works) -Step 2: Softmax packing writes BF16 P to TMEM (test this) -Step 3: PV MMA reads BF16 P from TMEM and V from SMEM, produces O - -But to isolate the bug, let me test just the PV MMA in isolation. -I'll write known BF16 values to TMEM using the softmax packing path, -then immediately read them back using the PV A-fragment path, -and compare. - -Actually, the simplest isolation test: -1. Do QK MMA to get S in TMEM (cosine 0.999999 verified) -2. Do softmax packing: S → P in TMEM (at offset 32) -3. Skip PV entirely — read P from TMEM using the C-fragment composition LOAD path -4. Output P to GMEM and compare against S.to(BF16) - -This tests whether the softmax packing writes P correctly to the same TMEM -that the PV would read from. - -But we can't easily read P from TMEM using the standard epilogue path -because the epilogue expects FP32 accumulator data. - -Alternative: Use the PV MMA with V=I (identity). If P is correct, -then P @ I = P. But V needs to be MN-major and (128, 128), not (128, 64). -The output would be (128, 128) which doesn't match our (128, 64) c tensor. - -Let me use V that selects the first 64 columns: V[k, n] = delta(k, n) for k in [0,63]. -This gives P @ V = P[:, :64], and the output is (128, 64). -But V is (128, 128) in the MMA K,N dims. V[k, n] for k in [0,127], n in [0,63]. -Hmm, this is getting complicated. Let me just do the identity approach with a (128, 128) output. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class TmemCol5_16_32: - """QK + softmax packing + PV with V=I to isolate PV MMA correctness. - Output should be P = S.to(BF16), i.e. (Q@K^T).bfloat16() - With V=I, O = P @ I = P. - But V is (K=128, N=128) in the MMA. We need a 128x128 identity in MN-major. - Output tensor is (128, 128). - """ - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False # needed by epilogue_tma_store - self.epilog_sync_bar_id = 1 # needed by epilogue_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = int(cute.size(qk_mma.shape_mnk, mode=[2])) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - # PV with V=I: output is (128, 128), same as QK - self.pv_mma_tiler = (self.qk_mma_tiler[0], qk_inst_k, self.qk_mma_tiler[1]) - # pv_mma_tiler = (128, 128, 128) since V is 128x128 - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - (self.pv_mma_tiler[0], self.pv_mma_tiler[1], self.pv_mma_tiler[2]), False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - # V SMEM layout - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - print(f"V: smem_size={int(cute.size(v_smem))}") - # O C-fragment - print(f"O: tOtO_size={int(cute.size(tOtO))} o_cols={o_cols}") - # PV tiler - print(f"PV: mma_tiler={self.pv_mma_tiler} epi_tile={self.epi_tile} cta_tile={self.cta_tile_shape_mnk}") - tP = cute.make_tensor(tStS.iterator, self.p_tmem_s.outer) - tOrP2 = pv_mma.get_slice(0).make_fragment_A(tP) - tOrP2_s = tOrP2[(None, None, None, 0)] - sz0 = int(cute.size(tOrP2_s, mode=[0])) - sz1 = int(cute.size(tOrP2_s, mode=[1])) - sz2 = int(cute.size(tOrP2_s, mode=[2])) if cute.rank(tOrP2_s.layout) >= 3 else 0 - print(sz0, sz1, sz2, int(cute.size(tStS)), int(cute.size(tP)), int(cute.size(tOrP2_s))) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - # ⛔⛔⛔ CRITICAL: num_tma_load_bytes MUST include ALL TMA-loaded tensors (Q + K + V). Missing V → DEADLOCK. See FOOTGUN #0 in README. - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - # PV with 128x128 output (V=I) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, (128, 16), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_tmem_col5_32.py b/tests/archive/test_tmem_col5_32.py deleted file mode 100644 index 4969eb0d..00000000 --- a/tests/archive/test_tmem_col5_32.py +++ /dev/null @@ -1,390 +0,0 @@ -""" -Minimal PV-only test: Load P from GMEM to TMEM via QK-style MMA, then PV from TMEM. -Step 1: QK MMA writes FP32 S to TMEM (we know this works) -Step 2: Softmax packing writes BF16 P to TMEM (test this) -Step 3: PV MMA reads BF16 P from TMEM and V from SMEM, produces O - -But to isolate the bug, let me test just the PV MMA in isolation. -I'll write known BF16 values to TMEM using the softmax packing path, -then immediately read them back using the PV A-fragment path, -and compare. - -Actually, the simplest isolation test: -1. Do QK MMA to get S in TMEM (cosine 0.999999 verified) -2. Do softmax packing: S → P in TMEM (at offset 32) -3. Skip PV entirely — read P from TMEM using the C-fragment composition LOAD path -4. Output P to GMEM and compare against S.to(BF16) - -This tests whether the softmax packing writes P correctly to the same TMEM -that the PV would read from. - -But we can't easily read P from TMEM using the standard epilogue path -because the epilogue expects FP32 accumulator data. - -Alternative: Use the PV MMA with V=I (identity). If P is correct, -then P @ I = P. But V needs to be MN-major and (128, 128), not (128, 64). -The output would be (128, 128) which doesn't match our (128, 64) c tensor. - -Let me use V that selects the first 64 columns: V[k, n] = delta(k, n) for k in [0,63]. -This gives P @ V = P[:, :64], and the output is (128, 64). -But V is (128, 128) in the MMA K,N dims. V[k, n] for k in [0,127], n in [0,63]. -Hmm, this is getting complicated. Let me just do the identity approach with a (128, 128) output. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class TmemCol5_32: - """QK + softmax packing + PV with V=I to isolate PV MMA correctness. - Output should be P = S.to(BF16), i.e. (Q@K^T).bfloat16() - With V=I, O = P @ I = P. - But V is (K=128, N=128) in the MMA. We need a 128x128 identity in MN-major. - Output tensor is (128, 128). - """ - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False # needed by epilogue_tma_store - self.epilog_sync_bar_id = 1 # needed by epilogue_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = int(cute.size(qk_mma.shape_mnk, mode=[2])) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - # PV with V=I: output is (128, 128), same as QK - self.pv_mma_tiler = (self.qk_mma_tiler[0], qk_inst_k, self.qk_mma_tiler[1]) - # pv_mma_tiler = (128, 128, 128) since V is 128x128 - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - (self.pv_mma_tiler[0], self.pv_mma_tiler[1], self.pv_mma_tiler[2]), False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - tP = cute.make_tensor(tStS.iterator, self.p_tmem_s.outer) - tOrP2 = pv_mma.get_slice(0).make_fragment_A(tP) - tOrP2_s = tOrP2[(None, None, None, 0)] - sz0 = int(cute.size(tOrP2_s, mode=[0])) - sz1 = int(cute.size(tOrP2_s, mode=[1])) - sz2 = int(cute.size(tOrP2_s, mode=[2])) if cute.rank(tOrP2_s.layout) >= 3 else 0 - print(sz0, sz1, sz2, int(cute.size(tStS)), int(cute.size(tP)), int(cute.size(tOrP2_s))) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - # ⛔⛔⛔ CRITICAL: num_tma_load_bytes MUST include ALL TMA-loaded tensors (Q + K + V). Missing V → DEADLOCK. See FOOTGUN #0 in README. - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - # PV with 128x128 output (V=I) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, (128, 32), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_tmem_col_offset.py b/tests/archive/test_tmem_col_offset.py deleted file mode 100644 index e926aa08..00000000 --- a/tests/archive/test_tmem_col_offset.py +++ /dev/null @@ -1,388 +0,0 @@ -""" -Minimal PV-only test: Load P from GMEM to TMEM via QK-style MMA, then PV from TMEM. -Step 1: QK MMA writes FP32 S to TMEM (we know this works) -Step 2: Softmax packing writes BF16 P to TMEM (test this) -Step 3: PV MMA reads BF16 P from TMEM and V from SMEM, produces O - -But to isolate the bug, let me test just the PV MMA in isolation. -I'll write known BF16 values to TMEM using the softmax packing path, -then immediately read them back using the PV A-fragment path, -and compare. - -Actually, the simplest isolation test: -1. Do QK MMA to get S in TMEM (cosine 0.999999 verified) -2. Do softmax packing: S → P in TMEM (at offset 32) -3. Skip PV entirely — read P from TMEM using the C-fragment composition LOAD path -4. Output P to GMEM and compare against S.to(BF16) - -This tests whether the softmax packing writes P correctly to the same TMEM -that the PV would read from. - -But we can't easily read P from TMEM using the standard epilogue path -because the epilogue expects FP32 accumulator data. - -Alternative: Use the PV MMA with V=I (identity). If P is correct, -then P @ I = P. But V needs to be MN-major and (128, 128), not (128, 64). -The output would be (128, 128) which doesn't match our (128, 64) c tensor. - -Let me use V that selects the first 64 columns: V[k, n] = delta(k, n) for k in [0,63]. -This gives P @ V = P[:, :64], and the output is (128, 64). -But V is (128, 128) in the MMA K,N dims. V[k, n] for k in [0,127], n in [0,63]. -Hmm, this is getting complicated. Let me just do the identity approach with a (128, 128) output. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class TmemColOffsetKernel: - """QK + softmax packing + PV with V=I to isolate PV MMA correctness. - Output should be P = S.to(BF16), i.e. (Q@K^T).bfloat16() - With V=I, O = P @ I = P. - But V is (K=128, N=128) in the MMA. We need a 128x128 identity in MN-major. - Output tensor is (128, 128). - """ - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False # needed by epilogue_tma_store - self.epilog_sync_bar_id = 1 # needed by epilogue_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = int(cute.size(qk_mma.shape_mnk, mode=[2])) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - # PV with V=I: output is (128, 128), same as QK - self.pv_mma_tiler = (self.qk_mma_tiler[0], self.qk_mma_tiler[1], self.qk_mma_tiler[1]) - # pv_mma_tiler = (128, 128, 128) since V is 128x128 - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - # Compare TMEM column offsets: QK C-fragment S vs PV A-fragment P - tP = cute.make_tensor(tStS.iterator, self.p_tmem_s.outer) - tOrP2 = pv_mma.get_slice(0).make_fragment_A(tP) - pv_col = find_tmem_tensor_col_offset(tOrP2) - print(f"TMEM col: S_base={s_cols}, P_Afrag_at_S_base={pv_col}, O_base={o_cols}") - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - # ⛔⛔⛔ CRITICAL: num_tma_load_bytes MUST include ALL TMA-loaded tensors (Q + K + V). Missing V → DEADLOCK. See FOOTGUN #0 in README. - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - # PV with 128x128 output (V=I) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, self.p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_tmem_copy_roundtrip.py b/tests/archive/test_tmem_copy_roundtrip.py deleted file mode 100644 index 3885f16c..00000000 --- a/tests/archive/test_tmem_copy_roundtrip.py +++ /dev/null @@ -1,271 +0,0 @@ -"""Minimal TMEM ld→st roundtrip. FP32 only, no BF16 cast. -Uses the fmha pattern: load to register, store via make_rmem_tensor + recast + load/store.""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - -class TMEMCopyRoundtrip: - def __init__(self, mma_tiler_mn): - self.qk_acc_dtype = Float32; self.q_dtype = BFloat16; self.o_dtype = BFloat16 - self.c_dtype = BFloat16; self.acc_dtype = Float32 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2; self.use_2cta_instrs = False - self.epilog_sync_bar_id = 1 - - def _setup(self, qk_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - self.mma_tiler = self.qk_mma_tiler - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - c_layout = LayoutEnum.ROW_MAJOR; self.c_layout = c_layout - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, False, c_layout, self.o_dtype) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, c_layout, self.epi_tile, 2) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - self.s_cols = find_tmem_tensor_col_offset(tStS) - self.tmem_alloc_cols = self.s_cols * 2 - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, - LayoutEnum.from_tensor(a).mma_major_mode(), - LayoutEnum.from_tensor(b).mma_major_mode(), - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, - tcgen05.OperandSource.SMEM) - self._setup(qk_mma) - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - tma_a, tma_ta = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - a, a_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_b, tma_tb = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - b, b_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, tma_a, tma_ta, tma_b, tma_tb, tma_c, tma_tc, - self.cluster_layout_vmnk, self.a_smem_s, self.b_smem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, tma_a, mA, tma_b, mB, tma_c, mC, cl_vmnk, - a_smem_s, b_smem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_a); cpasync.prefetch_descriptor(tma_b); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, len(self.epilogue_warp_id)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=False, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sA = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sB = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gA = cute.local_tile(mA, cute.slice_(self.mma_tiler, (None,0,None)), (None,None,None)) - gB = cute.local_tile(mB, cute.slice_(self.mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gA, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - tCgA = qk_thr.partition_A(gA); tCgB = qk_thr.partition_B(gB); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsA, tAgA = cpasync.tma_partition(tma_a, 0, a_lay, cute.group_modes(sA,0,3), cute.group_modes(tCgA,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsB, tBgB = cpasync.tma_partition(tma_b, 0, b_lay, cute.group_modes(sB,0,3), cute.group_modes(tCgB,0,3)) - tAgA = tAgA[(None,0,None,0)]; tBgB = tBgB[(None,0,None,0)] - - tCrA = qk_mma.make_fragment_A(sA); tCrB = qk_mma.make_fragment_B(sB) - - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator, tStS.layout) - tStS1 = cute.make_tensor(tStS.iterator + self.s_cols, tStS.layout) - - # Copy atoms - tmem_load_atom = cute.make_copy_atom( - tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS0) - sfw_idx = tidx % (32 * len(self.epilogue_warp_id)) - thr_load = tiled_tmem_load.get_slice(sfw_idx) - tTMEM_LOADtS0 = thr_load.partition_S(tStS0) - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tTMEM_LOADcS = thr_load.partition_D(tScS) - - tmem_store_atom = cute.make_copy_atom( - tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStS1) - thr_store = tiled_tmem_store.get_slice(sfw_idx) - tTMEM_STOREtS1 = thr_store.partition_D(tStS1) - tTMEM_STOREcS = thr_store.partition_S(tScS) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, 1)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ── TMA WARP ── - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_a, tAgA[(None,h.count)], tAsA[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_b, tBgB[(None,h.count)], tBsB[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_tmem_debug.py b/tests/archive/test_tmem_debug.py deleted file mode 100644 index cc78e78b..00000000 --- a/tests/archive/test_tmem_debug.py +++ /dev/null @@ -1,397 +0,0 @@ -""" -Minimal PV-only test: Load P from GMEM to TMEM via QK-style MMA, then PV from TMEM. -Step 1: QK MMA writes FP32 S to TMEM (we know this works) -Step 2: Softmax packing writes BF16 P to TMEM (test this) -Step 3: PV MMA reads BF16 P from TMEM and V from SMEM, produces O - -But to isolate the bug, let me test just the PV MMA in isolation. -I'll write known BF16 values to TMEM using the softmax packing path, -then immediately read them back using the PV A-fragment path, -and compare. - -Actually, the simplest isolation test: -1. Do QK MMA to get S in TMEM (cosine 0.999999 verified) -2. Do softmax packing: S → P in TMEM (at offset 32) -3. Skip PV entirely — read P from TMEM using the C-fragment composition LOAD path -4. Output P to GMEM and compare against S.to(BF16) - -This tests whether the softmax packing writes P correctly to the same TMEM -that the PV would read from. - -But we can't easily read P from TMEM using the standard epilogue path -because the epilogue expects FP32 accumulator data. - -Alternative: Use the PV MMA with V=I (identity). If P is correct, -then P @ I = P. But V needs to be MN-major and (128, 128), not (128, 64). -The output would be (128, 128) which doesn't match our (128, 64) c tensor. - -Let me use V that selects the first 64 columns: V[k, n] = delta(k, n) for k in [0,63]. -This gives P @ V = P[:, :64], and the output is (128, 64). -But V is (128, 128) in the MMA K,N dims. V[k, n] for k in [0,127], n in [0,63]. -Hmm, this is getting complicated. Let me just do the identity approach with a (128, 128) output. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class PvDiagKernel: - """QK + softmax packing + PV with V=I to isolate PV MMA correctness. - Output should be P = S.to(BF16), i.e. (Q@K^T).bfloat16() - With V=I, O = P @ I = P. - But V is (K=128, N=128) in the MMA. We need a 128x128 identity in MN-major. - Output tensor is (128, 128). - """ - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False # needed by epilogue_tma_store - self.epilog_sync_bar_id = 1 # needed by epilogue_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - # PV with V=I: output is (128, 128), same as QK - self.pv_mma_tiler = (self.qk_mma_tiler[0], self.qk_mma_tiler[1], self.qk_mma_tiler[1]) - # pv_mma_tiler = (128, 128, 128) since V is 128x128 - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - # Debug: print layout shapes - pv_thr2 = pv_mma.get_slice(0) - pv_acc2 = pv_thr2.partition_shape_C(self.pv_mma_tiler[:2]) - print(f"DEBUG: tStS.layout={tStS.layout}, p_tmem_s.outer shape mode 0={cute.size(p_tmem_s.outer, mode=[0])}, mode 1={cute.size(p_tmem_s.outer, mode=[1])}") - tP2 = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP2 = pv_thr2.make_fragment_A(tP2) - print(f"DEBUG: tOrP.layout={tOrP2.layout}, size={cute.size(tOrP2)}") - print(f"DEBUG: tStS size={cute.size(tStS)}, pv_acc_shape={pv_acc2}") - # Check how many K phases PV has - tOrP_sliced = tOrP2[(None, None, None, 0)] - print(f"DEBUG: tOrP_sliced size={cute.size(tOrP_sliced)}, ndim={tOrP_sliced.layout.ndim}") - if tOrP_sliced.layout.ndim >= 2: - print(f"DEBUG: tOrP_sliced shape: {[cute.size(tOrP_sliced, mode=[i]) for i in range(tOrP_sliced.layout.ndim)]}") - - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - # ⛔⛔⛔ CRITICAL: num_tma_load_bytes MUST include ALL TMA-loaded tensors (Q + K + V). Missing V → DEADLOCK. See FOOTGUN #0 in README. - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - # PV with 128x128 output (V=I) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_tmem_debug2.py b/tests/archive/test_tmem_debug2.py deleted file mode 100644 index 5ddea753..00000000 --- a/tests/archive/test_tmem_debug2.py +++ /dev/null @@ -1,392 +0,0 @@ -""" -Minimal PV-only test: Load P from GMEM to TMEM via QK-style MMA, then PV from TMEM. -Step 1: QK MMA writes FP32 S to TMEM (we know this works) -Step 2: Softmax packing writes BF16 P to TMEM (test this) -Step 3: PV MMA reads BF16 P from TMEM and V from SMEM, produces O - -But to isolate the bug, let me test just the PV MMA in isolation. -I'll write known BF16 values to TMEM using the softmax packing path, -then immediately read them back using the PV A-fragment path, -and compare. - -Actually, the simplest isolation test: -1. Do QK MMA to get S in TMEM (cosine 0.999999 verified) -2. Do softmax packing: S → P in TMEM (at offset 32) -3. Skip PV entirely — read P from TMEM using the C-fragment composition LOAD path -4. Output P to GMEM and compare against S.to(BF16) - -This tests whether the softmax packing writes P correctly to the same TMEM -that the PV would read from. - -But we can't easily read P from TMEM using the standard epilogue path -because the epilogue expects FP32 accumulator data. - -Alternative: Use the PV MMA with V=I (identity). If P is correct, -then P @ I = P. But V needs to be MN-major and (128, 128), not (128, 64). -The output would be (128, 128) which doesn't match our (128, 64) c tensor. - -Let me use V that selects the first 64 columns: V[k, n] = delta(k, n) for k in [0,63]. -This gives P @ V = P[:, :64], and the output is (128, 64). -But V is (128, 128) in the MMA K,N dims. V[k, n] for k in [0,127], n in [0,63]. -Hmm, this is getting complicated. Let me just do the identity approach with a (128, 128) output. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class PvDiagKernel: - """QK + softmax packing + PV with V=I to isolate PV MMA correctness. - Output should be P = S.to(BF16), i.e. (Q@K^T).bfloat16() - With V=I, O = P @ I = P. - But V is (K=128, N=128) in the MMA. We need a 128x128 identity in MN-major. - Output tensor is (128, 128). - """ - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False # needed by epilogue_tma_store - self.epilog_sync_bar_id = 1 # needed by epilogue_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - # PV with V=I: output is (128, 128), same as QK - self.pv_mma_tiler = (self.qk_mma_tiler[0], self.qk_mma_tiler[1], self.qk_mma_tiler[1]) - # pv_mma_tiler = (128, 128, 128) since V is 128x128 - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - pv_thr2 = pv_mma.get_slice(0) - pv_acc2 = pv_thr2.partition_shape_C(self.pv_mma_tiler[:2]) - tP2 = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP2 = pv_thr2.make_fragment_A(tP2) - tOrP2_s = tOrP2[(None, None, None, 0)] - print(int(cute.size(tStS)), int(cute.size(tP2)), int(cute.size(tOrP2)), int(cute.size(tOrP2_s))) - if tOrP2_s.layout.ndim >= 2: - print(int(cute.size(tOrP2_s, mode=[0])), int(cute.size(tOrP2_s, mode=[1])), int(cute.size(tOrP2_s, mode=[2])) if tOrP2_s.layout.ndim >= 3 else "2d") - - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - # ⛔⛔⛔ CRITICAL: num_tma_load_bytes MUST include ALL TMA-loaded tensors (Q + K + V). Missing V → DEADLOCK. See FOOTGUN #0 in README. - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - # PV with 128x128 output (V=I) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_tmem_fp32_roundtrip.py b/tests/archive/test_tmem_fp32_roundtrip.py deleted file mode 100644 index 0c2dd017..00000000 --- a/tests/archive/test_tmem_fp32_roundtrip.py +++ /dev/null @@ -1,255 +0,0 @@ -"""Minimal: Q@K^T → TMEM, ld FP32 from S0, st FP32 to S1, epi reads S1. -NO bf16 cast at all. Pure FP32 ld→st roundtrip.""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - -class FP32Roundtrip: - def __init__(self, mma_tiler_mn): - self.qk_acc_dtype = Float32; self.q_dtype = BFloat16; self.o_dtype = BFloat16 - self.c_dtype = BFloat16; self.acc_dtype = Float32 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2; self.use_2cta_instrs = False - self.epilog_sync_bar_id = 1 - - def _setup(self, qk_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - self.mma_tiler = self.qk_mma_tiler - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - c_layout = LayoutEnum.ROW_MAJOR; self.c_layout = c_layout - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, False, c_layout, self.o_dtype) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, c_layout, self.epi_tile, 2) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - self.s_cols = find_tmem_tensor_col_offset(tStS) - self.tmem_alloc_cols = self.s_cols * 2 - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, - LayoutEnum.from_tensor(a).mma_major_mode(), - LayoutEnum.from_tensor(b).mma_major_mode(), - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, - tcgen05.OperandSource.SMEM) - self._setup(qk_mma) - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - tma_a, tma_ta = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - a, a_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_b, tma_tb = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - b, b_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - self._kernel(qk_mma, tma_a, tma_ta, tma_b, tma_tb, tma_c, tma_tc, - self.cluster_layout_vmnk, self.a_smem_s, self.b_smem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, tma_a, mA, tma_b, mB, tma_c, mC, cl_vmnk, - a_smem_s, b_smem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_a); cpasync.prefetch_descriptor(tma_b); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, len(self.epilogue_warp_id)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=False, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sA = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sB = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gA = cute.local_tile(mA, cute.slice_(self.mma_tiler, (None,0,None)), (None,None,None)) - gB = cute.local_tile(mB, cute.slice_(self.mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gA, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - tCgA = qk_thr.partition_A(gA); tCgB = qk_thr.partition_B(gB); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsA, tAgA = cpasync.tma_partition(tma_a, 0, a_lay, cute.group_modes(sA,0,3), cute.group_modes(tCgA,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsB, tBgB = cpasync.tma_partition(tma_b, 0, b_lay, cute.group_modes(sB,0,3), cute.group_modes(tCgB,0,3)) - tAgA = tAgA[(None,0,None,0)]; tBgB = tBgB[(None,0,None,0)] - tCrA = qk_mma.make_fragment_A(sA); tCrB = qk_mma.make_fragment_B(sB) - - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator, tStS.layout) - tStS1 = cute.make_tensor(tStS.iterator + self.s_cols, tStS.layout) - - # Single tiled copy that can both ld and st - tmem_ld_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tmem_st_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_ld = tcgen05.make_tmem_copy(tmem_ld_atom, tStS0) - tiled_st = tcgen05.make_tmem_copy(tmem_st_atom, tStS1) - sfw = tidx % (32 * len(self.epilogue_warp_id)) - thr_ld = tiled_ld.get_slice(sfw) - thr_st = tiled_st.get_slice(sfw) - - tTMEM_LDtS = thr_ld.partition_S(tStS0) - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tTMEM_LDcS = thr_ld.partition_D(tScS) - tTMEM_STtS = thr_st.partition_D(tStS1) - tTMEM_STcS = thr_st.partition_S(tScS) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, 1)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # TMA - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_a, tAgA[(None,h.count)], tAsA[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_b, tBgB[(None,h.count)], tBsB[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_tmem_layout_diag.py b/tests/archive/test_tmem_layout_diag.py deleted file mode 100644 index aa293de8..00000000 --- a/tests/archive/test_tmem_layout_diag.py +++ /dev/null @@ -1,81 +0,0 @@ -""" -Diagnostic: Compare TMEM column mapping between QK C-fragment and PV A-fragment. -Write the TMEM column index for each logical element to GMEM. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils -from cutlass.cute.nvgpu import tcgen05 -from cutlass import Float32, BFloat16, Int32 -from cutlass.utils import LayoutEnum -import cutlass.torch as ct -import cuda.bindings.driver as cuda - - -class TmemLayoutDiag: - def __init__(self): - self.threads_per_cta = 128 # just 1 warp for simplicity - - @cute.jit - def __call__(self, qk_cols, pv16_cols, stream): - # qk_cols: GMEM tensor (128,) to store QK C-fragment TMEM column indices - # pv16_cols: GMEM tensor (128,) to store PV A-fragment TMEM column indices - - # Create QK MMA (128,128) - qk_mma = utils.sm100.make_trivial_tiled_mma( - BFloat16, BFloat16, LayoutEnum.ROW_MAJOR, LayoutEnum.ROW_MAJOR, - Float32, tcgen05.CtaGroup.ONE, (128, 128), tcgen05.OperandSource.SMEM) - # Create PV (128,16) MMA - pv16_mma = utils.sm100.make_trivial_tiled_mma( - BFloat16, BFloat16, cute.nvgpu.OperandMajorMode.K, LayoutEnum.ROW_MAJOR, - Float32, tcgen05.CtaGroup.ONE, (128, 16), tcgen05.OperandSource.TMEM) - - qk_inst_k = int(cute.size(qk_mma.shape_mnk, mode=[2])) - qk_mma_tiler = (128, 128, int(qk_inst_k * 4)) - pv16_mma_tiler = (128, qk_inst_k, 128) - - # Create layouts - p_tmem_16 = utils.sm100.make_smem_layout_a(pv16_mma, pv16_mma_tiler, BFloat16, 1) - - # Create tStS (QK C-fragment) - qk_thr = qk_mma.get_slice(0) - qk_acc = qk_thr.partition_shape_C(qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc) - - # Create tP for PV16 - tP16 = cute.make_tensor(tStS.iterator, p_tmem_16.outer) - pv16_thr = pv16_mma.get_slice(0) - tOrP16 = pv16_thr.make_fragment_A(tP16) - - # Write sizes to GMEM - # For thread 0 only - tidx, _, _ = cute.arch.thread_idx() - if tidx == 0: - # Write tStS size and tP16 size - qk_cols[0] = Int32(cute.size(tStS)) - pv16_cols[0] = Int32(cute.size(tP16)) - # Write qk_inst_k - qk_cols[1] = Int32(qk_inst_k) - # Write p_tmem_16 outer shape - pv16_cols[1] = Int32(cute.size(p_tmem_16.outer, mode=[0])) - pv16_cols[2] = Int32(cute.size(p_tmem_16.outer, mode=[1])) if p_tmem_16.outer.ndim >= 2 else Int32(0) - # Write QK acc shape - qk_cols[2] = Int32(qk_acc[0]) if hasattr(qk_acc, '__getitem__') else Int32(0) - - def run(self): - qk_cols = torch.zeros(4, dtype=torch.int32, device='cuda') - pv16_cols = torch.zeros(4, dtype=torch.int32, device='cuda') - - mQ = ct.from_dlpack(qk_cols).mark_layout_dynamic() - mP = ct.from_dlpack(pv16_cols).mark_layout_dynamic() - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - print("Compiling...", flush=True) - compiled = cute.compile(self, mQ, mP, stream) - compiled(mQ, mP, stream) - torch.cuda.synchronize() - - print(f"QK results: {qk_cols.tolist()}") - print(f"PV16 results: {pv16_cols.tolist()}") - - -if __name__ == "__main__": - TmemLayoutDiag().run() diff --git a/tests/archive/test_tmem_pure_fp32.py b/tests/archive/test_tmem_pure_fp32.py deleted file mode 100644 index 9bc9e98a..00000000 --- a/tests/archive/test_tmem_pure_fp32.py +++ /dev/null @@ -1,237 +0,0 @@ -"""Absolute minimal: ld FP32 from S0, st FP32 to S1, epi reads S1. -No recast, no BF16, no packing. Pure FP32 copy between TMEM regions.""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - -class PureFP32Copy: - def __init__(self, mma_tiler_mn): - self.qk_acc_dtype = Float32; self.q_dtype = BFloat16; self.o_dtype = BFloat16 - self.c_dtype = BFloat16; self.acc_dtype = Float32 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2; self.use_2cta_instrs = False - self.epilog_sync_bar_id = 1 - - def _setup(self, qk_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - self.mma_tiler = self.qk_mma_tiler - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - c_layout = LayoutEnum.ROW_MAJOR; self.c_layout = c_layout - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, False, c_layout, self.o_dtype) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, c_layout, self.epi_tile, 2) - self.num_ab_stage = 1; self.num_acc_stage = 1 - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - self.s_cols = find_tmem_tensor_col_offset(tStS) - self.tmem_alloc_cols = self.s_cols * 2 - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream): - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, - LayoutEnum.from_tensor(a).mma_major_mode(), - LayoutEnum.from_tensor(b).mma_major_mode(), - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, - tcgen05.OperandSource.SMEM) - self._setup(qk_mma) - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - tma_a, tma_ta = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - a, a_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_b, tma_tb = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - b, b_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - self._kernel(qk_mma, tma_a, tma_ta, tma_b, tma_tb, tma_c, tma_tc, - self.cluster_layout_vmnk, self.a_smem_s, self.b_smem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, tma_a, mA, tma_b, mB, tma_c, mC, cl_vmnk, - a_smem_s, b_smem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_a); cpasync.prefetch_descriptor(tma_b); cpasync.prefetch_descriptor(tma_c) - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, len(self.epilogue_warp_id)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=False, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - sA = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sB = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - gA = cute.local_tile(mA, cute.slice_(self.mma_tiler, (None,0,None)), (None,None,None)) - gB = cute.local_tile(mB, cute.slice_(self.mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gA, mode=[3]) - qk_thr = qk_mma.get_slice(0) - tCgA = qk_thr.partition_A(gA); tCgB = qk_thr.partition_B(gB); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsA, tAgA = cpasync.tma_partition(tma_a, 0, a_lay, cute.group_modes(sA,0,3), cute.group_modes(tCgA,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsB, tBgB = cpasync.tma_partition(tma_b, 0, b_lay, cute.group_modes(sB,0,3), cute.group_modes(tCgB,0,3)) - tAgA = tAgA[(None,0,None,0)]; tBgB = tBgB[(None,0,None,0)] - tCrA = qk_mma.make_fragment_A(sA); tCrB = qk_mma.make_fragment_B(sB) - qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator, tStS.layout) - tStS1 = cute.make_tensor(tStS.iterator + self.s_cols, tStS.layout) - - # LD and ST on same layout - tmem_ld = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tmem_st = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_ld = tcgen05.make_tmem_copy(tmem_ld, tStS0) - tiled_st = tcgen05.make_tmem_copy(tmem_st, tStS1) - sfw = tidx % (32 * len(self.epilogue_warp_id)) - thr_ld = tiled_ld.get_slice(sfw) - thr_st = tiled_st.get_slice(sfw) - tLdS = thr_ld.partition_S(tStS0) - tStS = thr_st.partition_D(tStS1) - cS_id = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS_id) - tLdcS = thr_ld.partition_D(tScS) - tStcS = thr_st.partition_S(tScS) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, 1)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_a, tAgA[(None,h.count)], tAsA[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_b, tBgB[(None,h.count)], tBsB[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_uniform_fp4.py b/tests/archive/test_uniform_fp4.py deleted file mode 100644 index 08e7e574..00000000 --- a/tests/archive/test_uniform_fp4.py +++ /dev/null @@ -1,30 +0,0 @@ -"""Test: uniform FP4 + uniform SF, different from all-ones. -If all E2M1 values are the same (e.g. value 3 = 1.5) and all SF=1.0, -then x = 1.5 for all elements, w = 1.5 for all elements. -GEMM output = (1.5^2) * K = 2.25 * 32 = 72.0 for every element. -""" -import torch, sys -sys.path.insert(0, 'src') -from nvfp4_megamoe_kernel.cutlass_nvfp4_gemm.kernel import cutlass_nvfp4_blockscaled_gemm - -device = "cuda" -M, N, K = 1, 32, 32 - -# Create packed FP4 where every nibble = 3 (E2M1 value 1.5) -# Packing: (nibbles[..., 1] << 4) | nibbles[..., 0] -# For both nibbles = 3: byte = (3 << 4) | 3 = 0x33 -byte_val = (3 << 4) | 3 # 0x33 -x_fp4 = torch.full((M, K // 2), byte_val, dtype=torch.int8, device=device) -w_fp4 = torch.full((K // 2, N), byte_val, dtype=torch.int8, device=device) - -# Uniform SF = 1.0 -x_sf = torch.ones(M, K // 16, dtype=torch.float8_e4m3fn, device=device) -w_sf = torch.ones(K // 16, N, dtype=torch.float8_e4m3fn, device=device) - -out = cutlass_nvfp4_blockscaled_gemm(x_fp4, x_sf, w_fp4, w_sf, M, N, K, alpha=1.0) - -# Reference: all x = 1.5, all w = 1.5, output = 1.5 * 1.5 * 32 = 72.0 -print(f"NVFP4 output first 8: {out[0, :8].tolist()}") -print(f"Expected: 72.0 for all elements") -print(f"Actual mean: {out.float().mean().item():.4f}") -print(f"All same? {torch.allclose(out, out[0,0].expand_as(out), atol=0.01)}") diff --git a/tests/archive/test_v28c_noepi.py b/tests/archive/test_v28c_noepi.py deleted file mode 100644 index b220af59..00000000 --- a/tests/archive/test_v28c_noepi.py +++ /dev/null @@ -1,315 +0,0 @@ -""" -Stage B v28c: Debug deadlock. -Skip the epilogue. Just do QK + softmax packing + PV MMA. -Verify PV MMA completes (no deadlock). -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - - -class StageBNoEpi: - def __init__(self, pv_mn, use_2cta_instrs=False): - self.pv_mn = pv_mn - self.qk_acc_dtype = Float32; self.q_dtype = BFloat16; self.b_dtype = BFloat16 - self.o_dtype = BFloat16; self.c_dtype = BFloat16; self.acc_dtype = Float32 - self.use_2cta_instrs = use_2cta_instrs - self.mma_tiler_mn = (128, 128); self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.tmem_alloc_sync_bar_id = 2; self.tmem_dealloc_sync_bar_id = 3 - self.num_ab_stage = 1; self.num_acc_stage = 1 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - self.pv_mma_tiler = (self.qk_mma_tiler[0], self.qk_mma_tiler[2], self.qk_mma_tiler[1]) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.b_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.b_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - print(f"[v28c] qk_mma_tiler={self.qk_mma_tiler} pv_mma_tiler={self.pv_mma_tiler}") - print(f"[v28c] s_cols={s_cols} o_cols={o_cols} num_tmem={self.num_tmem_alloc_cols}") - - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.b_dtype, b_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q: cute.Tensor, k: cute.Tensor, v: cute.Tensor, stream: cuda.CUstream): - self.q_dtype = q.element_type; self.b_dtype = k.element_type - self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.b_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - pv_mma_tiler = (qk_mma_tiler[0], qk_mma_tiler[2], qk_mma_tiler[1]) - - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.b_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, self.pv_mn, tcgen05.OperandSource.TMEM) - - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - self.cluster_layout_vmnk, self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - tmem_bar = pipeline.NamedBarrier(barrier_id=self.tmem_alloc_sync_bar_id, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.b_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.b_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ===== TMA LOAD WARP ===== - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1> 4) & 0xF).long()] - O, I2 = w.shape; I = I2*2 - u = torch.empty(O, I, dtype=torch.float32, device=d) - u[:,0::2] = lo; u[:,1::2] = hi - bs = sf.float().repeat_interleave(16, dim=1)[:O,:I] - return (u * bs * gs).to(torch.bfloat16) - -def rms(x, w, eps=1e-6): - v = x.float().pow(2).mean(-1, keepdim=True) - return (w.float() * (x * torch.rsqrt(v+eps)).float()).to(x.dtype) - -def make_runner(w, sf, gs_t, inf, outf, fused=False, lw=None): - from dsv4.layers.linear import Nvfp4Linear - fp4 = w.view(torch.float4_e2m1fn_x2).permute(1,0).contiguous() - s = sf.to(torch.float8_e4m3fn) if sf.dtype != torch.float8_e4m3fn else sf - s = s.permute(1,0).contiguous() - if fused and gs_t.numel() == 2: - g1,g2 = gs_t[0].item(), gs_t[1].item(); gs = max(g1,g2) - if g1 != g2: - s32 = s.float(); sp = lw[0] if lw else outf//2 - s32[:sp] *= g1/gs; s32[sp:] *= g2/gs; s = s32.to(torch.float8_e4m3fn) - else: - gs = gs_t.max().item() if gs_t.numel() > 1 else gs_t.item() - r = Nvfp4Linear(in_features=inf, out_features=outf, max_num_tokens=8192, device=str(w.device)) - r.fp4 = [fp4]; r.sf = [s]; r.gs = [gs] - r.finalize_weights(); r._ensure_initialized() - return r - -def apply_gptj_rope(x, positions, cos_sin, nope, rope): - if rope == 0 or x.numel() == 0: return x - half = rope // 2 - cos = cos_sin[positions, :half].to(x.dtype) - sin = cos_sin[positions, half:2*half].to(x.dtype) - if x.dim() == 3: cos = cos.unsqueeze(1); sin = sin.unsqueeze(1) - x_rope = x[..., nope:].clone() - even = x_rope[..., 0::2]; odd = x_rope[..., 1::2] - out = x.clone() - out[..., nope:][..., 0::2] = even * cos - odd * sin - out[..., nope:][..., 1::2] = even * sin + odd * cos - return out - -def apply_inv_gptj_rope(x, positions, cos_sin, nope, rope): - """Inverse RoPE: same as forward but sin → -sin.""" - if rope == 0 or x.numel() == 0: return x - half = rope // 2 - cos = cos_sin[positions, :half].to(x.dtype) - sin = cos_sin[positions, half:2*half].to(x.dtype) - if x.dim() == 3: cos = cos.unsqueeze(1); sin = sin.unsqueeze(1) - x_rope = x[..., nope:].clone() - even = x_rope[..., 0::2]; odd = x_rope[..., 1::2] - out = x.clone() - out[..., nope:][..., 0::2] = even * cos + odd * sin - out[..., nope:][..., 1::2] = -even * sin + odd * cos - return out - -def build_cos_sin(max_pos=4096, rope_dim=ROPE): - half = rope_dim // 2 - inv_freq = 1.0 / (10000.0 ** (torch.arange(0, half, dtype=torch.float32) / half)) - freqs = torch.outer(torch.arange(max_pos, dtype=torch.float32), inv_freq) - return torch.cat([freqs.cos(), freqs.sin()], dim=-1) - - -def swa_attention(q, kv, scale, window_size=WINDOW): - """Sliding window attention using SDPA. - - q: (T, NH, HD) with RoPE - kv: (T, HD) shared KV latent - For SWA: attend to last window_size tokens only. - """ - T, NH, HD = q.shape - if T <= window_size: - # Full attention within window - return full_causal_attention(q, kv, scale) - - # For long sequences, only attend to window - # This is a simplified version — production would use paged cache - q_2d = q.reshape(T * NH, HD) - kv_exp = kv.unsqueeze(1).expand(-1, NH, -1).contiguous() - k_2d = kv_exp.permute(1, 0, 2).unsqueeze(1).expand(NH, T, T, -1).contiguous().reshape(T * NH, T, HD) - v_2d = k_2d.clone() - scores = torch.matmul(q_2d.unsqueeze(1), k_2d.transpose(-1, -2)) * scale - query_pos = torch.arange(T, device=q.device).unsqueeze(1).repeat(1, NH).reshape(T * NH) - kv_pos = torch.arange(T, device=q.device).unsqueeze(0) - causal = kv_pos <= query_pos.unsqueeze(1) - window = kv_pos >= (query_pos.unsqueeze(1) - window_size + 1) - mask = causal & window - scores = scores.squeeze(1).masked_fill(~mask, float('-inf')) - weights = F.softmax(scores.float(), dim=-1).to(q.dtype) - out = torch.matmul(weights.unsqueeze(1), v_2d).squeeze(1) - return out.reshape(T, NH, HD) - - -def full_causal_attention(q, kv, scale): - """Full causal self-attention (for testing with T <= window_size).""" - T, NH, HD = q.shape - q_2d = q.reshape(T * NH, HD) - kv_exp = kv.unsqueeze(1).expand(-1, NH, -1).contiguous() - k_2d = kv_exp.permute(1, 0, 2).unsqueeze(1).expand(NH, T, T, -1).contiguous().reshape(T * NH, T, HD) - v_2d = k_2d.clone() - scores = torch.matmul(q_2d.unsqueeze(1), k_2d.transpose(-1, -2)) * scale - query_pos = torch.arange(T, device=q.device).unsqueeze(1).repeat(1, NH).reshape(T * NH) - kv_pos = torch.arange(T, device=q.device).unsqueeze(0) - causal = kv_pos <= query_pos.unsqueeze(1) - scores = scores.squeeze(1).masked_fill(~causal, float('-inf')) - weights = F.softmax(scores.float(), dim=-1).to(q.dtype) - out = torch.matmul(weights.unsqueeze(1), v_2d).squeeze(1) - return out.reshape(T, NH, HD) - - -def test_layer(layer_id, compress_ratio): - """Test the full attention pipeline for a specific layer.""" - torch.cuda.set_device(0) - torch.manual_seed(42) - torch.cuda.empty_cache() - - with open(os.path.join(MODEL, "model.safetensors.index.json")) as f: - wm = json.load(f)["weight_map"] - G = lambda k: P(k, wm, MODEL).to(DEV) - - p = f"model.layers.{layer_id}"; a = f"{p}.self_attn" - layer_type = "SWA" if compress_ratio <= 1 else f"CSA(c={compress_ratio})" - - print(f"\n{'='*70}") - print(f" Layer {layer_id} — {layer_type}") - print(f"{'='*70}") - - # Load weights - emb = G("model.embed_tokens.weight") - anorm = G(f"{p}.input_layernorm.weight") - qn = G(f"{a}.q_a_norm.weight"); kvn = G(f"{a}.kv_norm.weight") - woa = G(f"{a}.o_a_proj.weight") # (16384, 8192) BF16 - - qa_w = G(f"{a}.q_a_proj.weight"); qa_sf = G(f"{a}.q_a_proj.weight_scale"); qa_gs = G(f"{a}.q_a_proj.weight_scale_2") - qb_w = G(f"{a}.q_b_proj.weight"); qb_sf = G(f"{a}.q_b_proj.weight_scale"); qb_gs = G(f"{a}.q_b_proj.weight_scale_2") - kv_w = G(f"{a}.kv_proj.weight"); kv_sf = G(f"{a}.kv_proj.weight_scale"); kv_gs = G(f"{a}.kv_proj.weight_scale_2") - wob_w = G(f"{a}.o_b_proj.weight"); wob_sf = G(f"{a}.o_b_proj.weight_scale"); wob_gs = G(f"{a}.o_b_proj.weight_scale_2") - sinks = G(f"{a}.sinks") - - # BF16 references - qa_bf16 = dequant(qa_w, qa_sf, qa_gs.item()) - qb_bf16 = dequant(qb_w, qb_sf, qb_gs.item()) - kv_bf16 = dequant(kv_w, kv_sf, kv_gs.item()) - wob_bf16 = dequant(wob_w, wob_sf, wob_gs.item()) - - # CuTeDSL runners - r_qa = make_runner(qa_w, qa_sf, qa_gs, H, qa_w.shape[0]) - r_qb = make_runner(qb_w, qb_sf, qb_gs, QL, qb_w.shape[0]) - r_kv = make_runner(kv_w, kv_sf, kv_gs, H, kv_w.shape[0]) - r_wob = make_runner(wob_w, wob_sf, wob_gs, OG*OL, wob_w.shape[0]) - - # Input - NT = 6 - token_ids = torch.tensor([1, 450, 8403, 315, 5413, 374], dtype=torch.long, device=DEV) - cos_sin = build_cos_sin(max_pos=WINDOW + 256).to(DEV) - positions = torch.arange(NT, dtype=torch.int64, device=DEV) - - with torch.no_grad(): - hidden = emb[token_ids] - normed = rms(hidden, anorm, EPS) - - # ── CuTeDSL path ───────────────────────────────────────────── - qa_cute = r_qa.run(normed) - kv_cute = r_kv.run(normed) - qa_n = rms(qa_cute, qn, EPS) - kv_n = rms(kv_cute, kvn, EPS) - q_cute = r_qb.run(qa_n).view(NT, NH, HD) - q_rope = apply_gptj_rope(q_cute, positions, cos_sin, NOPE, ROPE) - - # SWA attention (for T=6, full causal within window) - o_attn = full_causal_attention(q_rope, kv_n, SCALE) - - # o_a: inverse RoPE + BMM - o_inv = apply_inv_gptj_rope(o_attn, positions, cos_sin, NOPE, ROPE) - o_grouped = o_inv.view(NT, OG, HPG * HD).permute(1, 0, 2) - woa_3d = woa.view(OG, OL, HPG * HD) - z_cute = torch.bmm(o_grouped, woa_3d.transpose(1, 2)).permute(1, 0, 2).reshape(NT, OG * OL) - - # o_b - attn_out = r_wob.run(z_cute) - - # ── BF16 reference ─────────────────────────────────────────── - qa_bf = normed @ qa_bf16.T - kv_bf = normed @ kv_bf16.T - qa_n_bf = rms(qa_bf, qn, EPS) - kv_n_bf = rms(kv_bf, kvn, EPS) - q_bf = (qa_n_bf @ qb_bf16.T).view(NT, NH, HD) - q_rope_bf = apply_gptj_rope(q_bf, positions, cos_sin, NOPE, ROPE) - o_attn_bf = full_causal_attention(q_rope_bf, kv_n_bf, SCALE) - o_inv_bf = apply_inv_gptj_rope(o_attn_bf, positions, cos_sin, NOPE, ROPE) - o_grouped_bf = o_inv_bf.view(NT, OG, HPG * HD).permute(1, 0, 2) - z_bf = torch.bmm(o_grouped_bf, woa_3d.transpose(1, 2)).permute(1, 0, 2).reshape(NT, OG * OL) - attn_bf = z_bf @ wob_bf16.T - - # ── Compare ────────────────────────────────────────────────── - c = F.cosine_similarity(attn_out.flatten().unsqueeze(0).float(), attn_bf.flatten().unsqueeze(0).float()).item() - print(f" CuTeDSL vs BF16 cosine: {c:.6f} {'✅' if c>=0.95 else '❌'}") - print(f" CuTeDSL amax: {attn_out.amax():.4f} BF16 amax: {attn_bf.amax():.4f}") - - # Full forward: attention → residual → norm → LM head - fnorm_w = G("model.norm.weight") - lm_head = G("lm_head.weight") - x = hidden + attn_out - x_normed = rms(x, fnorm_w, EPS) - logits = x_normed @ lm_head.T - top5 = torch.topk(logits[-1], 5) - log_std = logits[-1].float().std().item() - print(f" logits: amax={logits.amax():.4f} std={log_std:.4f} top5={top5.indices.tolist()}") - print(f" logit check: {'✅' if 0.5 < log_std < 50 else '❌'} (0.5 < std < 50)") - - # Cleanup - del r_qa, r_qb, r_kv, r_wob - torch.cuda.empty_cache() - return c - - -def main(): - print("=" * 70) - print(" DeepSeek-V4 CSA/HCA Attention Pipeline Test") - print(" (NOT MLA — Compressed Sparse Attention)") - print("=" * 70) - - # Test SWA layer (layer 60, compress_ratio=0) - c_swa = test_layer(60, 0) - - # Test C128A layer (layer 0, compress_ratio=128) - c_c128 = test_layer(0, 128) - - # Test C4A layer (layer 2, compress_ratio=4) - c_c4 = test_layer(2, 4) - - print(f"\n{'='*70}") - print(f" SUMMARY") - print(f" Layer 60 (SWA): {c_swa:.6f} {'✅' if c_swa>=0.95 else '❌'}") - print(f" Layer 0 (C128A/HCA): {c_c128:.6f} {'✅' if c_c128>=0.95 else '❌'}") - print(f" Layer 2 (C4A/CSA): {c_c4:.6f} {'✅' if c_c4>=0.95 else '❌'}") - print(f"{'='*70}") - - -if __name__ == "__main__": - main() diff --git a/tests/archive/test_v_mode_fix.py b/tests/archive/test_v_mode_fix.py deleted file mode 100644 index cd3f3932..00000000 --- a/tests/archive/test_v_mode_fix.py +++ /dev/null @@ -1,284 +0,0 @@ -""" -Test: V shape (HEAD_DIM, n, 1) instead of (n, HEAD_DIM, 1). -PV B operand wants (N=HEAD_DIM, K=seq_len). V must be MN-major (d, s_k). -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - -HEAD_DIM = 64 - -class FmhaV3: - def __init__(self): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0,1,2,3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.num_c_stage = 2 - self.kv_stage = 2; self.q_stage = 1; self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (128, 128, qk_ik * 4) - pv_ik = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (128, HEAD_DIM, pv_ik * (128 // pv_ik)) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0]//cute.size(qk_mma.thr_id.shape), HEAD_DIM, self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - self.q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, self.q_dtype, self.q_stage) - self.k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, self.q_dtype, self.kv_stage) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, self.kv_stage) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - qk_thr = qk_mma.get_slice(0); qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - pv_thr = pv_mma.get_slice(0); pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - self.tmem_s0_offset = 0; self.tmem_p0_offset = 32 - self.tmem_o0_offset = find_tmem_tensor_col_offset(tOtO) - tCS = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage)) - tCO = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCS, tCO], arch="sm_100") - cta = cute.size(qk_mma.thr_id.shape) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)) - self.q_tx_bytes = cute.size_in_bytes(self.q_dtype, q_s) * cta - self.kv_tx_bytes = cute.size_in_bytes(self.q_dtype, k_s) * cta - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, self.a_major, self.b_major, self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, self.qk_acc_dtype, self.cta_group, (128,HEAD_DIM), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)); v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id),q,q_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id),k,k_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape) - epi_s = cute.select(self.c_smem_s,mode=[0,1]) - tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile) - self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.v_smem_s,self.p_tmem_s,self.c_smem_s,self.epi_tile).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, q_smem_s, k_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx,_,_ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k); cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - q_bar: cute.struct.MemRange[cutlass.Int64, self.q_stage*2] - kv_bar: cute.struct.MemRange[cutlass.Int64, self.kv_stage*2] - s_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage*2] - tmem_dealloc: cutlass.Int64; holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - qp,qc = pipeline.PipelineTmaUmma.create(barrier_storage=st.q_bar.data_ptr(),num_stages=self.q_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.q_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - kvp,kvc = pipeline.PipelineTmaUmma.create(barrier_storage=st.kv_bar.data_ptr(),num_stages=self.kv_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.kv_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - s_prod,s_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.s_bar.data_ptr(),num_stages=1,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.epilogue_warp_id))).make_participants() - softmax_done_bar = pipeline.NamedBarrier(barrier_id=3, num_threads=32 + 32*len(self.epilogue_warp_id)) - acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(),num_stages=self.num_acc_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,len(self.epilogue_warp_id)),cta_layout_vmnk=cl_vmnk,defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2,num_threads=32*len((self.mma_warp_id,*self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr,barrier_for_retrieve=tmem_bar,allocator_warp_id=self.epilogue_warp_id[0],is_two_cta=cute.size(qk_mma.thr_id.shape)==2,two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk,is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype,layout=q_smem_s.outer,byte_alignment=128,swizzle=q_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype,layout=k_smem_s.outer,byte_alignment=128,swizzle=k_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype,layout=v_smem_s.outer,byte_alignment=128,swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype,layout=c_smem_s.outer,byte_alignment=128,swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None)) - gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None)) - gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None)) - gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None)) - n_kv_tiles = cute.size(gK, mode=[3]) - - qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - tCgV = pv_thr.partition_B(gV); tCgC = pv_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,0,None,0)).shape) - tAsQ,tAgQ = cpasync.tma_partition(tma_q,0,a_lay,cute.group_modes(sQ,0,3),cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape) - tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3)) - tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)]; tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - # PV read view (MMA only) - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None,None,None,0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # TMA LOAD - if warp_idx == self.tma_warp_id: - qp.reset(); qh = qp.acquire_and_advance() - cute.copy(tma_q,tAgQ[(None,qh.count)],tAsQ[(None,qh.index)],tma_bar_ptr=qh.barrier) - qp.tail() - kvp.reset(); pk = kvp.try_acquire() - for kt in cutlass.range(n_kv_tiles,unroll=1): - kh = kvp.acquire_and_advance(pk) - cute.copy(tma_k,tBgK[(None,kh.count)],tBsK[(None,kh.index)],tma_bar_ptr=kh.barrier) - pk = cutlass.Boolean(1) - vh = kvp.acquire_and_advance(pk) - cute.copy(tma_v,tVgV[(None,vh.count)],tVsV[(None,vh.index)],tma_bar_ptr=vh.barrier) - pk = cutlass.Boolean(1) - kvp.tail() - - # MMA - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - qc.reset(); qh = qc.wait_and_advance(); qh.release() - kvc.reset(); pk = kvc.try_wait() - acc_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_stage) - acc_pipe.producer_acquire(acc_st) - for kt in range(n_kv_tiles): - kh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - sh = s_prod.acquire_and_advance() - qk_mma.set(tcgen05.Field.ACCUMULATE, False) - for kb in cutlass.range(cute.size(tCrQ,mode=[2]), unroll_full=True): - cute.gemm(qk_mma, tStS0, tCrQ[(None,None,kb,0)], tCrK[(None,None,kb,kh.index)], tStS0) - qk_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - sh.commit(); kh.release() - softmax_done_bar.arrive_and_wait() - vh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - pv_mma.set(tcgen05.Field.ACCUMULATE, kt != 0) - for kb in cutlass.range(cute.size(tOrP0,mode=[2]), unroll_full=True): - cute.gemm(pv_mma, tOtO0, tOrP0[(None,None,kb)], tCrV[(None,None,kb,vh.index)], tOtO0) - pv_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - vh.release() - acc_pipe.producer_commit(acc_st); acc_st.advance() - acc_pipe.producer_tail(acc_st) - - # EPILOGUE - if warp_idx < self.mma_warp_id: - tmem.allocate(self.num_tmem_alloc_cols) - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) - sfw_idx = tidx % (32 * len(self.epilogue_warp_id)) - - # S load - tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS0) - thr_load = tiled_tmem_load.get_slice(sfw_idx) - tTMEM_LOADtS = thr_load.partition_S(tStS0) - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tTMEM_LOADcS = thr_load.partition_D(tScS) - - # P store — QK C-fragment composition (FMHA pattern) - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - tStP_layout = cute.composition( - tStS.layout, - cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32)), - ) - tStP0 = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStP_layout) - tmem_store_atom = cute.make_copy_atom( - tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), - self.qk_acc_dtype, - ) - tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStP0) - thr_store = tiled_tmem_store.get_slice(sfw_idx) - tTMEM_STOREtP = thr_store.partition_D(tStP0) - tScP_layout = cute.composition( - tScS.layout, - cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32)), - ) - tScP = cute.make_tensor(tScS.iterator, tScP_layout) - tTMEM_STOREcP = thr_store.partition_S(tScP) - - for kt in range(n_kv_tiles): - si_handle = s_cons.wait_and_advance() - tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS) - # Register bridge - rP_words = cute.make_rmem_tensor(tTMEM_STOREcP.shape, self.qk_acc_dtype) - rP_bf16 = cute.make_tensor( - cute.recast_ptr(rP_words.iterator, dtype=self.q_dtype), - tTMEM_LOADrS.layout, - ) - frg_cnt = 4 - frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt - tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile)) - rP_bf16_frg = cute.logical_divide(rP_bf16, cute.make_layout(frg_tile)) - for j in range(frg_cnt): - s_vec = tTMEM_LOADrS_frg[None, j].load() - rP_bf16_frg[None, j].store(s_vec.to(self.q_dtype)) - cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP) - cute.arch.fence_view_async_tmem_store() - si_handle.release() - softmax_done_bar.arrive() - - tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout) - acc_cons_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage) - c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) - c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp) - acc_cons_st = utils.gemm.sm100.epilogue_tma_store(self, tidx, warp_idx, tma_c, tCtO_base, sC, tCgC, epi_tile, 0, const_expr(lambda x: x), (0,0,0), acc_cons_st, acc_pipe, c_pipe) - c_pipe.producer_tail() - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def test(): - torch.manual_seed(42) - m, n, hd = 128, 128, HEAD_DIM - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(n, hd, 1, dtype=torch.bfloat16, device='cuda') - - # FIX: V is (HEAD_DIM, n, 1) — PV B operand wants (d, s_k) = (HEAD_DIM, seq_len) - v_base = torch.ones(n, hd, dtype=torch.bfloat16, device='cuda') - v_for_kernel = v_base.T.unsqueeze(-1) # shape (hd, n, 1), stride (1, hd) - print(f"V shape: {v_for_kernel.shape} stride: {v_for_kernel.stride()}", flush=True) - - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - qf = q[:,:,0].float(); kf = k[:,:,0].float() - ref = (qf @ kf.T).bfloat16().float() @ v_base.float() - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v_for_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_for_kernel)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - kernel = FmhaV3() - print('Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - print(f'pv_mma_tiler={kernel.pv_mma_tiler} p0={kernel.tmem_p0_offset} o0={kernel.tmem_o0_offset} alloc={kernel.num_tmem_alloc_cols}', flush=True) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - out = c[:,:,0].float() - cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - print(f'V=(hd,n) cosine {cos:.6f} {"PASS" if cos >= 0.99 else "FAIL"}') - if cos < 0.99: - print(f' out[0,:4]={out[0,:4].tolist()} ref[0,:4]={ref[0,:4].tolist()}') - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_vllm_codepaths_b200.py b/tests/archive/test_vllm_codepaths_b200.py deleted file mode 100644 index 505a193e..00000000 --- a/tests/archive/test_vllm_codepaths_b200.py +++ /dev/null @@ -1,252 +0,0 @@ -#!/usr/bin/env python3 -""" -Test the EXACT code paths used in vLLM's Blackwell attention. - -Imports the actual functions from csa_attention.py and blackwell_attention.py -and verifies they produce correct output with real weights. - -This is the closest possible test to what runs in the container. -""" -import sys, os, json, torch, torch.nn.functional as F -from safetensors import safe_open - -REPO = "/root/nvfp4-megamoe-kernel" -sys.path.insert(0, REPO) -MODEL = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4" -DEV = "cuda:0" - -H = 7168; NH = 128; HD = 512; NOPE = 448; ROPE = 64 -QL = 1536; OL = 1024; OG = 16; HPG = NH // OG -EPS = 1e-6; WINDOW = 128; SCALE = HD ** -0.5 - -E2M1 = torch.tensor([0,.5,1.,1.5,2.,3.,4.,6.,-0,-.5,-1.,-1.5,-2.,-3.,-4.,-6.], dtype=torch.float32) - -_cache = {} -def P(k, wm, md): - if k in _cache: return _cache[k] - with safe_open(os.path.join(md, wm[k]), framework="pt") as f: - t = f.get_tensor(k) - _cache[k] = t - return t - -def rms(x, w, eps=1e-6): - v = x.float().pow(2).mean(-1, keepdim=True) - return (w.float() * (x * torch.rsqrt(v+eps)).float()).to(x.dtype) - -def make_runner(w, sf, gs_t, inf, outf): - from dsv4.layers.linear import Nvfp4Linear - fp4 = w.view(torch.float4_e2m1fn_x2).permute(1,0).contiguous() - s = sf.to(torch.float8_e4m3fn) if sf.dtype != torch.float8_e4m3fn else sf - s = s.permute(1,0).contiguous() - gs = gs_t.max().item() if gs_t.numel() > 1 else gs_t.item() - r = Nvfp4Linear(in_features=inf, out_features=outf, max_num_tokens=8192, device=str(w.device)) - r.fp4 = [fp4]; r.sf = [s]; r.gs = [gs] - r.finalize_weights(); r._ensure_initialized() - return r - -def build_cos_sin(max_pos=4096, rope_dim=ROPE): - half = rope_dim // 2 - inv_freq = 1.0 / (10000.0 ** (torch.arange(0, half, dtype=torch.float32) / half)) - freqs = torch.outer(torch.arange(max_pos, dtype=torch.float32), inv_freq) - return torch.cat([freqs.cos(), freqs.sin()], dim=-1) - -def apply_gptj_rope(x, positions, cos_sin, nope_dim, rope_dim): - if rope_dim == 0 or x.numel() == 0: return x - half = rope_dim // 2 - cos = cos_sin[positions, :half].to(x.dtype) - sin = cos_sin[positions, half:2*half].to(x.dtype) - if x.dim() == 3: cos = cos.unsqueeze(1); sin = sin.unsqueeze(1) - x_rope = x[..., nope_dim:].clone() - even = x_rope[..., 0::2]; odd = x_rope[..., 1::2] - out = x.clone() - out[..., nope_dim:][..., 0::2] = even * cos - odd * sin - out[..., nope_dim:][..., 1::2] = even * sin + odd * cos - return out - -def apply_inv_gptj_rope(x, positions, cos_sin, nope_dim, rope_dim): - if rope_dim == 0 or x.numel() == 0: return x - half = rope_dim // 2 - cos = cos_sin[positions, :half].to(x.dtype) - sin = cos_sin[positions, half:2*half].to(x.dtype) - if x.dim() == 3: cos = cos.unsqueeze(1); sin = sin.unsqueeze(1) - x_rope = x[..., nope_dim:].clone() - even = x_rope[..., 0::2]; odd = x_rope[..., 1::2] - out = x.clone() - out[..., nope_dim:][..., 0::2] = even * cos + odd * sin - out[..., nope_dim:][..., 1::2] = -even * sin + odd * cos - return out - -def kv_quantize_fp8(kv_bf16): - amax = kv_bf16.float().abs().amax(dim=-1, keepdim=True).clamp(min=1e-12) - fp8_max = torch.tensor(448.0, dtype=torch.float32, device=kv_bf16.device) - scale = fp8_max / amax - kv_fp8 = (kv_bf16.float() * scale).to(torch.float8_e4m3fn) - inv_scale = (amax / fp8_max).to(torch.bfloat16) - return kv_fp8, inv_scale - -def kv_dequantize_fp8(kv_fp8, inv_scale): - return (kv_fp8.to(torch.bfloat16) * inv_scale).to(torch.bfloat16) - -def causal_prefill_attention(q, kv, scale): - T, NH, HD = q.shape - q_t = q.permute(1, 0, 2) - kv_exp = kv.unsqueeze(0).expand(NH, -1, -1) - out = F.scaled_dot_product_attention(q_t, kv_exp, kv_exp, is_causal=True, scale=scale) - return out.permute(1, 0, 2) - - -def main(): - """Test the exact csa_attention.py code paths used in the container.""" - from dsv4.reference.attention import ( - apply_gptj_rope, - apply_inv_gptj_rope, - ) - # Import the vLLM patch version (the actual code used in the container) - sys.path.insert(0, os.path.join(REPO, "vllm", "patches", "layers")) - from csa_attention import ( - fused_qnorm_rope_kv_insert_py, - blackwell_attention_kv_write as vllm_kv_write, - blackwell_attention_decode as vllm_decode, - kv_quantize_fp8 as vllm_kv_quantize, - kv_dequantize_fp8 as vllm_kv_dequantize, - causal_prefill_attention, - ) - - torch.cuda.set_device(0) - - with open(os.path.join(MODEL, "model.safetensors.index.json")) as f: - wm = json.load(f)["weight_map"] - G = lambda k: P(k, wm, MODEL).to(DEV) - - # Test layer 60 (SWA) - layer_id = 60 - p = f"model.layers.{layer_id}"; a = f"{p}.self_attn" - - emb = G("model.embed_tokens.weight") - anorm = G(f"{p}.input_layernorm.weight") - qn = G(f"{a}.q_a_norm.weight"); kvn = G(f"{a}.kv_norm.weight") - woa = G(f"{a}.o_a_proj.weight") - qa_w = G(f"{a}.q_a_proj.weight"); qa_sf = G(f"{a}.q_a_proj.weight_scale"); qa_gs = G(f"{a}.q_a_proj.weight_scale_2") - qb_w = G(f"{a}.q_b_proj.weight"); qb_sf = G(f"{a}.q_b_proj.weight_scale"); qb_gs = G(f"{a}.q_b_proj.weight_scale_2") - kv_w = G(f"{a}.kv_proj.weight"); kv_sf = G(f"{a}.kv_proj.weight_scale"); kv_gs = G(f"{a}.kv_proj.weight_scale_2") - wob_w = G(f"{a}.o_b_proj.weight"); wob_sf = G(f"{a}.o_b_proj.weight_scale"); wob_gs = G(f"{a}.o_b_proj.weight_scale_2") - - r_qa = make_runner(qa_w, qa_sf, qa_gs, H, qa_w.shape[0]) - r_qb = make_runner(qb_w, qb_sf, qb_gs, QL, qb_w.shape[0]) - r_kv = make_runner(kv_w, kv_sf, kv_gs, H, kv_w.shape[0]) - r_wob = make_runner(wob_w, wob_sf, wob_gs, OG*OL, wob_w.shape[0]) - cos_sin = build_cos_sin(max_pos=4096).to(DEV) - woa_3d = woa.view(OG, OL, HPG * HD) - - N = 8 - token_ids = torch.tensor([1, 450, 8403, 315, 5413, 374, 2198, 643], dtype=torch.long, device=DEV) - - with torch.no_grad(): - # ── Test 1: Verify fused_qnorm_rope_kv_insert_py ────────── - print("=== Test 1: fused_qnorm_rope_kv_insert_py ===") - positions_p = torch.arange(N, dtype=torch.int64, device=DEV) - hidden_p = emb[token_ids] - normed_p = rms(hidden_p, anorm, EPS) - qa_p = r_qa.run(normed_p) - kv_p = r_kv.run(normed_p) - - # Manual Q norm + RoPE (reference) - qa_n_ref = rms(qa_p, qn, EPS) - q_ref = r_qb.run(qa_n_ref).view(N, NH, HD) - q_rope_ref = apply_gptj_rope(q_ref, positions_p, cos_sin, NOPE, ROPE) - - # Using fused_qnorm_rope_kv_insert_py - q_test = r_qb.run(qa_n_ref).view(N, NH, HD) - fused_qnorm_rope_kv_insert_py( - q_test, kv_p, None, None, positions_p, - cos_sin, EPS, 64, # block_size - nope_dim=NOPE, rope_dim=ROPE, - ) - - c = F.cosine_similarity(q_rope_ref.flatten().unsqueeze(0).float(), q_test.flatten().unsqueeze(0).float()).item() - print(f" fused_qnorm_rope vs manual: cosine = {c:.6f} str('PASS' if c>=0.999 else 'FAIL')") - - # ── Test 2: Verify blackwell_attention_kv_write ─────────── - print("\n=== Test 2: blackwell_attention_kv_write ===") - block_size = 64; max_tokens = 256 - num_blocks = (max_tokens + block_size - 1) // block_size - - # uint8 cache (like vLLM uses) - swa_cache = torch.zeros(num_blocks, block_size, HD, dtype=torch.uint8, device=DEV) - inv_scale_cache = torch.zeros(max_tokens, 1, dtype=torch.bfloat16, device=DEV) - slot_mapping = positions_p # Simple: slot = position - - # Manual KV RoPE + fp8 quant - kv_n = rms(kv_p, kvn, EPS) - kv_rope_manual = apply_gptj_rope(kv_n.unsqueeze(1), positions_p, cos_sin, NOPE, ROPE).squeeze(1) - kv_fp8_manual, inv_s_manual = kv_quantize_fp8(kv_rope_manual) - - # Write using vLLM's function - vllm_kv_write( - kv_n, positions_p, swa_cache, inv_scale_cache, - slot_mapping, block_size, cos_sin, - nope_dim=NOPE, rope_dim=ROPE, - ) - - # Read back and compare - bi = slot_mapping // block_size; oi = slot_mapping % block_size - kv_read = swa_cache[bi, oi].view(torch.float8_e4m3fn) - inv_read = inv_scale_cache[slot_mapping] - kv_dequant = kv_dequantize_fp8(kv_read, inv_read) - - c = F.cosine_similarity(kv_rope_manual.flatten().unsqueeze(0).float(), kv_dequant.flatten().unsqueeze(0).float()).item() - print(f" vllm_kv_write roundtrip: cosine = {c:.6f} str('PASS' if c>=0.99 else 'FAIL')") - - # ── Test 3: Decode attention using swa_indices ──────────── - print("\n=== Test 3: Decode attention with swa_indices ===") - decode_id = torch.tensor([991], dtype=torch.long, device=DEV) - pos_d = torch.tensor([N], dtype=torch.int64, device=DEV) - - # Write decode KV to cache - hidden_d = emb[decode_id] - normed_d = rms(hidden_d, anorm, EPS) - qa_d = r_qa.run(normed_d); kv_d = r_kv.run(normed_d) - qa_n_d = rms(qa_d, qn, EPS); kv_n_d = rms(kv_d, kvn, EPS) - q_d = r_qb.run(qa_n_d).view(1, NH, HD) - q_rope_d = apply_gptj_rope(q_d, pos_d, cos_sin, NOPE, ROPE) - - vllm_kv_write(kv_n_d, pos_d, swa_cache, inv_scale_cache, - pos_d, block_size, cos_sin, nope_dim=NOPE, rope_dim=ROPE) - - # swa_indices: simulate vLLM's pre-computed indices - # These are flat slot indices for each decode token's window - all_slots = torch.arange(N + 1, dtype=torch.int64, device=DEV) - swa_indices = all_slots.unsqueeze(0) # (1, N+1) — all tokens in window - swa_lens = torch.tensor([N + 1], dtype=torch.int64, device=DEV) - - o_decode = vllm_decode( - q_rope_d, pos_d, swa_cache, inv_scale_cache, - pos_d, block_size, SCALE, WINDOW, - swa_indices=swa_indices, - swa_lens=swa_lens, - decode_token_idx=0, - ) - print(f" Decode output: amax={o_decode.amax():.4f} NaN={torch.isnan(o_decode).any()}") - - # ── Reference: full prefill attention ──────────────────── - all_ids = torch.cat([token_ids, decode_id]) - all_pos = torch.arange(N + 1, dtype=torch.int64, device=DEV) - hidden_ref = emb[all_ids] - normed_ref = rms(hidden_ref, anorm, EPS) - qa_ref = r_qa.run(normed_ref); kv_ref = r_kv.run(normed_ref) - qa_n_ref = rms(qa_ref, qn, EPS); kv_n_ref = rms(kv_ref, kvn, EPS) - q_ref = r_qb.run(qa_n_ref).view(N + 1, NH, HD) - q_rope_ref = apply_gptj_rope(q_ref, all_pos, cos_sin, NOPE, ROPE) - kv_rope_ref = apply_gptj_rope(kv_n_ref.unsqueeze(1), all_pos, cos_sin, NOPE, ROPE).squeeze(1) - o_ref = causal_prefill_attention(q_rope_ref, kv_rope_ref, SCALE) - o_ref_decode = o_ref[-1:] - - c = F.cosine_similarity(o_decode.flatten().unsqueeze(0).float(), o_ref_decode.flatten().unsqueeze(0).float()).item() - status = "PASS" if c >= 0.98 else "FAIL" - print(f" Decode vs reference cosine: {c:.6f} {status}") - - print("\n=== DONE ===") - - -if __name__ == "__main__": - main() diff --git a/tests/archive/test_vsmem_diag.py b/tests/archive/test_vsmem_diag.py deleted file mode 100644 index 11c3ef95..00000000 --- a/tests/archive/test_vsmem_diag.py +++ /dev/null @@ -1,385 +0,0 @@ -""" -Minimal PV-only test: Load P from GMEM to TMEM via QK-style MMA, then PV from TMEM. -Step 1: QK MMA writes FP32 S to TMEM (we know this works) -Step 2: Softmax packing writes BF16 P to TMEM (test this) -Step 3: PV MMA reads BF16 P from TMEM and V from SMEM, produces O - -But to isolate the bug, let me test just the PV MMA in isolation. -I'll write known BF16 values to TMEM using the softmax packing path, -then immediately read them back using the PV A-fragment path, -and compare. - -Actually, the simplest isolation test: -1. Do QK MMA to get S in TMEM (cosine 0.999999 verified) -2. Do softmax packing: S → P in TMEM (at offset 32) -3. Skip PV entirely — read P from TMEM using the C-fragment composition LOAD path -4. Output P to GMEM and compare against S.to(BF16) - -This tests whether the softmax packing writes P correctly to the same TMEM -that the PV would read from. - -But we can't easily read P from TMEM using the standard epilogue path -because the epilogue expects FP32 accumulator data. - -Alternative: Use the PV MMA with V=I (identity). If P is correct, -then P @ I = P. But V needs to be MN-major and (128, 128), not (128, 64). -The output would be (128, 128) which doesn't match our (128, 64) c tensor. - -Let me use V that selects the first 64 columns: V[k, n] = delta(k, n) for k in [0,63]. -This gives P @ V = P[:, :64], and the output is (128, 64). -But V is (128, 128) in the MMA K,N dims. V[k, n] for k in [0,127], n in [0,63]. -Hmm, this is getting complicated. Let me just do the identity approach with a (128, 128) output. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class VSmemDiag: - """QK + softmax packing + PV with V=I to isolate PV MMA correctness. - Output should be P = S.to(BF16), i.e. (Q@K^T).bfloat16() - With V=I, O = P @ I = P. - But V is (K=128, N=128) in the MMA. We need a 128x128 identity in MN-major. - Output tensor is (128, 128). - """ - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False # needed by epilogue_tma_store - self.epilog_sync_bar_id = 1 # needed by epilogue_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - # PV with V=I: output is (128, 128), same as QK - self.pv_mma_tiler = (self.qk_mma_tiler[0], self.qk_mma_tiler[1], self.qk_mma_tiler[1]) - # pv_mma_tiler = (128, 128, 128) since V is 128x128 - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - print(int(cute.size(v_smem)), int(cute.size(v_smem, mode=[0])), int(cute.size(v_smem, mode=[1]))) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - # ⛔⛔⛔ CRITICAL: num_tma_load_bytes MUST include ALL TMA-loaded tensors (Q + K + V). Missing V → DEADLOCK. See FOOTGUN #0 in README. - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - # PV with 128x128 output (V=I) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - - -if __name__ == '__main__': - test() diff --git a/tests/archive/test_warmup_gs.py b/tests/archive/test_warmup_gs.py deleted file mode 100644 index a4517d53..00000000 --- a/tests/archive/test_warmup_gs.py +++ /dev/null @@ -1,206 +0,0 @@ -#!/usr/bin/env python3 -""" -Test C: Warmup-based gs computation — verify that exact warmup gs values -produce good cosine when used with quantize_activation_nvfp4. - -The warmup runs quantize_to_nvfp4 (dynamic gs) on representative input, -captures the exact gs for both L1 and L2, then feeds those values to -quantize_activation_nvfp4 (fixed gs, cudagraph-safe). - -Usage (on B200): - source /root/nvfp4-megamoe-kernel/tests/.venv/bin/activate - python3 tests/test_warmup_gs.py -""" -import torch, sys, os, json -sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) - -from dsv4.ops.quantize import ( - quantize_to_nvfp4, - quantize_activation_nvfp4, -) -from dsv4.ops.layouts import ( - make_b_k_major, - assemble_scales_2d_side, - assemble_scales_3d_side, - compute_expert_offsets, -) -from dsv4.ops.gemm_runner import ( - run_nvfp4_grouped_gemm, -) -from dsv4.reference.moe_pipeline import run_nvfp4_moe -from vllm.nvfp4_cutedsl import Nvfp4MoE -from safetensors import safe_open - -MODEL_DIR = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4" -DEVICE = "cuda" -E2M1_LUT = torch.tensor([0.,0.5,1.,1.5,2.,3.,4.,6.,-0.,-.5,-1.,-1.5,-2.,-3.,-4.,-6.], dtype=torch.float32) - - -def dequant(w, sf, gs): - dev = w.device - lo = E2M1_LUT.to(dev)[(w & 0xF).long()] - up = E2M1_LUT.to(dev)[((w >> 4) & 0xF).long()] - o = torch.empty(w.shape[0], w.shape[1]*2, dtype=torch.float32, device=dev) - o[:, 0::2] = lo; o[:, 1::2] = up - return (o * sf.float().repeat_interleave(16, dim=1)[:, :o.shape[1]] * gs).to(torch.bfloat16) - - -def load_tensor(key): - with open(os.path.join(MODEL_DIR, "model.safetensors.index.json")) as f: - wm = json.load(f)["weight_map"] - shard = os.path.join(MODEL_DIR, wm.get(key, "")) - if not os.path.exists(shard): return None - with safe_open(shard, framework="pt") as f: - if key in f.keys(): return f.get_tensor(key).to(DEVICE) - return None - - -def load_layer0_experts(expert_indices): - l1_fp4, l1_sf, l1_gs = [], [], [] - l2_fp4, l2_sf, l2_gs = [], [], [] - for e in expert_indices: - gw = load_tensor(f"model.layers.0.mlp.experts.{e}.gate_proj.weight") - uw = load_tensor(f"model.layers.0.mlp.experts.{e}.up_proj.weight") - gsf = load_tensor(f"model.layers.0.mlp.experts.{e}.gate_proj.weight_scale") - usf = load_tensor(f"model.layers.0.mlp.experts.{e}.up_proj.weight_scale") - ggs = load_tensor(f"model.layers.0.mlp.experts.{e}.gate_proj.weight_scale_2").item() - ugs = load_tensor(f"model.layers.0.mlp.experts.{e}.up_proj.weight_scale_2").item() - fw = torch.cat([gw, uw], dim=0) - fw4 = fw.view(torch.float4_e2m1fn_x2).permute(1, 0).contiguous() - fs = torch.cat([gsf, usf], dim=0).permute(1, 0).contiguous() - mgs = max(ggs, ugs) - if ggs != ugs: - f32 = fs.float() - f32[:, :3072] *= (ggs / mgs) - f32[:, 3072:] *= (ugs / mgs) - fs = f32.to(torch.float8_e4m3fn) - l1_fp4.append(fw4); l1_sf.append(fs); l1_gs.append(mgs) - dw = load_tensor(f"model.layers.0.mlp.experts.{e}.down_proj.weight") - dsf = load_tensor(f"model.layers.0.mlp.experts.{e}.down_proj.weight_scale") - dgs = load_tensor(f"model.layers.0.mlp.experts.{e}.down_proj.weight_scale_2").item() - l2_fp4.append(dw.view(torch.float4_e2m1fn_x2).permute(1, 0).contiguous()) - l2_sf.append(dsf.permute(1, 0).contiguous()); l2_gs.append(dgs) - return l1_fp4, l1_sf, l1_gs, l2_fp4, l2_sf, l2_gs - - -def warmup_compute_gs(runner, hidden_states, topk_weights, topk_ids): - """Run a full forward pass with quantize_to_nvfp4 (dynamic gs) - to capture the exact gs values for L1 and L2.""" - device = hidden_states.device - num_tokens = hidden_states.shape[0] - top_k = topk_ids.shape[1] - - # Build slot mapping (same as runner.run()) - runner._ensure_stacked() - flat_ids = topk_ids.reshape(-1) - num_slots = num_tokens * top_k - token_indices = runner._token_indices[:num_slots] - sort_idx = flat_ids.argsort(stable=True) - sorted_ids = flat_ids[sort_idx] - sorted_token_ids = token_indices[sort_idx] - slot_hidden = hidden_states[sorted_token_ids] - - # L1: dynamic gs - _, _, l1_gs = quantize_to_nvfp4(slot_hidden) - - # Run L1 GEMM with dynamic gs to get L1 output - x_fp4, x_sf = quantize_activation_nvfp4(slot_hidden, l1_gs) - - expert_id_range = runner._expert_id_range - tokens_per_expert = (sorted_ids.unsqueeze(1) == expert_id_range.unsqueeze(0)).sum(dim=0).int() - expert_offsets = runner._expert_offsets_buf - expert_offsets.zero_() - expert_offsets[1:runner.num_experts + 1] = tokens_per_expert.cumsum(0) - - l1_scale_a = runner._assemble_scales_cudagraph_safe( - x_sf, expert_offsets[:runner.num_experts + 1], - runner._padded_x_sf_buf_l1, runner._per_expert_scale_bufs_l1 - ) - l1_gsa = torch.full((runner.num_experts,), l1_gs, dtype=torch.float32, device=device) - - l1_out = run_nvfp4_grouped_gemm( - mat_a=x_fp4, mat_b=runner._l1_mat_b, - scale_a=l1_scale_a, scale_b=runner._l1_scale_b, - expert_offsets=expert_offsets[1:runner.num_experts + 1], - global_scale_a=l1_gsa, global_scale_b=runner._l1_gsb, - ) - - # L2: compute gs from actual L1 output - gate = l1_out[:, :runner.intermediate_size] - up = l1_out[:, runner.intermediate_size:] - activated = torch.nn.functional.silu(gate) * up - _, _, l2_gs = quantize_to_nvfp4(activated) - - return l1_gs, l2_gs - - -def main(): - expert_indices = [0, 1, 2] - num_experts = len(expert_indices) - hidden_size = 7168 - intermediate_size = 3072 - - print("Loading weights...") - l1_fp4, l1_sf, l1_gs, l2_fp4, l2_sf, l2_gs = load_layer0_experts(expert_indices) - - torch.manual_seed(42) - hidden_states = torch.randn(4, hidden_size, dtype=torch.bfloat16, device=DEVICE) * 2.0 - topk_ids = torch.tensor([[0, 1]] * 4, dtype=torch.int32, device=DEVICE) - topk_weights = torch.tensor([[0.6, 0.4]] * 4, dtype=torch.float32, device=DEVICE) - - # Pipeline reference - weights = {'l1_fp4': l1_fp4, 'l1_sf': l1_sf, 'l1_gs': l1_gs, - 'l2_fp4': l2_fp4, 'l2_sf': l2_sf, 'l2_gs': l2_gs} - ref = run_nvfp4_moe(hidden_states.clone(), topk_ids.clone(), topk_weights.clone(), weights, expert_indices) - print(f"Pipeline: amax={ref.abs().max():.4f}, mean={ref.float().mean():.6f}") - - # ── Test 1: Runner with warmup gs (no safety margin) ── - print("\n--- Test 1: Warmup gs, no safety margin ---") - runner = Nvfp4MoE(num_experts, hidden_size, intermediate_size, device=DEVICE) - runner.prepare_weights_direct( - [w.clone() for w in l1_fp4], [w.clone() for w in l1_sf], list(l1_gs), - [w.clone() for w in l2_fp4], [w.clone() for w in l2_sf], list(l2_gs), - ) - - # Use the runner's built-in warmup method - runner.compute_activation_global_scales(hidden_states.clone(), topk_weights, topk_ids) - - result = runner.run(hidden_states.clone(), topk_weights, topk_ids) - - cos = torch.nn.functional.cosine_similarity( - result.flatten().unsqueeze(0).float(), ref.flatten().unsqueeze(0).float() - ).item() - print(f" Cosine: {cos:.6f}, amax={result.abs().max():.4f}") - - # ── Test 2: Runner with warmup gs + safety margins ── - for safety in [1.0, 1.1, 1.2, 1.5, 2.0]: - runner2 = Nvfp4MoE(num_experts, hidden_size, intermediate_size, device=DEVICE) - runner2.prepare_weights_direct( - [w.clone() for w in l1_fp4], [w.clone() for w in l1_sf], list(l1_gs), - [w.clone() for w in l2_fp4], [w.clone() for w in l2_sf], list(l2_gs), - ) - runner2._l1_activation_global_scale = l1_gs_val * safety - runner2._l2_activation_global_scale = l2_gs_val * safety - result2 = runner2.run(hidden_states.clone(), topk_weights, topk_ids) - cos2 = torch.nn.functional.cosine_similarity( - result2.flatten().unsqueeze(0).float(), ref.flatten().unsqueeze(0).float() - ).item() - print(f" Safety {safety:.1f}x: cosine={cos2:.6f}, amax={result2.abs().max():.4f}") - - # ── Test 3: Different input (verify warmup gs generalizes) ── - print("\n--- Test 3: Different input with same warmup gs ---") - torch.manual_seed(99) - hidden_states2 = torch.randn(4, hidden_size, dtype=torch.bfloat16, device=DEVICE) * 2.0 - topk_ids2 = torch.tensor([[0, 1]] * 4, dtype=torch.int32, device=DEVICE) - topk_weights2 = torch.tensor([[0.6, 0.4]] * 4, dtype=torch.float32, device=DEVICE) - - ref2 = run_nvfp4_moe(hidden_states2.clone(), topk_ids2.clone(), topk_weights2.clone(), weights, expert_indices) - result3 = runner.run(hidden_states2.clone(), topk_weights2, topk_ids2) - cos3 = torch.nn.functional.cosine_similarity( - result3.flatten().unsqueeze(0).float(), ref2.flatten().unsqueeze(0).float() - ).item() - print(f" Different input: cosine={cos3:.6f}") - - -if __name__ == "__main__": - main() diff --git a/tests/archive/test_wo_a.py b/tests/archive/test_wo_a.py deleted file mode 100644 index 71f086d9..00000000 --- a/tests/archive/test_wo_a.py +++ /dev/null @@ -1,171 +0,0 @@ -"""Unit test: wo_a NVFP4 grouped linear + inverse RoPE. - -Tests the CuTeDSL NVFP4 grouped GEMM that replaces DeepGEMM's fp8_einsum -for the wo_a (o-projection first half) in DeepSeek V4 attention. - -Also tests inverse_rope_bf16 against a synthetic reference. - -Usage (B200): python3 tests/test_wo_a.py - -Requires: CuTeDSL, CUDA, Blackwell GPU -""" - -import sys -import os -import torch -import torch.nn.functional as F - -# Add repo root to path -sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) - -from dsv4.ops.rope import inverse_rope_bf16 -from dsv4.layers.grouped_linear import Nvfp4GroupedLinear - -DEVICE = "cuda:0" - -# DeepSeek V4 Pro dimensions -N_LOCAL_GROUPS = 8 -HEADS_PER_GROUP = 16 # 128 heads / 8 groups -HEAD_DIM = 512 -NOPE_DIM = 448 -ROPE_DIM = 64 -O_LORA_RANK = 1536 -GROUP_IN = HEADS_PER_GROUP * HEAD_DIM # 8192 -NUM_TOKENS = 4 - - -def test_inverse_rope(): - """Test inverse_rope_bf16: apply RoPE then inverse → should recover original.""" - print("\n=== Test: inverse_rope_bf16 ===") - - torch.manual_seed(42) - num_tokens = 4 - num_heads = N_LOCAL_GROUPS * HEADS_PER_GROUP - max_pos = 128 - - # Build cos_sin_cache (same format as vLLM: cos||sin concatenated) - rope_dim = ROPE_DIM - half_rope = rope_dim // 2 - base = 10000.0 - inv_freq = 1.0 / (base ** (torch.arange(0, half_rope, dtype=torch.float32, device=DEVICE) / half_rope)) - - pos = torch.arange(max_pos, dtype=torch.float32, device=DEVICE) - freqs = torch.outer(pos, inv_freq) # (max_pos, half_rope) - cos_sin_cache = torch.cat([freqs.cos(), freqs.sin(), ], dim=-1) # (max_pos, rope_dim) - - # Random attention output - o = torch.randn(num_tokens, num_heads, HEAD_DIM, dtype=torch.bfloat16, device=DEVICE) * 2.0 - positions = torch.randint(0, max_pos, (num_tokens,), dtype=torch.int64, device=DEVICE) - - # Apply RoPE (forward), then inverse - # Forward RoPE (GPT-J interleaved): - o_rope = o[:, :, NOPE_DIM:].clone() - cos_all = cos_sin_cache[positions, :half_rope].unsqueeze(1).to(o.dtype) - sin_all = cos_sin_cache[positions, half_rope:].unsqueeze(1).to(o.dtype) - o_even = o_rope[:, :, 0::2] - o_odd = o_rope[:, :, 1::2] - rope_even = o_even * cos_all - o_odd * sin_all - rope_odd = o_even * sin_all + o_odd * cos_all - o_fwd = o.clone() - o_fwd[:, :, NOPE_DIM:][:, :, 0::2] = rope_even - o_fwd[:, :, NOPE_DIM:][:, :, 1::2] = rope_odd - - # Apply inverse RoPE - o_inv = inverse_rope_bf16(o_fwd, positions, cos_sin_cache, NOPE_DIM, ROPE_DIM) - - # Compare with original - cos = F.cosine_similarity( - o.flatten().unsqueeze(0).float(), - o_inv.flatten().unsqueeze(0).float() - ).item() - mse = (o.float() - o_inv.float()).pow(2).mean().item() - status = "✅" if cos > 0.999 else "❌" - print(f" inverse_rope → original: cosine={cos:.6f} MSE={mse:.6e} {status}") - return cos - - -def test_wo_a_grouped_linear(): - """Test CuTeDSL NVFP4 wo_a grouped linear against BF16 reference.""" - print("\n=== Test: wo_a NVFP4 Grouped Linear ===") - - torch.manual_seed(42) - num_tokens = NUM_TOKENS - - # Random attention output (after inverse RoPE) - o = torch.randn(num_tokens, N_LOCAL_GROUPS * HEADS_PER_GROUP, HEAD_DIM, - dtype=torch.bfloat16, device=DEVICE) * 2.0 - - # Random wo_a weight (BF16, grouped format) - # In vLLM, wo_a is ColumnParallelLinear with is_bmm=True - # Weight shape: (n_local_groups, heads_per_group * head_dim, o_lora_rank) - wo_a_weight = torch.randn( - N_LOCAL_GROUPS, GROUP_IN, O_LORA_RANK, - dtype=torch.bfloat16, device=DEVICE - ) * 0.1 - - # BF16 reference: grouped matmul - o_grouped = o.reshape(num_tokens, N_LOCAL_GROUPS, GROUP_IN) - z_ref = torch.empty(num_tokens, N_LOCAL_GROUPS, O_LORA_RANK, - dtype=torch.bfloat16, device=DEVICE) - for g in range(N_LOCAL_GROUPS): - # (tokens, GROUP_IN) × (GROUP_IN, O_LORA_RANK) → (tokens, O_LORA_RANK) - z_ref[:, g, :] = o_grouped[:, g, :] @ wo_a_weight[g] - - # CuTeDSL NVFP4 runner - runner = Nvfp4GroupedLinear( - n_local_groups=N_LOCAL_GROUPS, - heads_per_group=HEADS_PER_GROUP, - head_dim=HEAD_DIM, - o_lora_rank=O_LORA_RANK, - max_num_tokens=8192, - device=DEVICE, - ) - runner.set_bf16_weight(wo_a_weight) - runner.finalize_weights() - - # Warmup + compute activation global scale - runner._ensure_initialized() - runner.compute_activation_global_scale(o) - - # Run - with torch.no_grad(): - z_out = runner.run(o) - - # Compare - cos = F.cosine_similarity( - z_ref.flatten().unsqueeze(0).float(), - z_out.flatten().unsqueeze(0).float() - ).item() - mse = (z_ref.float() - z_out.float()).pow(2).mean().item() - status = "✅" if cos >= 0.98 else "❌" - print(f" wo_a grouped linear: cosine={cos:.6f} MSE={mse:.6e} {status}") - print(f" z_ref amax={z_ref.amax():.4f} z_out amax={z_out.amax():.4f}") - - return cos - - -def main(): - torch.cuda.set_device(0) - print("=== wo_a NVFP4 Grouped Linear + Inverse RoPE Tests ===") - - cos_rope = test_inverse_rope() - cos_woa = test_wo_a_grouped_linear() - - print(f"\n=== SUMMARY ===") - results = {"inverse_rope": cos_rope, "wo_a_grouped_linear": cos_woa} - all_pass = True - for name, cos in results.items(): - threshold = 0.999 if name == "inverse_rope" else 0.98 - status = "✅" if cos >= threshold else "❌" - if cos < threshold: - all_pass = False - print(f" {name}: cosine={cos:.6f} {status}") - - if all_pass: - print("\n✅ ALL PASS") - else: - print("\n❌ SOME FAILED") - - -if __name__ == "__main__": - main() diff --git a/tests/archive/test_wo_a_bmm.py b/tests/archive/test_wo_a_bmm.py deleted file mode 100644 index 9770669b..00000000 --- a/tests/archive/test_wo_a_bmm.py +++ /dev/null @@ -1,79 +0,0 @@ -"""Unit test: wo_a BF16 BMM reshape logic (CPU only). - -Verifies that the per-group BMM reshape in the BF16 wo_a path -produces the same result as the flat linear (when no TP sharding). - -Usage: python3 tests/test_wo_a_bmm.py -""" -import torch -import torch.nn.functional as F - - -def test_bmm_vs_flat(): - """Compare per-group BMM vs flat linear for wo_a.""" - # Simulate: n_local_groups=2, heads_per_group=8, head_dim=512, o_lora_rank=1024 - n_local_groups = 2 - heads_per_group = 8 - head_dim = 512 - o_lora_rank = 1024 - num_tokens = 4 - in_features = heads_per_group * head_dim # 4096 - out_features = n_local_groups * o_lora_rank # 2048 - - torch.manual_seed(42) - - # Random attention output after inverse RoPE - # Shape: (num_tokens, n_local_heads, head_dim) where n_local_heads = n_local_groups * heads_per_group - o_inv = torch.randn(num_tokens, n_local_groups * heads_per_group, head_dim, dtype=torch.bfloat16) - - # Random wo_a weight (ColumnParallelLinear, no TP sharding for this test) - # Weight shape: (out_features, in_features) = (2048, 4096) - wo_a_weight = torch.randn(out_features, in_features, dtype=torch.bfloat16) * 0.02 - - # Flat linear (the OLD broken way - would give wrong result if in_features != n_local_heads * head_dim) - # This test just verifies the BMM matches the flat case when dimensions align - - # BMM approach (NEW way): - # Reshape o_inv: (num_tokens, n_local_groups, heads_per_group * head_dim) - # -> permute: (n_local_groups, num_tokens, in_features) - o_grouped = o_inv.view(num_tokens, n_local_groups, heads_per_group * head_dim).permute(1, 0, 2) - - # Reshape weight: (out_features, in_features) -> (n_local_groups, o_lora_rank, in_features) - wo_a_w = wo_a_weight.view(n_local_groups, o_lora_rank, in_features) - - # BMM: (n_local_groups, num_tokens, in) @ (n_local_groups, in, o_lora_rank) - z_bmm = torch.bmm(o_grouped, wo_a_w.transpose(1, 2)) - # -> permute: (num_tokens, n_local_groups, o_lora_rank) - z_bmm = z_bmm.permute(1, 0, 2).reshape(num_tokens, n_local_groups * o_lora_rank) - - # Reference: per-group matmul (the ground truth) - z_ref = torch.zeros(num_tokens, n_local_groups, o_lora_rank, dtype=torch.bfloat16) - for g in range(n_local_groups): - # (num_tokens, in_features) @ (in_features, o_lora_rank) - z_ref[:, g, :] = o_grouped[g] @ wo_a_w[g].T - z_ref = z_ref.reshape(num_tokens, n_local_groups * o_lora_rank) - - cos = F.cosine_similarity(z_bmm.flatten().unsqueeze(0).float(), - z_ref.flatten().unsqueeze(0).float()).item() - mse = (z_bmm.float() - z_ref.float()).pow(2).mean().item() - - status = "✅" if cos > 0.9999 else "❌" - print(f"BMM vs flat: cosine={cos:.8f} MSE={mse:.2e} {status}") - - # Also verify shapes - assert o_grouped.shape == (n_local_groups, num_tokens, in_features), \ - f"o_grouped shape: {o_grouped.shape}" - assert wo_a_w.shape == (n_local_groups, o_lora_rank, in_features), \ - f"wo_a_w shape: {wo_a_w.shape}" - assert z_bmm.shape == (num_tokens, out_features), \ - f"z_bmm shape: {z_bmm.shape}" - - return cos - - -if __name__ == "__main__": - cos = test_bmm_vs_flat() - if cos > 0.9999: - print("\n✅ PASS") - else: - print("\n❌ FAIL") diff --git a/tests/archive/unit_test_128_128_vdiag.py b/tests/archive/unit_test_128_128_vdiag.py deleted file mode 100644 index 634df139..00000000 --- a/tests/archive/unit_test_128_128_vdiag.py +++ /dev/null @@ -1,385 +0,0 @@ -""" -Minimal PV-only test: Load P from GMEM to TMEM via QK-style MMA, then PV from TMEM. -Step 1: QK MMA writes FP32 S to TMEM (we know this works) -Step 2: Softmax packing writes BF16 P to TMEM (test this) -Step 3: PV MMA reads BF16 P from TMEM and V from SMEM, produces O - -But to isolate the bug, let me test just the PV MMA in isolation. -I'll write known BF16 values to TMEM using the softmax packing path, -then immediately read them back using the PV A-fragment path, -and compare. - -Actually, the simplest isolation test: -1. Do QK MMA to get S in TMEM (cosine 0.999999 verified) -2. Do softmax packing: S → P in TMEM (at offset 32) -3. Skip PV entirely — read P from TMEM using the C-fragment composition LOAD path -4. Output P to GMEM and compare against S.to(BF16) - -This tests whether the softmax packing writes P correctly to the same TMEM -that the PV would read from. - -But we can't easily read P from TMEM using the standard epilogue path -because the epilogue expects FP32 accumulator data. - -Alternative: Use the PV MMA with V=I (identity). If P is correct, -then P @ I = P. But V needs to be MN-major and (128, 128), not (128, 64). -The output would be (128, 128) which doesn't match our (128, 64) c tensor. - -Let me use V that selects the first 64 columns: V[k, n] = delta(k, n) for k in [0,63]. -This gives P @ V = P[:, :64], and the output is (128, 64). -But V is (128, 128) in the MMA K,N dims. V[k, n] for k in [0,127], n in [0,63]. -Hmm, this is getting complicated. Let me just do the identity approach with a (128, 128) output. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - - -class VDiag128: - """QK + softmax packing + PV with V=I to isolate PV MMA correctness. - Output should be P = S.to(BF16), i.e. (Q@K^T).bfloat16() - With V=I, O = P @ I = P. - But V is (K=128, N=128) in the MMA. We need a 128x128 identity in MN-major. - Output tensor is (128, 128). - """ - def __init__(self, mma_tiler_mn): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) - self.use_2cta_instrs = False # needed by epilogue_tma_store - self.epilog_sync_bar_id = 1 # needed by epilogue_tma_store - self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0, 1, 2, 3) - self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192 - self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) - # PV with V=I: output is (128, 128), same as QK - self.pv_mma_tiler = (self.qk_mma_tiler[0], self.qk_mma_tiler[1], self.qk_mma_tiler[1]) - # pv_mma_tiler = (128, 128, 128) since V is 128x128 - self.mma_tiler = self.qk_mma_tiler - - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = ( - self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), - self.qk_mma_tiler[1], self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape( - self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0) - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - s_cols = find_tmem_tensor_col_offset(tStS) - pv_thr = pv_mma.get_slice(0) - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - o_cols = find_tmem_tensor_col_offset(tOtO) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0 - self.tmem_p0_offset = 32 - self.tmem_o0_offset = s_cols - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") - - # ⛔⛔⛔ CRITICAL: num_tma_load_bytes MUST include ALL TMA-loaded tensors (Q + K + V). Missing V → DEADLOCK. See FOOTGUN #0 in README. - a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - self.num_tma_load_bytes = ( - cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + - cute.size_in_bytes(self.q_dtype, v_smem) - ) * cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) - # PV with 128x128 output (V=I) - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - - q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) - k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) - v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) - tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id), - q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id), - k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) - tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), - v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) - epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) - tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) - - self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, - tma_c, tma_tc, self.cluster_layout_vmnk, - self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, - tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - use_2cta = cute.size(qk_mma.thr_id.shape) == 2 - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - ab_p, ab_c = pipeline.PipelineTmaUmma.create( - barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ).make_participants() - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) - tCgV = pv_thr.partition_B(gV) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_acc_shape) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_acc_shape) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ═══ TMA LOAD WARP ═══ - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt, unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) - cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else 'FAIL')) - - -if __name__ == '__main__': - test() diff --git a/tests/archive/unit_test_fmha_v3_debug.py b/tests/archive/unit_test_fmha_v3_debug.py deleted file mode 100644 index 91ae2a68..00000000 --- a/tests/archive/unit_test_fmha_v3_debug.py +++ /dev/null @@ -1,469 +0,0 @@ -""" -FMHA v3 + Stage C: QK -> online softmax -> PV with KV-tile interleaving. -Stage C: row_max, exp2, O rescale, row_sum, final normalization. -FMHA pattern P store preserved from Stage B. -""" -import math -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - -HEAD_DIM = 64 - -class FmhaV3Softmax: - def __init__(self): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0,1,2,3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.num_c_stage = 2 - self.kv_stage = 2; self.q_stage = 1; self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (128, 128, qk_ik * 4) - pv_ik = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (128, HEAD_DIM, pv_ik * (128 // pv_ik)) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0]//cute.size(qk_mma.thr_id.shape), HEAD_DIM, self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - self.q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, self.q_dtype, self.q_stage) - self.k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, self.q_dtype, self.kv_stage) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, self.kv_stage) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - qk_thr = qk_mma.get_slice(0); qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - pv_thr = pv_mma.get_slice(0); pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - self.tmem_s0_offset = 0; self.tmem_p0_offset = 32 - # P occupies [tmem_p0_offset, tmem_p0_offset + p_cols_fp32) - # S occupies [0, qk_mma_tiler[1]) = [0, 128) - # O must NOT overlap P. Place O after max(S end, P end), aligned to 32. - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - p_end = self.tmem_p0_offset + p_cols_fp32 # 32 + 64 = 96 - s_cols = self.qk_mma_tiler[1] # 128 - o_after = max(s_cols, p_end) # 128 - self.tmem_o0_offset = ((o_after + 31) // 32) * 32 - self.tmem_vec_offset = 0 # Reuse S region for per-row inv_row_sum vector # align to 32 = 128 - self.tmem_vec_offset = 0 # Reuse S region (free after softmax loop) - o_cols = find_tmem_tensor_col_offset(tOtO) # footprint of O - total = self.tmem_o0_offset + o_cols - # Must be multiple of 32 AND power of 2 - self.num_tmem_alloc_cols = 1 - while self.num_tmem_alloc_cols < total: - self.num_tmem_alloc_cols *= 2 - cta = cute.size(qk_mma.thr_id.shape) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)) - self.q_tx_bytes = cute.size_in_bytes(self.q_dtype, q_s) * cta - self.kv_tx_bytes = cute.size_in_bytes(self.q_dtype, k_s) * cta - self.scale_softmax_log2 = Float32(1.0 / math.sqrt(HEAD_DIM) * math.log2(math.e)) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - # # s_k hardcoded # BROKEN in @cute.jit - # FMHA-style V: reconstruct as (HEAD_DIM, s_k, 1) MN-major - v_fmha = cute.make_tensor( - v.iterator, - cute.make_layout( - (HEAD_DIM, 128, 1), - stride=(1, HEAD_DIM, HEAD_DIM * 128), - ), - ) - self.v_major = LayoutEnum.from_tensor(v_fmha).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, self.a_major, self.b_major, self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, self.qk_acc_dtype, self.cta_group, (128,HEAD_DIM), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)); v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id),q,q_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id),k,k_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v_fmha,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape) - epi_s = cute.select(self.c_smem_s,mode=[0,1]) - tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile) - self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.v_smem_s,self.p_tmem_s,self.c_smem_s,self.epi_tile).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, q_smem_s, k_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx,_,_ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k); cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - q_bar: cute.struct.MemRange[cutlass.Int64, self.q_stage*2] - kv_bar: cute.struct.MemRange[cutlass.Int64, self.kv_stage*2] - s_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage*2] - tmem_dealloc: cutlass.Int64; holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - qp,qc = pipeline.PipelineTmaUmma.create(barrier_storage=st.q_bar.data_ptr(),num_stages=self.q_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.q_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - kvp,kvc = pipeline.PipelineTmaUmma.create(barrier_storage=st.kv_bar.data_ptr(),num_stages=self.kv_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.kv_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - s_prod,s_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.s_bar.data_ptr(),num_stages=1,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.epilogue_warp_id))).make_participants() - softmax_done_bar = pipeline.NamedBarrier(barrier_id=3, num_threads=32 + 32*len(self.epilogue_warp_id)) - pv_done_bar = pipeline.NamedBarrier(barrier_id=4, num_threads=32 + 32*len(self.epilogue_warp_id)) - acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(),num_stages=self.num_acc_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,len(self.epilogue_warp_id)),cta_layout_vmnk=cl_vmnk,defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2,num_threads=32*len((self.mma_warp_id,*self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr,barrier_for_retrieve=tmem_bar,allocator_warp_id=self.epilogue_warp_id[0],is_two_cta=cute.size(qk_mma.thr_id.shape)==2,two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk,is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype,layout=q_smem_s.outer,byte_alignment=128,swizzle=q_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype,layout=k_smem_s.outer,byte_alignment=128,swizzle=k_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype,layout=v_smem_s.outer,byte_alignment=128,swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype,layout=c_smem_s.outer,byte_alignment=128,swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None)) - gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None)) - gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None)) - gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None)) - n_kv_tiles = cute.size(gK, mode=[3]) - - qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - tCgV = pv_thr.partition_B(gV); tCgC = pv_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,0,None,0)).shape) - tAsQ,tAgQ = cpasync.tma_partition(tma_q,0,a_lay,cute.group_modes(sQ,0,3),cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape) - tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3)) - tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)]; tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - # --- PV read view (for MMA only, NOT for softmax store) --- - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None,None,None,0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # TMA LOAD - if warp_idx == self.tma_warp_id: - qp.reset(); qh = qp.acquire_and_advance() - cute.copy(tma_q,tAgQ[(None,qh.count)],tAsQ[(None,qh.index)],tma_bar_ptr=qh.barrier) - qp.tail() - kvp.reset(); pk = kvp.try_acquire() - for kt in cutlass.range(n_kv_tiles,unroll=1): - kh = kvp.acquire_and_advance(pk) - cute.copy(tma_k,tBgK[(None,kh.count)],tBsK[(None,kh.index)],tma_bar_ptr=kh.barrier) - pk = cutlass.Boolean(1) - vh = kvp.acquire_and_advance(pk) - cute.copy(tma_v,tVgV[(None,vh.count)],tVsV[(None,vh.index)],tma_bar_ptr=vh.barrier) - pk = cutlass.Boolean(1) - kvp.tail() - - # MMA - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - qc.reset(); qh = qc.wait_and_advance(); qh.release() - kvc.reset(); pk = kvc.try_wait() - acc_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_stage) - acc_pipe.producer_acquire(acc_st) - for kt in range(n_kv_tiles): - kh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - sh = s_prod.acquire_and_advance() - qk_mma.set(tcgen05.Field.ACCUMULATE, False) - for kb in cutlass.range(cute.size(tCrQ,mode=[2]), unroll_full=True): - cute.gemm(qk_mma, tStS0, tCrQ[(None,None,kb,0)], tCrK[(None,None,kb,kh.index)], tStS0) - qk_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - sh.commit(); kh.release() - softmax_done_bar.arrive_and_wait() - vh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - pv_mma.set(tcgen05.Field.ACCUMULATE, kt != 0) - for kb in cutlass.range(cute.size(tOrP0,mode=[2]), unroll_full=True): - cute.gemm(pv_mma, tOtO0, tOrP0[(None,None,kb)], tCrV[(None,None,kb,vh.index)], tOtO0) - pv_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - vh.release() - pv_done_bar.arrive() - acc_pipe.producer_commit(acc_st); acc_st.advance() - acc_pipe.producer_tail(acc_st) - - # ===================== EPILOGUE WARPS (STAGE C: ONLINE SOFTMAX) ===================== - if warp_idx < self.mma_warp_id: - tmem.allocate(self.num_tmem_alloc_cols) - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) - sfw_idx = tidx % (32 * len(self.epilogue_warp_id)) - - # --- S load (QK C-fragment) --- - tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS0) - thr_load = tiled_tmem_load.get_slice(sfw_idx) - tTMEM_LOADtS = thr_load.partition_S(tStS0) - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tTMEM_LOADcS = thr_load.partition_D(tScS) - - # --- P store (QK C-fragment composition, FMHA pattern) --- - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - tStP_layout = cute.composition(tStS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tStP0 = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStP_layout) - tmem_store_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStP0) - thr_store = tiled_tmem_store.get_slice(sfw_idx) - tTMEM_STOREtP = thr_store.partition_D(tStP0) - tScP_layout = cute.composition(tScS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tScP = cute.make_tensor(tScS.iterator, tScP_layout) - tTMEM_STOREcP = thr_store.partition_S(tScP) - - # --- Vector TMEM (per-row row_sum storage, FMHA pattern) --- - # composition(tStS.layout, (128, 2)) = 2 FP32 columns per logical row - # vec[0] = row_sum (final, after loop), vec[1] = unused - # Reuses S TMEM region (offset 0), free after softmax loop writes - - tStS_vec_layout = cute.composition(tStS.layout, cute.make_layout((128, 2))) - tStS_vec = cute.make_tensor(tStS.iterator + self.tmem_vec_offset, tStS_vec_layout) - tScS_vec_layout = cute.composition(tScS.layout, cute.make_layout((128, 2))) - tScS_vec = cute.make_tensor(tScS.iterator, tScS_vec_layout) - tmem_store_vec_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(2)), self.qk_acc_dtype) - tiled_tmem_store_vec = tcgen05.make_tmem_copy(tmem_store_vec_atom, tStS_vec) - thr_tmem_store_vec = tiled_tmem_store_vec.get_slice(sfw_idx) - tTMEM_STORE_VECtS = thr_tmem_store_vec.partition_D(tStS_vec) - tTMEM_STORE_VECcS = thr_tmem_store_vec.partition_S(tScS_vec) - tmem_load_vec_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(2)), self.qk_acc_dtype) - tiled_tmem_load_vec = tcgen05.make_tmem_copy(tmem_load_vec_atom, tStS_vec) - thr_tmem_load_vec = tiled_tmem_load_vec.get_slice(sfw_idx) - tTMEM_LOAD_VECtS = thr_tmem_load_vec.partition_S(tStS_vec) - tTMEM_LOAD_VECcS = thr_tmem_load_vec.partition_D(tScS_vec) - - # --- C6: O TMEM load/store for rescale (correction_rescale pattern) --- - corr_tile_size = 16 - cO = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1])) - tOcO = pv_thr.partition_C(cO) - o_tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), self.qk_acc_dtype) - o_tmem_store_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), self.qk_acc_dtype) - tOtO_i_layout = cute.composition(tOtO0.layout, cute.make_layout((128, corr_tile_size))) - tOcO_i_layout = cute.composition(tOcO.layout, cute.make_layout((128, corr_tile_size))) - tOtO_i = cute.make_tensor(tOtO0.iterator, tOtO_i_layout) - tOcO_i = cute.make_tensor(tOcO.iterator, tOcO_i_layout) - o_tiled_tmem_load = tcgen05.make_tmem_copy(o_tmem_load_atom, tOtO_i) - o_tiled_tmem_store = tcgen05.make_tmem_copy(o_tmem_store_atom, tOtO_i) - o_thr_load = o_tiled_tmem_load.get_slice(sfw_idx) - o_thr_store = o_tiled_tmem_store.get_slice(sfw_idx) - tTMEM_LOADtO = o_thr_load.partition_S(tOtO_i) - tTMEM_LOADcO = o_thr_load.partition_D(tOcO_i) - tTMEM_STOREtO = o_thr_store.partition_D(tOtO_i) - o_col_tiles = self.pv_mma_tiler[1] // corr_tile_size - - # --- C2: Per-thread row state (persist across KV tiles) --- - row_max = -cutlass.Float32.inf - row_sum = cutlass.Float32(0.0) - - # --- C3: QK scale = 1/sqrt(HEAD_DIM) * log2(e) for exp2 --- - scale = self.scale_softmax_log2 - - # ============================================================= - # Per-KV-tile online softmax loop - # ============================================================= - for kt in range(n_kv_tiles): - si_handle = s_cons.wait_and_advance() - - # Load S from TMEM (FP32, QK C-fragment layout) - tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS) - - # --- C4: Compute tile_max via .reduce(MAX) --- - old_row_max = row_max - row_max = tTMEM_LOADrS.load().reduce(cute.ReductionOp.MAX, row_max, 0) - row_max_safe = row_max - if row_max == -cutlass.Float32.inf: - row_max_safe = cutlass.Float32(0.0) - - # --- C5: Compute rescale factor --- - acc_scale = cute.math.exp2(scale * (old_row_max - row_max_safe), fastmath=True) - - # --- C6: Rescale O in TMEM (load O, multiply by acc_scale, store O) --- - # acc_scale belongs to QK row (N//4), but O rows are in PV partition (N). - # Store acc_scale to vector by QK row, read by PV row. - if kt > 0: - pv_done_bar.arrive_and_wait() - - # Store acc_scale to vector indexed by QK logical row - qk_row_c6 = tTMEM_LOADcS[0][0] - thr_vs_c6 = tiled_tmem_store_vec.get_slice(qk_row_c6) - tVStore_c6 = thr_vs_c6.partition_D(tStS_vec) - tVStoreSrc_c6 = thr_vs_c6.partition_S(tScS_vec) - tVStoreRmem_c6 = cute.make_rmem_tensor(tVStoreSrc_c6.shape, self.qk_acc_dtype) - tVStoreRmem_c6[0] = acc_scale - cute.copy(tiled_tmem_store_vec, tVStoreRmem_c6, tVStore_c6) - cute.arch.fence_view_async_tmem_store() - - # Read acc_scale from vector indexed by PV logical row - pv_row_c6 = tTMEM_LOADcO[0][0] - thr_vl_c6 = tiled_tmem_load_vec.get_slice(pv_row_c6) - tVLoad_c6 = thr_vl_c6.partition_S(tStS_vec) - tVLoadDst_c6 = thr_vl_c6.partition_D(tScS_vec) - tVLoadRmem_c6 = cute.make_rmem_tensor(tVLoadDst_c6.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load_vec, tVLoad_c6, tVLoadRmem_c6) - cute.arch.fence_view_async_tmem_load() - acc_scale_pv = tVLoadRmem_c6[0] - - tTMrO = cute.make_rmem_tensor((tTMEM_LOADcO.shape, o_col_tiles), self.qk_acc_dtype) - for i in range(o_col_tiles): - tTMrO_i_ = tTMrO[None, i] - tTMrO_i_layout = cute.composition(tTMrO_i_.layout, cute.make_layout(tTMrO.shape[0])) - tTMrO_i = cute.make_tensor(tTMrO_i_.iterator, tTMrO_i_layout) - tTMEM_LOADtO_i = cute.make_tensor(tTMEM_LOADtO.iterator + i * corr_tile_size, tTMEM_LOADtO.layout) - tTMEM_STOREtO_i = cute.make_tensor(tTMEM_STOREtO.iterator + i * corr_tile_size, tTMEM_STOREtO.layout) - cute.copy(o_tiled_tmem_load, tTMEM_LOADtO_i, tTMrO_i) - for j in cutlass.range(cute.size(tTMrO_i), vectorize=True): - tTMrO_i[j] = tTMrO_i[j] * acc_scale_pv - cute.copy(o_tiled_tmem_store, tTMrO_i, tTMEM_STOREtO_i) - cute.arch.fence_view_async_tmem_store() - - # Rescale row_sum - row_sum = row_sum * acc_scale - - # --- C7: Compute P = exp2((S - row_max_safe) * scale) --- - minus_row_max_scale = (cutlass.Float32(0.0) - row_max_safe) * scale - - # Register bridge (FMHA pattern: FP32 backing + BF16 view) - rP_words = cute.make_rmem_tensor(tTMEM_STOREcP.shape, self.qk_acc_dtype) - rP_bf16 = cute.make_tensor(cute.recast_ptr(rP_words.iterator, dtype=self.q_dtype), tTMEM_LOADrS.layout) - - frg_cnt = 4 - frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt - tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile)) - rP_bf16_frg = cute.logical_divide(rP_bf16, cute.make_layout(frg_tile)) - - # Scale S, compute exp2, store through register bridge - for j in range(frg_cnt): - for k in cutlass.range(cute.size(tTMEM_LOADrS_frg, mode=[0]), vectorize=True): - tTMEM_LOADrS_frg[k, j] = tTMEM_LOADrS_frg[k, j] * scale + minus_row_max_scale - tTMEM_LOADrS_frg[k, j] = cute.math.exp2(tTMEM_LOADrS_frg[k, j], fastmath=True) - s_vec = tTMEM_LOADrS_frg[None, j].load() - rP_bf16_frg[None, j].store(s_vec.to(self.q_dtype)) - - # Store P to TMEM - cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP) - cute.arch.fence_view_async_tmem_store() - si_handle.release() - softmax_done_bar.arrive() - - # --- C8: Row sum accumulation (CUTLASS FMHA packed f32x2 pattern) --- - # P values still in tTMEM_LOADrS registers. - # 4 accumulators for 4 reduction_unroll columns. - local_row_sum_0 = (cutlass.Float32(0.0), cutlass.Float32(0.0)) - local_row_sum_1 = (cutlass.Float32(0.0), cutlass.Float32(0.0)) - local_row_sum_2 = (cutlass.Float32(0.0), cutlass.Float32(0.0)) - local_row_sum_3 = (cutlass.Float32(0.0), cutlass.Float32(0.0)) - - reduction_unroll = 4 - rfrg_tile = cute.size(tTMEM_LOADrS) // reduction_unroll - tTMEM_LOADrS_rfrg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(rfrg_tile)) - - for j in cutlass.range_constexpr(0, cute.size(tTMEM_LOADrS_rfrg, mode=[0]), 2): - local_row_sum_0 = cute.arch.add_packed_f32x2( - local_row_sum_0, (tTMEM_LOADrS_rfrg[j, 0], tTMEM_LOADrS_rfrg[j + 1, 0])) - local_row_sum_1 = cute.arch.add_packed_f32x2( - local_row_sum_1, (tTMEM_LOADrS_rfrg[j, 1], tTMEM_LOADrS_rfrg[j + 1, 1])) - local_row_sum_2 = cute.arch.add_packed_f32x2( - local_row_sum_2, (tTMEM_LOADrS_rfrg[j, 2], tTMEM_LOADrS_rfrg[j + 1, 2])) - local_row_sum_3 = cute.arch.add_packed_f32x2( - local_row_sum_3, (tTMEM_LOADrS_rfrg[j, 3], tTMEM_LOADrS_rfrg[j + 1, 3])) - - local_row_sum_0 = cute.arch.add_packed_f32x2(local_row_sum_0, local_row_sum_1) - local_row_sum_2 = cute.arch.add_packed_f32x2(local_row_sum_2, local_row_sum_3) - local_row_sum_0 = cute.arch.add_packed_f32x2(local_row_sum_0, local_row_sum_2) - tile_sum = local_row_sum_0[0] + local_row_sum_0[1] - - row_sum = row_sum + tile_sum - - # --- C9: Final normalization via O TMEM rescale --- - pv_done_bar.arrive_and_wait() - - # Use QK-accumulated row_sum directly (DEBUG: check if row mapping matches PV) - inv_row_sum = cutlass.Float32(1.0) / row_sum - - # Normalize O in TMEM using PV-correct inv_row_sum - tTMrO_final = cute.make_rmem_tensor((tTMEM_LOADcO.shape, o_col_tiles), self.qk_acc_dtype) - for i in range(o_col_tiles): - tTMrO_i_ = tTMrO_final[None, i] - tTMrO_i_layout = cute.composition(tTMrO_i_.layout, cute.make_layout(tTMrO_final.shape[0])) - tTMrO_i = cute.make_tensor(tTMrO_i_.iterator, tTMrO_i_layout) - tTMEM_LOADtO_i = cute.make_tensor( - tTMEM_LOADtO.iterator + i * corr_tile_size, tTMEM_LOADtO.layout) - tTMEM_STOREtO_i = cute.make_tensor( - tTMEM_STOREtO.iterator + i * corr_tile_size, tTMEM_STOREtO.layout) - cute.copy(o_tiled_tmem_load, tTMEM_LOADtO_i, tTMrO_i) - for j in cutlass.range(cute.size(tTMrO_i), vectorize=True): - tTMrO_i[j] = tTMrO_i[j] * inv_row_sum - cute.copy(o_tiled_tmem_store, tTMrO_i, tTMEM_STOREtO_i) - cute.arch.fence_view_async_tmem_store() - - # Now O in TMEM is normalized. Use standard epilogue_tma_store with identity. - tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout) - acc_cons_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage) - c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) - c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp) - acc_cons_st = utils.gemm.sm100.epilogue_tma_store( - self, tidx, warp_idx, tma_c, tCtO_base, sC, tCgC, epi_tile, 0, - const_expr(lambda x: x), - (0,0,0), acc_cons_st, acc_pipe, c_pipe) - c_pipe.producer_tail() - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def test(): - import math - torch.manual_seed(42) - for n in [128, 256, 384]: - m, hd = 128, HEAD_DIM - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device="cuda") - k = torch.randn(n, hd, 1, dtype=torch.bfloat16, device="cuda") - v = torch.randn(n, hd, dtype=torch.bfloat16, device="cuda") - v_kernel = v.unsqueeze(-1) - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device="cuda") - qf = q[:,:,0].float(); kf = k[:,:,0].float() - attn = qf @ kf.T / math.sqrt(hd) - ref = torch.softmax(attn, dim=-1) @ v.float() - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - kernel = FmhaV3Softmax() - print(f"n={n}: Compiling...", flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - print(f"n={n}: tmem: s0={kernel.tmem_s0_offset} p0={kernel.tmem_p0_offset} o0={kernel.tmem_o0_offset} vec={kernel.tmem_vec_offset} alloc={kernel.num_tmem_alloc_cols}", flush=True) - print(f"n={n}: Running...", flush=True) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - out = c[:,:,0].float() - cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - max_err = (out - ref).abs().max().item() - print(f"FMHA softmax n={n}: cosine {cos:.6f} max_err {max_err:.6f} {'PASS' if cos >= 0.999 else 'FAIL'}", flush=True) - -if __name__ == "__main__": - test() - - diff --git a/tests/archive/unit_test_fmha_v3_fixed_v.py b/tests/archive/unit_test_fmha_v3_fixed_v.py deleted file mode 100644 index a9e968f6..00000000 --- a/tests/archive/unit_test_fmha_v3_fixed_v.py +++ /dev/null @@ -1,512 +0,0 @@ -""" -FMHA v3 + Stage C: QK -> online softmax -> PV with KV-tile interleaving. -Stage C: row_max, exp2, O rescale, row_sum, final normalization. -FMHA pattern P store preserved from Stage B. -""" -import math -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - -HEAD_DIM = 64 - -class FmhaV3Softmax: - def __init__(self, s_k=128): - self.s_k = s_k - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0,1,2,3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.num_c_stage = 2 - self.kv_stage = 2; self.q_stage = 1; self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (128, 128, qk_ik * 4) - pv_ik = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (128, HEAD_DIM, pv_ik * (128 // pv_ik)) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0]//cute.size(qk_mma.thr_id.shape), HEAD_DIM, self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - self.q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, self.q_dtype, self.q_stage) - self.k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, self.q_dtype, self.kv_stage) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, self.kv_stage) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - qk_thr = qk_mma.get_slice(0); qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - pv_thr = pv_mma.get_slice(0); pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - self.tmem_s0_offset = 0; self.tmem_p0_offset = 32 - # P occupies [tmem_p0_offset, tmem_p0_offset + p_cols_fp32) - # S occupies [0, qk_mma_tiler[1]) = [0, 128) - # O must NOT overlap P. Place O after max(S end, P end), aligned to 32. - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - p_end = self.tmem_p0_offset + p_cols_fp32 # 32 + 64 = 96 - s_cols = self.qk_mma_tiler[1] # 128 - o_after = max(s_cols, p_end) # 128 - self.tmem_o0_offset = ((o_after + 31) // 32) * 32 - self.tmem_vec_offset = 0 # Reuse S region for per-row inv_row_sum vector # align to 32 = 128 - self.tmem_vec_offset = 0 # Reuse S region (free after softmax loop) - o_cols = find_tmem_tensor_col_offset(tOtO) # footprint of O - total = self.tmem_o0_offset + o_cols - # Must be multiple of 32 AND power of 2 - self.num_tmem_alloc_cols = 1 - while self.num_tmem_alloc_cols < total: - self.num_tmem_alloc_cols *= 2 - cta = cute.size(qk_mma.thr_id.shape) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)) - self.q_tx_bytes = cute.size_in_bytes(self.q_dtype, q_s) * cta - self.kv_tx_bytes = cute.size_in_bytes(self.q_dtype, k_s) * cta - self.scale_softmax_log2 = Float32(1.0 / math.sqrt(HEAD_DIM) * math.log2(math.e)) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - # # s_k hardcoded # BROKEN in @cute.jit - # FMHA-style V: reconstruct as (HEAD_DIM, s_k, 1) MN-major - v_fmha = cute.make_tensor( - v.iterator, - cute.make_layout( - (HEAD_DIM, 128, 1), - stride=(1, HEAD_DIM, HEAD_DIM * 128), - ), - ) - self.v_major = LayoutEnum.from_tensor(v_fmha).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, self.a_major, self.b_major, self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, self.qk_acc_dtype, self.cta_group, (128,HEAD_DIM), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)); v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id),q,q_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id),k,k_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v_fmha,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape) - epi_s = cute.select(self.c_smem_s,mode=[0,1]) - tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile) - self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.v_smem_s,self.p_tmem_s,self.c_smem_s,self.epi_tile).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, q_smem_s, k_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx,_,_ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k); cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - q_bar: cute.struct.MemRange[cutlass.Int64, self.q_stage*2] - kv_bar: cute.struct.MemRange[cutlass.Int64, self.kv_stage*2] - s_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage*2] - tmem_dealloc: cutlass.Int64; holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - qp,qc = pipeline.PipelineTmaUmma.create(barrier_storage=st.q_bar.data_ptr(),num_stages=self.q_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.q_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - kvp,kvc = pipeline.PipelineTmaUmma.create(barrier_storage=st.kv_bar.data_ptr(),num_stages=self.kv_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.kv_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - s_prod,s_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.s_bar.data_ptr(),num_stages=1,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.epilogue_warp_id))).make_participants() - softmax_done_bar = pipeline.NamedBarrier(barrier_id=3, num_threads=32 + 32*len(self.epilogue_warp_id)) - pv_done_bar = pipeline.NamedBarrier(barrier_id=4, num_threads=32 + 32*len(self.epilogue_warp_id)) - acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(),num_stages=self.num_acc_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,len(self.epilogue_warp_id)),cta_layout_vmnk=cl_vmnk,defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2,num_threads=32*len((self.mma_warp_id,*self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr,barrier_for_retrieve=tmem_bar,allocator_warp_id=self.epilogue_warp_id[0],is_two_cta=cute.size(qk_mma.thr_id.shape)==2,two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk,is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype,layout=q_smem_s.outer,byte_alignment=128,swizzle=q_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype,layout=k_smem_s.outer,byte_alignment=128,swizzle=k_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype,layout=v_smem_s.outer,byte_alignment=128,swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype,layout=c_smem_s.outer,byte_alignment=128,swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None)) - gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None)) - gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None)) - gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None)) - n_kv_tiles = cute.size(gK, mode=[3]) - - qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - tCgV = pv_thr.partition_B(gV); tCgC = pv_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,0,None,0)).shape) - tAsQ,tAgQ = cpasync.tma_partition(tma_q,0,a_lay,cute.group_modes(sQ,0,3),cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape) - tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3)) - tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)]; tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - # --- PV read view (for MMA only, NOT for softmax store) --- - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None,None,None,0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # TMA LOAD - if warp_idx == self.tma_warp_id: - qp.reset(); qh = qp.acquire_and_advance() - cute.copy(tma_q,tAgQ[(None,qh.count)],tAsQ[(None,qh.index)],tma_bar_ptr=qh.barrier) - qp.tail() - kvp.reset(); pk = kvp.try_acquire() - for kt in cutlass.range(n_kv_tiles,unroll=1): - kh = kvp.acquire_and_advance(pk) - cute.copy(tma_k,tBgK[(None,kh.count)],tBsK[(None,kh.index)],tma_bar_ptr=kh.barrier) - pk = cutlass.Boolean(1) - vh = kvp.acquire_and_advance(pk) - cute.copy(tma_v,tVgV[(None,vh.count)],tVsV[(None,vh.index)],tma_bar_ptr=vh.barrier) - pk = cutlass.Boolean(1) - kvp.tail() - - # MMA - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - qc.reset(); qh = qc.wait_and_advance(); qh.release() - kvc.reset(); pk = kvc.try_wait() - acc_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_stage) - acc_pipe.producer_acquire(acc_st) - for kt in range(n_kv_tiles): - kh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - sh = s_prod.acquire_and_advance() - qk_mma.set(tcgen05.Field.ACCUMULATE, False) - for kb in cutlass.range(cute.size(tCrQ,mode=[2]), unroll_full=True): - cute.gemm(qk_mma, tStS0, tCrQ[(None,None,kb,0)], tCrK[(None,None,kb,kh.index)], tStS0) - qk_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - sh.commit(); kh.release() - softmax_done_bar.arrive_and_wait() - vh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - pv_mma.set(tcgen05.Field.ACCUMULATE, kt != 0) - for kb in cutlass.range(cute.size(tOrP0,mode=[2]), unroll_full=True): - cute.gemm(pv_mma, tOtO0, tOrP0[(None,None,kb)], tCrV[(None,None,kb,vh.index)], tOtO0) - pv_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - vh.release() - pv_done_bar.arrive() - acc_pipe.producer_commit(acc_st); acc_st.advance() - acc_pipe.producer_tail(acc_st) - - # ===================== EPILOGUE WARPS (STAGE C: ONLINE SOFTMAX) ===================== - if warp_idx < self.mma_warp_id: - tmem.allocate(self.num_tmem_alloc_cols) - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) - sfw_idx = tidx % (32 * len(self.epilogue_warp_id)) - - # --- S load (QK C-fragment) --- - tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS0) - thr_load = tiled_tmem_load.get_slice(sfw_idx) - tTMEM_LOADtS = thr_load.partition_S(tStS0) - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tTMEM_LOADcS = thr_load.partition_D(tScS) - - # --- P store (QK C-fragment composition, FMHA pattern) --- - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - tStP_layout = cute.composition(tStS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tStP0 = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStP_layout) - tmem_store_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStP0) - thr_store = tiled_tmem_store.get_slice(sfw_idx) - tTMEM_STOREtP = thr_store.partition_D(tStP0) - tScP_layout = cute.composition(tScS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tScP = cute.make_tensor(tScS.iterator, tScP_layout) - tTMEM_STOREcP = thr_store.partition_S(tScP) - - # --- Vector TMEM (per-row row_sum storage, FMHA pattern) --- - # composition(tStS.layout, (128, 2)) = 2 FP32 columns per logical row - # vec[0] = row_sum (final, after loop), vec[1] = unused - # Reuses S TMEM region (offset 0), free after softmax loop writes - - tStS_vec_layout = cute.composition(tStS.layout, cute.make_layout((128, 2))) - tStS_vec = cute.make_tensor(tStS.iterator + self.tmem_vec_offset, tStS_vec_layout) - tScS_vec_layout = cute.composition(tScS.layout, cute.make_layout((128, 2))) - tScS_vec = cute.make_tensor(tScS.iterator, tScS_vec_layout) - tmem_store_vec_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(2)), self.qk_acc_dtype) - tiled_tmem_store_vec = tcgen05.make_tmem_copy(tmem_store_vec_atom, tStS_vec) - thr_tmem_store_vec = tiled_tmem_store_vec.get_slice(sfw_idx) - tTMEM_STORE_VECtS = thr_tmem_store_vec.partition_D(tStS_vec) - tTMEM_STORE_VECcS = thr_tmem_store_vec.partition_S(tScS_vec) - tmem_load_vec_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(2)), self.qk_acc_dtype) - tiled_tmem_load_vec = tcgen05.make_tmem_copy(tmem_load_vec_atom, tStS_vec) - thr_tmem_load_vec = tiled_tmem_load_vec.get_slice(sfw_idx) - tTMEM_LOAD_VECtS = thr_tmem_load_vec.partition_S(tStS_vec) - tTMEM_LOAD_VECcS = thr_tmem_load_vec.partition_D(tScS_vec) - - # --- C6: O TMEM load/store for rescale (correction_rescale pattern) --- - corr_tile_size = 16 - cO = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1])) - tOcO = pv_thr.partition_C(cO) - o_tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), self.qk_acc_dtype) - o_tmem_store_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), self.qk_acc_dtype) - tOtO_i_layout = cute.composition(tOtO0.layout, cute.make_layout((128, corr_tile_size))) - tOcO_i_layout = cute.composition(tOcO.layout, cute.make_layout((128, corr_tile_size))) - tOtO_i = cute.make_tensor(tOtO0.iterator, tOtO_i_layout) - tOcO_i = cute.make_tensor(tOcO.iterator, tOcO_i_layout) - o_tiled_tmem_load = tcgen05.make_tmem_copy(o_tmem_load_atom, tOtO_i) - o_tiled_tmem_store = tcgen05.make_tmem_copy(o_tmem_store_atom, tOtO_i) - o_thr_load = o_tiled_tmem_load.get_slice(sfw_idx) - o_thr_store = o_tiled_tmem_store.get_slice(sfw_idx) - tTMEM_LOADtO = o_thr_load.partition_S(tOtO_i) - tTMEM_LOADcO = o_thr_load.partition_D(tOcO_i) - tTMEM_STOREtO = o_thr_store.partition_D(tOtO_i) - o_col_tiles = self.pv_mma_tiler[1] // corr_tile_size - - # --- C2: Per-thread row state (persist across KV tiles) --- - row_max = -cutlass.Float32.inf - row_sum = cutlass.Float32(0.0) - - # --- C3: QK scale = 1/sqrt(HEAD_DIM) * log2(e) for exp2 --- - scale = self.scale_softmax_log2 - - # ============================================================= - # Per-KV-tile online softmax loop - # ============================================================= - for kt in range(n_kv_tiles): - si_handle = s_cons.wait_and_advance() - - # Load S from TMEM (FP32, QK C-fragment layout) - tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS) - - # --- C4: Compute tile_max via .reduce(MAX) --- - old_row_max = row_max - row_max = tTMEM_LOADrS.load().reduce(cute.ReductionOp.MAX, row_max, 0) - row_max_safe = row_max - if row_max == -cutlass.Float32.inf: - row_max_safe = cutlass.Float32(0.0) - - # --- C5: Compute rescale factor --- - acc_scale = cute.math.exp2(scale * (old_row_max - row_max_safe), fastmath=True) - - # --- C6: Rescale O in TMEM (load O, multiply by acc_scale, store O) --- - # acc_scale belongs to QK row (N//4), but O rows are in PV partition (N). - # Store acc_scale to vector by QK row, read by PV row. - if kt > 0: - pv_done_bar.arrive_and_wait() - - # Store acc_scale to vector indexed by QK logical row - qk_row_c6 = tTMEM_LOADcS[0][0] - thr_vs_c6 = tiled_tmem_store_vec.get_slice(qk_row_c6) - tVStore_c6 = thr_vs_c6.partition_D(tStS_vec) - tVStoreSrc_c6 = thr_vs_c6.partition_S(tScS_vec) - tVStoreRmem_c6 = cute.make_rmem_tensor(tVStoreSrc_c6.shape, self.qk_acc_dtype) - tVStoreRmem_c6[0] = acc_scale - cute.copy(tiled_tmem_store_vec, tVStoreRmem_c6, tVStore_c6) - cute.arch.fence_view_async_tmem_store() - - # Read acc_scale from vector indexed by PV logical row - pv_row_c6 = tTMEM_LOADcO[0][0] - thr_vl_c6 = tiled_tmem_load_vec.get_slice(pv_row_c6) - tVLoad_c6 = thr_vl_c6.partition_S(tStS_vec) - tVLoadDst_c6 = thr_vl_c6.partition_D(tScS_vec) - tVLoadRmem_c6 = cute.make_rmem_tensor(tVLoadDst_c6.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load_vec, tVLoad_c6, tVLoadRmem_c6) - cute.arch.fence_view_async_tmem_load() - acc_scale_pv = tVLoadRmem_c6[0] - - tTMrO = cute.make_rmem_tensor((tTMEM_LOADcO.shape, o_col_tiles), self.qk_acc_dtype) - for i in range(o_col_tiles): - tTMrO_i_ = tTMrO[None, i] - tTMrO_i_layout = cute.composition(tTMrO_i_.layout, cute.make_layout(tTMrO.shape[0])) - tTMrO_i = cute.make_tensor(tTMrO_i_.iterator, tTMrO_i_layout) - tTMEM_LOADtO_i = cute.make_tensor(tTMEM_LOADtO.iterator + i * corr_tile_size, tTMEM_LOADtO.layout) - tTMEM_STOREtO_i = cute.make_tensor(tTMEM_STOREtO.iterator + i * corr_tile_size, tTMEM_STOREtO.layout) - cute.copy(o_tiled_tmem_load, tTMEM_LOADtO_i, tTMrO_i) - for j in cutlass.range(cute.size(tTMrO_i), vectorize=True): - tTMrO_i[j] = tTMrO_i[j] * acc_scale_pv - cute.copy(o_tiled_tmem_store, tTMrO_i, tTMEM_STOREtO_i) - cute.arch.fence_view_async_tmem_store() - - # Rescale row_sum - row_sum = row_sum * acc_scale - - # --- C7: Compute P = exp2((S - row_max_safe) * scale) --- - minus_row_max_scale = (cutlass.Float32(0.0) - row_max_safe) * scale - - # Register bridge (FMHA pattern: FP32 backing + BF16 view) - rP_words = cute.make_rmem_tensor(tTMEM_STOREcP.shape, self.qk_acc_dtype) - rP_bf16 = cute.make_tensor(cute.recast_ptr(rP_words.iterator, dtype=self.q_dtype), tTMEM_LOADrS.layout) - - frg_cnt = 4 - frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt - tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile)) - rP_bf16_frg = cute.logical_divide(rP_bf16, cute.make_layout(frg_tile)) - - # Scale S, compute exp2, store through register bridge - for j in range(frg_cnt): - for k in cutlass.range(cute.size(tTMEM_LOADrS_frg, mode=[0]), vectorize=True): - tTMEM_LOADrS_frg[k, j] = tTMEM_LOADrS_frg[k, j] * scale + minus_row_max_scale - tTMEM_LOADrS_frg[k, j] = cute.math.exp2(tTMEM_LOADrS_frg[k, j], fastmath=True) - s_vec = tTMEM_LOADrS_frg[None, j].load() - rP_bf16_frg[None, j].store(s_vec.to(self.q_dtype)) - - # Store P to TMEM - cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP) - cute.arch.fence_view_async_tmem_store() - si_handle.release() - softmax_done_bar.arrive() - - # --- C8: Row sum accumulation (CUTLASS FMHA packed f32x2 pattern) --- - # P values still in tTMEM_LOADrS registers. - # 4 accumulators for 4 reduction_unroll columns. - local_row_sum_0 = (cutlass.Float32(0.0), cutlass.Float32(0.0)) - local_row_sum_1 = (cutlass.Float32(0.0), cutlass.Float32(0.0)) - local_row_sum_2 = (cutlass.Float32(0.0), cutlass.Float32(0.0)) - local_row_sum_3 = (cutlass.Float32(0.0), cutlass.Float32(0.0)) - - reduction_unroll = 4 - rfrg_tile = cute.size(tTMEM_LOADrS) // reduction_unroll - tTMEM_LOADrS_rfrg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(rfrg_tile)) - - for j in cutlass.range_constexpr(0, cute.size(tTMEM_LOADrS_rfrg, mode=[0]), 2): - local_row_sum_0 = cute.arch.add_packed_f32x2( - local_row_sum_0, (tTMEM_LOADrS_rfrg[j, 0], tTMEM_LOADrS_rfrg[j + 1, 0])) - local_row_sum_1 = cute.arch.add_packed_f32x2( - local_row_sum_1, (tTMEM_LOADrS_rfrg[j, 1], tTMEM_LOADrS_rfrg[j + 1, 1])) - local_row_sum_2 = cute.arch.add_packed_f32x2( - local_row_sum_2, (tTMEM_LOADrS_rfrg[j, 2], tTMEM_LOADrS_rfrg[j + 1, 2])) - local_row_sum_3 = cute.arch.add_packed_f32x2( - local_row_sum_3, (tTMEM_LOADrS_rfrg[j, 3], tTMEM_LOADrS_rfrg[j + 1, 3])) - - local_row_sum_0 = cute.arch.add_packed_f32x2(local_row_sum_0, local_row_sum_1) - local_row_sum_2 = cute.arch.add_packed_f32x2(local_row_sum_2, local_row_sum_3) - local_row_sum_0 = cute.arch.add_packed_f32x2(local_row_sum_0, local_row_sum_2) - tile_sum = local_row_sum_0[0] + local_row_sum_0[1] - - row_sum = row_sum + tile_sum - - # --- C9: Final normalization via O TMEM rescale --- - pv_done_bar.arrive_and_wait() - - # Compute inv_row_sum from P in TMEM using PV partition. - # P was stored by softmax loop into TMEM at offset tmem_p0_offset. - # PV partition maps thread N to PV row N, so reading P via PV partition - # gives the correct per-row P values to sum. - # This avoids the QK→PV row mapping mismatch (QK: N->N//4, PV: N->N). - - # P is stored as BF16 in TMEM at tmem_p0_offset. - # We need to read it via PV TMEM load and sum the values. - # P has shape (128, HEAD_DIM//2) in FP32 columns (64 BF16 = 32 FP32 cols). - # Use the P TMEM load partition (PV A-fragment read). - - # Actually, P was stored via QK C-fragment store (St32x32bOp Repetition(32)). - # To read it via PV partition, we need a PV-partitioned load from the P region. - # Let's use the same o_tiled_tmem_load but pointed at P's TMEM offset. - - # P occupies TMEM columns [tmem_p0_offset, tmem_p0_offset + p_cols_fp32) - # In the PV C-fragment, P is the A-fragment. We can use tOrP0's layout. - # tOrP0 was set up with offset for PV MMA read. - - # Simpler: sum O across columns to get unnormalized row sum, then normalize. - # For V=identity, O = P@V = sum(P per row). So O.sum(dim=-1) = row_sum. - # For arbitrary V, O = P@V. O.sum(dim=-1) = sum_j(P@V)[j] = sum_j(sum_i P[i]*V[i,j]) - # This is NOT sum(P). So this trick only works for V=identity. - - # Correct approach: read P from TMEM, sum it per PV row. - # P is at TMEM offset tmem_p0_offset, stored as BF16 with St32x32bOp. - # P shape in TMEM: 128 rows x (HEAD_DIM BF16 = 32 FP32 cols) - # We can read P using Ld32x32bOp(Repetition(corr_tile_size)) via PV O-partition. - - # Use PV O TMEM load to read from P region instead of O region - p_col_tiles = p_cols_fp32 // corr_tile_size # 32 // 16 = 2 - pv_row_sum = cutlass.Float32(0.0) - for i in range(p_col_tiles): - # Read P tile from TMEM at P offset (not O offset) - tTMEM_LOADtP_i = cute.make_tensor( - tTMEM_LOADtO.iterator + (self.tmem_p0_offset - self.tmem_o0_offset) + i * corr_tile_size, - tTMEM_LOADtO.layout) - tTMrP_i = cute.make_rmem_tensor(tTMEM_LOADcO.shape, self.qk_acc_dtype) - cute.copy(o_tiled_tmem_load, tTMEM_LOADtP_i, tTMrP_i) - # Use .reduce(SUM) instead of scalar accumulation (vectorizer can't handle scalar in vectorized loop) - tile_p_sum = tTMrP_i.load().reduce(cute.ReductionOp.ADD, cutlass.Float32(0.0), 0) - pv_row_sum = pv_row_sum + tile_p_sum - - inv_row_sum = cutlass.Float32(1.0) / pv_row_sum - - # Normalize O in TMEM using PV-correct inv_row_sum - tTMrO_final = cute.make_rmem_tensor((tTMEM_LOADcO.shape, o_col_tiles), self.qk_acc_dtype) - for i in range(o_col_tiles): - tTMrO_i_ = tTMrO_final[None, i] - tTMrO_i_layout = cute.composition(tTMrO_i_.layout, cute.make_layout(tTMrO_final.shape[0])) - tTMrO_i = cute.make_tensor(tTMrO_i_.iterator, tTMrO_i_layout) - tTMEM_LOADtO_i = cute.make_tensor( - tTMEM_LOADtO.iterator + i * corr_tile_size, tTMEM_LOADtO.layout) - tTMEM_STOREtO_i = cute.make_tensor( - tTMEM_STOREtO.iterator + i * corr_tile_size, tTMEM_STOREtO.layout) - cute.copy(o_tiled_tmem_load, tTMEM_LOADtO_i, tTMrO_i) - for j in cutlass.range(cute.size(tTMrO_i), vectorize=True): - tTMrO_i[j] = tTMrO_i[j] * inv_row_sum - cute.copy(o_tiled_tmem_store, tTMrO_i, tTMEM_STOREtO_i) - cute.arch.fence_view_async_tmem_store() - - # Now O in TMEM is normalized. Use standard epilogue_tma_store with identity. - tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout) - acc_cons_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage) - c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) - c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp) - acc_cons_st = utils.gemm.sm100.epilogue_tma_store( - self, tidx, warp_idx, tma_c, tCtO_base, sC, tCgC, epi_tile, 0, - const_expr(lambda x: x), - (0,0,0), acc_cons_st, acc_pipe, c_pipe) - c_pipe.producer_tail() - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def test(): - import math - torch.manual_seed(42) - for n in [128, 256, 384]: - m, hd = 128, HEAD_DIM - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device="cuda") - k = torch.randn(n, hd, 1, dtype=torch.bfloat16, device="cuda") - v = torch.randn(n, hd, dtype=torch.bfloat16, device="cuda") - v_kernel = v.unsqueeze(-1) - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device="cuda") - qf = q[:,:,0].float(); kf = k[:,:,0].float() - attn = qf @ kf.T / math.sqrt(hd) - ref = torch.softmax(attn, dim=-1) @ v.float() - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - kernel = FmhaV3Softmax(s_k=n) - print(f"n={n}: Compiling...", flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - print(f"n={n}: tmem: s0={kernel.tmem_s0_offset} p0={kernel.tmem_p0_offset} o0={kernel.tmem_o0_offset} vec={kernel.tmem_vec_offset} alloc={kernel.num_tmem_alloc_cols}", flush=True) - print(f"n={n}: Running...", flush=True) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - out = c[:,:,0].float() - cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - max_err = (out - ref).abs().max().item() - print(f"FMHA softmax n={n}: cosine {cos:.6f} max_err {max_err:.6f} {'PASS' if cos >= 0.999 else 'FAIL'}", flush=True) - -if __name__ == "__main__": - test() - - diff --git a/tests/archive/unit_test_fmha_v3_noop_c9.py b/tests/archive/unit_test_fmha_v3_noop_c9.py deleted file mode 100644 index b9a11c6c..00000000 --- a/tests/archive/unit_test_fmha_v3_noop_c9.py +++ /dev/null @@ -1,469 +0,0 @@ -""" -FMHA v3 + Stage C: QK -> online softmax -> PV with KV-tile interleaving. -Stage C: row_max, exp2, O rescale, row_sum, final normalization. -FMHA pattern P store preserved from Stage B. -""" -import math -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - -HEAD_DIM = 64 - -class FmhaV3Softmax: - def __init__(self): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0,1,2,3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.num_c_stage = 2 - self.kv_stage = 2; self.q_stage = 1; self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (128, 128, qk_ik * 4) - pv_ik = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (128, HEAD_DIM, pv_ik * (128 // pv_ik)) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0]//cute.size(qk_mma.thr_id.shape), HEAD_DIM, self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - self.q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, self.q_dtype, self.q_stage) - self.k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, self.q_dtype, self.kv_stage) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, self.kv_stage) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - qk_thr = qk_mma.get_slice(0); qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - pv_thr = pv_mma.get_slice(0); pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - self.tmem_s0_offset = 0; self.tmem_p0_offset = 32 - # P occupies [tmem_p0_offset, tmem_p0_offset + p_cols_fp32) - # S occupies [0, qk_mma_tiler[1]) = [0, 128) - # O must NOT overlap P. Place O after max(S end, P end), aligned to 32. - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - p_end = self.tmem_p0_offset + p_cols_fp32 # 32 + 64 = 96 - s_cols = self.qk_mma_tiler[1] # 128 - o_after = max(s_cols, p_end) # 128 - self.tmem_o0_offset = ((o_after + 31) // 32) * 32 - self.tmem_vec_offset = 0 # Reuse S region for per-row inv_row_sum vector # align to 32 = 128 - self.tmem_vec_offset = 0 # Reuse S region (free after softmax loop) - o_cols = find_tmem_tensor_col_offset(tOtO) # footprint of O - total = self.tmem_o0_offset + o_cols - # Must be multiple of 32 AND power of 2 - self.num_tmem_alloc_cols = 1 - while self.num_tmem_alloc_cols < total: - self.num_tmem_alloc_cols *= 2 - cta = cute.size(qk_mma.thr_id.shape) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)) - self.q_tx_bytes = cute.size_in_bytes(self.q_dtype, q_s) * cta - self.kv_tx_bytes = cute.size_in_bytes(self.q_dtype, k_s) * cta - self.scale_softmax_log2 = Float32(1.0 / math.sqrt(HEAD_DIM) * math.log2(math.e)) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - # # s_k hardcoded # BROKEN in @cute.jit - # FMHA-style V: reconstruct as (HEAD_DIM, s_k, 1) MN-major - v_fmha = cute.make_tensor( - v.iterator, - cute.make_layout( - (HEAD_DIM, 128, 1), - stride=(1, HEAD_DIM, HEAD_DIM * 128), - ), - ) - self.v_major = LayoutEnum.from_tensor(v_fmha).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, self.a_major, self.b_major, self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, self.qk_acc_dtype, self.cta_group, (128,HEAD_DIM), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)); v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id),q,q_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id),k,k_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v_fmha,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape) - epi_s = cute.select(self.c_smem_s,mode=[0,1]) - tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile) - self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.v_smem_s,self.p_tmem_s,self.c_smem_s,self.epi_tile).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, q_smem_s, k_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx,_,_ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k); cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - q_bar: cute.struct.MemRange[cutlass.Int64, self.q_stage*2] - kv_bar: cute.struct.MemRange[cutlass.Int64, self.kv_stage*2] - s_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage*2] - tmem_dealloc: cutlass.Int64; holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - qp,qc = pipeline.PipelineTmaUmma.create(barrier_storage=st.q_bar.data_ptr(),num_stages=self.q_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.q_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - kvp,kvc = pipeline.PipelineTmaUmma.create(barrier_storage=st.kv_bar.data_ptr(),num_stages=self.kv_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.kv_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - s_prod,s_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.s_bar.data_ptr(),num_stages=1,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.epilogue_warp_id))).make_participants() - softmax_done_bar = pipeline.NamedBarrier(barrier_id=3, num_threads=32 + 32*len(self.epilogue_warp_id)) - pv_done_bar = pipeline.NamedBarrier(barrier_id=4, num_threads=32 + 32*len(self.epilogue_warp_id)) - acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(),num_stages=self.num_acc_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,len(self.epilogue_warp_id)),cta_layout_vmnk=cl_vmnk,defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2,num_threads=32*len((self.mma_warp_id,*self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr,barrier_for_retrieve=tmem_bar,allocator_warp_id=self.epilogue_warp_id[0],is_two_cta=cute.size(qk_mma.thr_id.shape)==2,two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk,is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype,layout=q_smem_s.outer,byte_alignment=128,swizzle=q_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype,layout=k_smem_s.outer,byte_alignment=128,swizzle=k_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype,layout=v_smem_s.outer,byte_alignment=128,swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype,layout=c_smem_s.outer,byte_alignment=128,swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None)) - gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None)) - gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None)) - gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None)) - n_kv_tiles = cute.size(gK, mode=[3]) - - qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - tCgV = pv_thr.partition_B(gV); tCgC = pv_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,0,None,0)).shape) - tAsQ,tAgQ = cpasync.tma_partition(tma_q,0,a_lay,cute.group_modes(sQ,0,3),cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape) - tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3)) - tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)]; tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - # --- PV read view (for MMA only, NOT for softmax store) --- - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None,None,None,0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # TMA LOAD - if warp_idx == self.tma_warp_id: - qp.reset(); qh = qp.acquire_and_advance() - cute.copy(tma_q,tAgQ[(None,qh.count)],tAsQ[(None,qh.index)],tma_bar_ptr=qh.barrier) - qp.tail() - kvp.reset(); pk = kvp.try_acquire() - for kt in cutlass.range(n_kv_tiles,unroll=1): - kh = kvp.acquire_and_advance(pk) - cute.copy(tma_k,tBgK[(None,kh.count)],tBsK[(None,kh.index)],tma_bar_ptr=kh.barrier) - pk = cutlass.Boolean(1) - vh = kvp.acquire_and_advance(pk) - cute.copy(tma_v,tVgV[(None,vh.count)],tVsV[(None,vh.index)],tma_bar_ptr=vh.barrier) - pk = cutlass.Boolean(1) - kvp.tail() - - # MMA - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - qc.reset(); qh = qc.wait_and_advance(); qh.release() - kvc.reset(); pk = kvc.try_wait() - acc_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_stage) - acc_pipe.producer_acquire(acc_st) - for kt in range(n_kv_tiles): - kh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - sh = s_prod.acquire_and_advance() - qk_mma.set(tcgen05.Field.ACCUMULATE, False) - for kb in cutlass.range(cute.size(tCrQ,mode=[2]), unroll_full=True): - cute.gemm(qk_mma, tStS0, tCrQ[(None,None,kb,0)], tCrK[(None,None,kb,kh.index)], tStS0) - qk_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - sh.commit(); kh.release() - softmax_done_bar.arrive_and_wait() - vh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - pv_mma.set(tcgen05.Field.ACCUMULATE, kt != 0) - for kb in cutlass.range(cute.size(tOrP0,mode=[2]), unroll_full=True): - cute.gemm(pv_mma, tOtO0, tOrP0[(None,None,kb)], tCrV[(None,None,kb,vh.index)], tOtO0) - pv_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - vh.release() - pv_done_bar.arrive() - acc_pipe.producer_commit(acc_st); acc_st.advance() - acc_pipe.producer_tail(acc_st) - - # ===================== EPILOGUE WARPS (STAGE C: ONLINE SOFTMAX) ===================== - if warp_idx < self.mma_warp_id: - tmem.allocate(self.num_tmem_alloc_cols) - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) - sfw_idx = tidx % (32 * len(self.epilogue_warp_id)) - - # --- S load (QK C-fragment) --- - tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS0) - thr_load = tiled_tmem_load.get_slice(sfw_idx) - tTMEM_LOADtS = thr_load.partition_S(tStS0) - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tTMEM_LOADcS = thr_load.partition_D(tScS) - - # --- P store (QK C-fragment composition, FMHA pattern) --- - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - tStP_layout = cute.composition(tStS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tStP0 = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStP_layout) - tmem_store_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStP0) - thr_store = tiled_tmem_store.get_slice(sfw_idx) - tTMEM_STOREtP = thr_store.partition_D(tStP0) - tScP_layout = cute.composition(tScS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tScP = cute.make_tensor(tScS.iterator, tScP_layout) - tTMEM_STOREcP = thr_store.partition_S(tScP) - - # --- Vector TMEM (per-row row_sum storage, FMHA pattern) --- - # composition(tStS.layout, (128, 2)) = 2 FP32 columns per logical row - # vec[0] = row_sum (final, after loop), vec[1] = unused - # Reuses S TMEM region (offset 0), free after softmax loop writes - - tStS_vec_layout = cute.composition(tStS.layout, cute.make_layout((128, 2))) - tStS_vec = cute.make_tensor(tStS.iterator + self.tmem_vec_offset, tStS_vec_layout) - tScS_vec_layout = cute.composition(tScS.layout, cute.make_layout((128, 2))) - tScS_vec = cute.make_tensor(tScS.iterator, tScS_vec_layout) - tmem_store_vec_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(2)), self.qk_acc_dtype) - tiled_tmem_store_vec = tcgen05.make_tmem_copy(tmem_store_vec_atom, tStS_vec) - thr_tmem_store_vec = tiled_tmem_store_vec.get_slice(sfw_idx) - tTMEM_STORE_VECtS = thr_tmem_store_vec.partition_D(tStS_vec) - tTMEM_STORE_VECcS = thr_tmem_store_vec.partition_S(tScS_vec) - tmem_load_vec_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(2)), self.qk_acc_dtype) - tiled_tmem_load_vec = tcgen05.make_tmem_copy(tmem_load_vec_atom, tStS_vec) - thr_tmem_load_vec = tiled_tmem_load_vec.get_slice(sfw_idx) - tTMEM_LOAD_VECtS = thr_tmem_load_vec.partition_S(tStS_vec) - tTMEM_LOAD_VECcS = thr_tmem_load_vec.partition_D(tScS_vec) - - # --- C6: O TMEM load/store for rescale (correction_rescale pattern) --- - corr_tile_size = 16 - cO = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1])) - tOcO = pv_thr.partition_C(cO) - o_tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), self.qk_acc_dtype) - o_tmem_store_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), self.qk_acc_dtype) - tOtO_i_layout = cute.composition(tOtO0.layout, cute.make_layout((128, corr_tile_size))) - tOcO_i_layout = cute.composition(tOcO.layout, cute.make_layout((128, corr_tile_size))) - tOtO_i = cute.make_tensor(tOtO0.iterator, tOtO_i_layout) - tOcO_i = cute.make_tensor(tOcO.iterator, tOcO_i_layout) - o_tiled_tmem_load = tcgen05.make_tmem_copy(o_tmem_load_atom, tOtO_i) - o_tiled_tmem_store = tcgen05.make_tmem_copy(o_tmem_store_atom, tOtO_i) - o_thr_load = o_tiled_tmem_load.get_slice(sfw_idx) - o_thr_store = o_tiled_tmem_store.get_slice(sfw_idx) - tTMEM_LOADtO = o_thr_load.partition_S(tOtO_i) - tTMEM_LOADcO = o_thr_load.partition_D(tOcO_i) - tTMEM_STOREtO = o_thr_store.partition_D(tOtO_i) - o_col_tiles = self.pv_mma_tiler[1] // corr_tile_size - - # --- C2: Per-thread row state (persist across KV tiles) --- - row_max = -cutlass.Float32.inf - row_sum = cutlass.Float32(0.0) - - # --- C3: QK scale = 1/sqrt(HEAD_DIM) * log2(e) for exp2 --- - scale = self.scale_softmax_log2 - - # ============================================================= - # Per-KV-tile online softmax loop - # ============================================================= - for kt in range(n_kv_tiles): - si_handle = s_cons.wait_and_advance() - - # Load S from TMEM (FP32, QK C-fragment layout) - tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS) - - # --- C4: Compute tile_max via .reduce(MAX) --- - old_row_max = row_max - row_max = tTMEM_LOADrS.load().reduce(cute.ReductionOp.MAX, row_max, 0) - row_max_safe = row_max - if row_max == -cutlass.Float32.inf: - row_max_safe = cutlass.Float32(0.0) - - # --- C5: Compute rescale factor --- - acc_scale = cute.math.exp2(scale * (old_row_max - row_max_safe), fastmath=True) - - # --- C6: Rescale O in TMEM (load O, multiply by acc_scale, store O) --- - # acc_scale belongs to QK row (N//4), but O rows are in PV partition (N). - # Store acc_scale to vector by QK row, read by PV row. - if kt > 0: - pv_done_bar.arrive_and_wait() - - # Store acc_scale to vector indexed by QK logical row - qk_row_c6 = tTMEM_LOADcS[0][0] - thr_vs_c6 = tiled_tmem_store_vec.get_slice(qk_row_c6) - tVStore_c6 = thr_vs_c6.partition_D(tStS_vec) - tVStoreSrc_c6 = thr_vs_c6.partition_S(tScS_vec) - tVStoreRmem_c6 = cute.make_rmem_tensor(tVStoreSrc_c6.shape, self.qk_acc_dtype) - tVStoreRmem_c6[0] = acc_scale - cute.copy(tiled_tmem_store_vec, tVStoreRmem_c6, tVStore_c6) - cute.arch.fence_view_async_tmem_store() - - # Read acc_scale from vector indexed by PV logical row - pv_row_c6 = tTMEM_LOADcO[0][0] - thr_vl_c6 = tiled_tmem_load_vec.get_slice(pv_row_c6) - tVLoad_c6 = thr_vl_c6.partition_S(tStS_vec) - tVLoadDst_c6 = thr_vl_c6.partition_D(tScS_vec) - tVLoadRmem_c6 = cute.make_rmem_tensor(tVLoadDst_c6.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load_vec, tVLoad_c6, tVLoadRmem_c6) - cute.arch.fence_view_async_tmem_load() - acc_scale_pv = tVLoadRmem_c6[0] - - tTMrO = cute.make_rmem_tensor((tTMEM_LOADcO.shape, o_col_tiles), self.qk_acc_dtype) - for i in range(o_col_tiles): - tTMrO_i_ = tTMrO[None, i] - tTMrO_i_layout = cute.composition(tTMrO_i_.layout, cute.make_layout(tTMrO.shape[0])) - tTMrO_i = cute.make_tensor(tTMrO_i_.iterator, tTMrO_i_layout) - tTMEM_LOADtO_i = cute.make_tensor(tTMEM_LOADtO.iterator + i * corr_tile_size, tTMEM_LOADtO.layout) - tTMEM_STOREtO_i = cute.make_tensor(tTMEM_STOREtO.iterator + i * corr_tile_size, tTMEM_STOREtO.layout) - cute.copy(o_tiled_tmem_load, tTMEM_LOADtO_i, tTMrO_i) - for j in cutlass.range(cute.size(tTMrO_i), vectorize=True): - tTMrO_i[j] = tTMrO_i[j] * acc_scale_pv - cute.copy(o_tiled_tmem_store, tTMrO_i, tTMEM_STOREtO_i) - cute.arch.fence_view_async_tmem_store() - - # Rescale row_sum - row_sum = row_sum * acc_scale - - # --- C7: Compute P = exp2((S - row_max_safe) * scale) --- - minus_row_max_scale = (cutlass.Float32(0.0) - row_max_safe) * scale - - # Register bridge (FMHA pattern: FP32 backing + BF16 view) - rP_words = cute.make_rmem_tensor(tTMEM_STOREcP.shape, self.qk_acc_dtype) - rP_bf16 = cute.make_tensor(cute.recast_ptr(rP_words.iterator, dtype=self.q_dtype), tTMEM_LOADrS.layout) - - frg_cnt = 4 - frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt - tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile)) - rP_bf16_frg = cute.logical_divide(rP_bf16, cute.make_layout(frg_tile)) - - # Scale S, compute exp2, store through register bridge - for j in range(frg_cnt): - for k in cutlass.range(cute.size(tTMEM_LOADrS_frg, mode=[0]), vectorize=True): - tTMEM_LOADrS_frg[k, j] = tTMEM_LOADrS_frg[k, j] * scale + minus_row_max_scale - tTMEM_LOADrS_frg[k, j] = cute.math.exp2(tTMEM_LOADrS_frg[k, j], fastmath=True) - s_vec = tTMEM_LOADrS_frg[None, j].load() - rP_bf16_frg[None, j].store(s_vec.to(self.q_dtype)) - - # Store P to TMEM - cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP) - cute.arch.fence_view_async_tmem_store() - si_handle.release() - softmax_done_bar.arrive() - - # --- C8: Row sum accumulation (CUTLASS FMHA packed f32x2 pattern) --- - # P values still in tTMEM_LOADrS registers. - # 4 accumulators for 4 reduction_unroll columns. - local_row_sum_0 = (cutlass.Float32(0.0), cutlass.Float32(0.0)) - local_row_sum_1 = (cutlass.Float32(0.0), cutlass.Float32(0.0)) - local_row_sum_2 = (cutlass.Float32(0.0), cutlass.Float32(0.0)) - local_row_sum_3 = (cutlass.Float32(0.0), cutlass.Float32(0.0)) - - reduction_unroll = 4 - rfrg_tile = cute.size(tTMEM_LOADrS) // reduction_unroll - tTMEM_LOADrS_rfrg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(rfrg_tile)) - - for j in cutlass.range_constexpr(0, cute.size(tTMEM_LOADrS_rfrg, mode=[0]), 2): - local_row_sum_0 = cute.arch.add_packed_f32x2( - local_row_sum_0, (tTMEM_LOADrS_rfrg[j, 0], tTMEM_LOADrS_rfrg[j + 1, 0])) - local_row_sum_1 = cute.arch.add_packed_f32x2( - local_row_sum_1, (tTMEM_LOADrS_rfrg[j, 1], tTMEM_LOADrS_rfrg[j + 1, 1])) - local_row_sum_2 = cute.arch.add_packed_f32x2( - local_row_sum_2, (tTMEM_LOADrS_rfrg[j, 2], tTMEM_LOADrS_rfrg[j + 1, 2])) - local_row_sum_3 = cute.arch.add_packed_f32x2( - local_row_sum_3, (tTMEM_LOADrS_rfrg[j, 3], tTMEM_LOADrS_rfrg[j + 1, 3])) - - local_row_sum_0 = cute.arch.add_packed_f32x2(local_row_sum_0, local_row_sum_1) - local_row_sum_2 = cute.arch.add_packed_f32x2(local_row_sum_2, local_row_sum_3) - local_row_sum_0 = cute.arch.add_packed_f32x2(local_row_sum_0, local_row_sum_2) - tile_sum = local_row_sum_0[0] + local_row_sum_0[1] - - row_sum = row_sum + tile_sum - - # --- C9: Final normalization via O TMEM rescale --- - pv_done_bar.arrive_and_wait() - - # DEBUG: hardcoded inv_row_sum = 1.0 (no normalization) - inv_row_sum = cutlass.Float32(1.0) - - # Normalize O in TMEM using PV-correct inv_row_sum - tTMrO_final = cute.make_rmem_tensor((tTMEM_LOADcO.shape, o_col_tiles), self.qk_acc_dtype) - for i in range(o_col_tiles): - tTMrO_i_ = tTMrO_final[None, i] - tTMrO_i_layout = cute.composition(tTMrO_i_.layout, cute.make_layout(tTMrO_final.shape[0])) - tTMrO_i = cute.make_tensor(tTMrO_i_.iterator, tTMrO_i_layout) - tTMEM_LOADtO_i = cute.make_tensor( - tTMEM_LOADtO.iterator + i * corr_tile_size, tTMEM_LOADtO.layout) - tTMEM_STOREtO_i = cute.make_tensor( - tTMEM_STOREtO.iterator + i * corr_tile_size, tTMEM_STOREtO.layout) - cute.copy(o_tiled_tmem_load, tTMEM_LOADtO_i, tTMrO_i) - for j in cutlass.range(cute.size(tTMrO_i), vectorize=True): - tTMrO_i[j] = tTMrO_i[j] * inv_row_sum - cute.copy(o_tiled_tmem_store, tTMrO_i, tTMEM_STOREtO_i) - cute.arch.fence_view_async_tmem_store() - - # Now O in TMEM is normalized. Use standard epilogue_tma_store with identity. - tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout) - acc_cons_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage) - c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) - c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp) - acc_cons_st = utils.gemm.sm100.epilogue_tma_store( - self, tidx, warp_idx, tma_c, tCtO_base, sC, tCgC, epi_tile, 0, - const_expr(lambda x: x), - (0,0,0), acc_cons_st, acc_pipe, c_pipe) - c_pipe.producer_tail() - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def test(): - import math - torch.manual_seed(42) - for n in [128, 256, 384]: - m, hd = 128, HEAD_DIM - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device="cuda") - k = torch.randn(n, hd, 1, dtype=torch.bfloat16, device="cuda") - v = torch.randn(n, hd, dtype=torch.bfloat16, device="cuda") - v_kernel = v.unsqueeze(-1) - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device="cuda") - qf = q[:,:,0].float(); kf = k[:,:,0].float() - attn = qf @ kf.T / math.sqrt(hd) - ref = torch.softmax(attn, dim=-1) @ v.float() - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - kernel = FmhaV3Softmax() - print(f"n={n}: Compiling...", flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - print(f"n={n}: tmem: s0={kernel.tmem_s0_offset} p0={kernel.tmem_p0_offset} o0={kernel.tmem_o0_offset} vec={kernel.tmem_vec_offset} alloc={kernel.num_tmem_alloc_cols}", flush=True) - print(f"n={n}: Running...", flush=True) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - out = c[:,:,0].float() - cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - max_err = (out - ref).abs().max().item() - print(f"FMHA softmax n={n}: cosine {cos:.6f} max_err {max_err:.6f} {'PASS' if cos >= 0.999 else 'FAIL'}", flush=True) - -if __name__ == "__main__": - test() - - diff --git a/tests/archive/unit_test_fmha_v3_per_row.py b/tests/archive/unit_test_fmha_v3_per_row.py deleted file mode 100644 index b7315966..00000000 --- a/tests/archive/unit_test_fmha_v3_per_row.py +++ /dev/null @@ -1,587 +0,0 @@ -""" -FMHA v3 + Stage C: QK -> online softmax -> PV with KV-tile interleaving. -Stage C: row_max, exp2, O rescale, row_sum, final normalization. -FMHA pattern P store preserved from Stage B. -""" -import math -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - -HEAD_DIM = 64 - -class FmhaV3Softmax: - def __init__(self, s_k: int = 128): - self.s_k = s_k - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0,1,2,3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.num_c_stage = 2 - self.kv_stage = 2; self.q_stage = 1; self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (128, 128, qk_ik * 4) - pv_ik = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (128, HEAD_DIM, pv_ik * (128 // pv_ik)) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0]//cute.size(qk_mma.thr_id.shape), HEAD_DIM, self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - self.q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, self.q_dtype, self.q_stage) - self.k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, self.q_dtype, self.kv_stage) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, self.kv_stage) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - qk_thr = qk_mma.get_slice(0); qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - pv_thr = pv_mma.get_slice(0); pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - self.tmem_s0_offset = 0; self.tmem_p0_offset = 32 - # P occupies [tmem_p0_offset, tmem_p0_offset + p_cols_fp32) - # S occupies [0, qk_mma_tiler[1]) = [0, 128) - # O must NOT overlap P. Place O after max(S end, P end), aligned to 32. - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - p_end = self.tmem_p0_offset + p_cols_fp32 # 32 + 64 = 96 - s_cols = self.qk_mma_tiler[1] # 128 - o_after = max(s_cols, p_end) # 128 - self.tmem_o0_offset = ((o_after + 31) // 32) * 32 - self.tmem_vec_offset = 0 # Reuse S region for per-row inv_row_sum vector # align to 32 = 128 - self.tmem_vec_offset = 0 # Reuse S region (free after softmax loop) - o_cols = find_tmem_tensor_col_offset(tOtO) # footprint of O - total = self.tmem_o0_offset + o_cols - # Must be multiple of 32 AND power of 2 - self.num_tmem_alloc_cols = 1 - while self.num_tmem_alloc_cols < total: - self.num_tmem_alloc_cols *= 2 - cta = cute.size(qk_mma.thr_id.shape) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)) - self.q_tx_bytes = cute.size_in_bytes(self.q_dtype, q_s) * cta - self.kv_tx_bytes = cute.size_in_bytes(self.q_dtype, k_s) * cta - self.scale_softmax_log2 = Float32(1.0 / math.sqrt(HEAD_DIM) * math.log2(math.e)) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - # # s_k hardcoded # BROKEN in @cute.jit - # FMHA-style V: reconstruct as (HEAD_DIM, s_k, 1) MN-major - v_fmha = cute.make_tensor( - v.iterator, - cute.make_layout( - (HEAD_DIM, self.s_k, 1), - stride=(1, HEAD_DIM, HEAD_DIM * self.s_k), - ), - ) - self.v_major = LayoutEnum.from_tensor(v_fmha).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, self.a_major, self.b_major, self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, self.qk_acc_dtype, self.cta_group, (128,HEAD_DIM), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)); v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id),q,q_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id),k,k_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v_fmha,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape) - epi_s = cute.select(self.c_smem_s,mode=[0,1]) - tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile) - self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.v_smem_s,self.p_tmem_s,self.c_smem_s,self.epi_tile).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, q_smem_s, k_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx,_,_ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k); cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - q_bar: cute.struct.MemRange[cutlass.Int64, self.q_stage*2] - kv_bar: cute.struct.MemRange[cutlass.Int64, self.kv_stage*2] - s_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage*2] - tmem_dealloc: cutlass.Int64; holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - qp,qc = pipeline.PipelineTmaUmma.create(barrier_storage=st.q_bar.data_ptr(),num_stages=self.q_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.q_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - kvp,kvc = pipeline.PipelineTmaUmma.create(barrier_storage=st.kv_bar.data_ptr(),num_stages=self.kv_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.kv_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - s_prod,s_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.s_bar.data_ptr(),num_stages=1,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.epilogue_warp_id))).make_participants() - softmax_done_bar = pipeline.NamedBarrier(barrier_id=3, num_threads=32 + 32*len(self.epilogue_warp_id)) - pv_done_bar = pipeline.NamedBarrier(barrier_id=4, num_threads=32 + 32*len(self.epilogue_warp_id)) - vec_handoff_bar = pipeline.NamedBarrier(barrier_id=5, num_threads=32*len(self.epilogue_warp_id)) - acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(),num_stages=self.num_acc_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.epilogue_warp_id)),cta_layout_vmnk=cl_vmnk,defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2,num_threads=32*len((self.mma_warp_id,*self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr,barrier_for_retrieve=tmem_bar,allocator_warp_id=self.epilogue_warp_id[0],is_two_cta=cute.size(qk_mma.thr_id.shape)==2,two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk,is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype,layout=q_smem_s.outer,byte_alignment=128,swizzle=q_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype,layout=k_smem_s.outer,byte_alignment=128,swizzle=k_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype,layout=v_smem_s.outer,byte_alignment=128,swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype,layout=c_smem_s.outer,byte_alignment=128,swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None)) - gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None)) - gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None)) - gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None)) - n_kv_tiles = cute.size(gK, mode=[3]) - - qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - tCgV = pv_thr.partition_B(gV); tCgC = pv_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,0,None,0)).shape) - tAsQ,tAgQ = cpasync.tma_partition(tma_q,0,a_lay,cute.group_modes(sQ,0,3),cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape) - tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3)) - tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)]; tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - # --- PV read view (for MMA only, NOT for softmax store) --- - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None,None,None,0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # TMA LOAD - if warp_idx == self.tma_warp_id: - qp.reset(); qh = qp.acquire_and_advance() - cute.copy(tma_q,tAgQ[(None,qh.count)],tAsQ[(None,qh.index)],tma_bar_ptr=qh.barrier) - qp.tail() - kvp.reset(); pk = kvp.try_acquire() - for kt in cutlass.range(n_kv_tiles,unroll=1): - kh = kvp.acquire_and_advance(pk) - cute.copy(tma_k,tBgK[(None,kh.count)],tBsK[(None,kh.index)],tma_bar_ptr=kh.barrier) - pk = cutlass.Boolean(1) - vh = kvp.acquire_and_advance(pk) - cute.copy(tma_v,tVgV[(None,vh.count)],tVsV[(None,vh.index)],tma_bar_ptr=vh.barrier) - pk = cutlass.Boolean(1) - kvp.tail() - - # MMA - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - qc.reset(); qh = qc.wait_and_advance(); qh.release() - kvc.reset(); pk = kvc.try_wait() - acc_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_stage) - acc_pipe.producer_acquire(acc_st) - for kt in range(n_kv_tiles): - kh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - sh = s_prod.acquire_and_advance() - qk_mma.set(tcgen05.Field.ACCUMULATE, False) - for kb in cutlass.range(cute.size(tCrQ,mode=[2]), unroll_full=True): - cute.gemm(qk_mma, tStS0, tCrQ[(None,None,kb,0)], tCrK[(None,None,kb,kh.index)], tStS0) - qk_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - sh.commit(); kh.release() - softmax_done_bar.arrive_and_wait() - vh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - pv_mma.set(tcgen05.Field.ACCUMULATE, kt != 0) - for kb in cutlass.range(cute.size(tOrP0,mode=[2]), unroll_full=True): - cute.gemm(pv_mma, tOtO0, tOrP0[(None,None,kb)], tCrV[(None,None,kb,vh.index)], tOtO0) - pv_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - vh.release() - pv_done_bar.arrive() - acc_pipe.producer_commit(acc_st); acc_st.advance() - acc_pipe.producer_tail(acc_st) - - # ===================== EPILOGUE WARPS (STAGE C: ONLINE SOFTMAX) ===================== - if warp_idx < self.mma_warp_id: - tmem.allocate(self.num_tmem_alloc_cols) - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) - sfw_idx = tidx % (32 * len(self.epilogue_warp_id)) - - # --- S load (QK C-fragment) --- - tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS0) - thr_load = tiled_tmem_load.get_slice(sfw_idx) - tTMEM_LOADtS = thr_load.partition_S(tStS0) - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tTMEM_LOADcS = thr_load.partition_D(tScS) - - # --- P store (QK C-fragment composition, FMHA pattern) --- - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - tStP_layout = cute.composition(tStS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tStP0 = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStP_layout) - tmem_store_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStP0) - thr_store = tiled_tmem_store.get_slice(sfw_idx) - tTMEM_STOREtP = thr_store.partition_D(tStP0) - tScP_layout = cute.composition(tScS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tScP = cute.make_tensor(tScS.iterator, tScP_layout) - tTMEM_STOREcP = thr_store.partition_S(tScP) - - # --- Vector TMEM (per-row row_sum storage, FMHA pattern) --- - # composition(tStS.layout, (128, 2)) = 2 FP32 columns per logical row - # vec[0] = row_sum (final, after loop), vec[1] = unused - # Reuses S TMEM region (offset 0), free after softmax loop writes - - tStS_vec_layout = cute.composition(tStS.layout, cute.make_layout((128, 2))) - tStS_vec = cute.make_tensor(tStS.iterator + self.tmem_vec_offset, tStS_vec_layout) - tScS_vec_layout = cute.composition(tScS.layout, cute.make_layout((128, 2))) - tScS_vec = cute.make_tensor(tScS.iterator, tScS_vec_layout) - tmem_store_vec_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(2)), self.qk_acc_dtype) - tiled_tmem_store_vec = tcgen05.make_tmem_copy(tmem_store_vec_atom, tStS_vec) - thr_tmem_store_vec = tiled_tmem_store_vec.get_slice(sfw_idx) - tTMEM_STORE_VECtS = thr_tmem_store_vec.partition_D(tStS_vec) - tTMEM_STORE_VECcS = thr_tmem_store_vec.partition_S(tScS_vec) - tmem_load_vec_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(2)), self.qk_acc_dtype) - tiled_tmem_load_vec = tcgen05.make_tmem_copy(tmem_load_vec_atom, tStS_vec) - thr_tmem_load_vec = tiled_tmem_load_vec.get_slice(sfw_idx) - tTMEM_LOAD_VECtS = thr_tmem_load_vec.partition_S(tStS_vec) - tTMEM_LOAD_VECcS = thr_tmem_load_vec.partition_D(tScS_vec) - - # --- C6: O TMEM load/store for rescale (correction_rescale pattern) --- - corr_tile_size = 16 - cO = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1])) - tOcO = pv_thr.partition_C(cO) - o_tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), self.qk_acc_dtype) - o_tmem_store_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), self.qk_acc_dtype) - tOtO_i_layout = cute.composition(tOtO0.layout, cute.make_layout((128, corr_tile_size))) - tOcO_i_layout = cute.composition(tOcO.layout, cute.make_layout((128, corr_tile_size))) - tOtO_i = cute.make_tensor(tOtO0.iterator, tOtO_i_layout) - tOcO_i = cute.make_tensor(tOcO.iterator, tOcO_i_layout) - o_tiled_tmem_load = tcgen05.make_tmem_copy(o_tmem_load_atom, tOtO_i) - o_tiled_tmem_store = tcgen05.make_tmem_copy(o_tmem_store_atom, tOtO_i) - o_thr_load = o_tiled_tmem_load.get_slice(sfw_idx) - o_thr_store = o_tiled_tmem_store.get_slice(sfw_idx) - tTMEM_LOADtO = o_thr_load.partition_S(tOtO_i) - tTMEM_LOADcO = o_thr_load.partition_D(tOcO_i) - tTMEM_STOREtO = o_thr_store.partition_D(tOtO_i) - o_col_tiles = self.pv_mma_tiler[1] // corr_tile_size - - # --- C2: Per-QK-fragment-row state (persist across KV tiles) --- - # The QK TMEM load fragment is logically 4 rows x 32 columns for each - # softmax thread. The old scalar row_max/row_sum reduced across all - # 4 rows and therefore produced a row_sum around 4.0. Keep one - # online-softmax state per local QK row. - qk_frg_cnt = 4 - qk_frg_tile = cute.size(tTMEM_LOADcS) // qk_frg_cnt - tTMEM_LOADcS_frg = cute.logical_divide(tTMEM_LOADcS, cute.make_layout(qk_frg_tile)) - - qk_row0 = tTMEM_LOADcS_frg[0, 0][0] - qk_row1 = tTMEM_LOADcS_frg[0, 1][0] - qk_row2 = tTMEM_LOADcS_frg[0, 2][0] - qk_row3 = tTMEM_LOADcS_frg[0, 3][0] - - row_max0 = -cutlass.Float32.inf - row_max1 = -cutlass.Float32.inf - row_max2 = -cutlass.Float32.inf - row_max3 = -cutlass.Float32.inf - - row_sum0 = cutlass.Float32(0.0) - row_sum1 = cutlass.Float32(0.0) - row_sum2 = cutlass.Float32(0.0) - row_sum3 = cutlass.Float32(0.0) - - # --- C3: QK scale = 1/sqrt(HEAD_DIM) * log2(e) for exp2 --- - scale = self.scale_softmax_log2 - - # ============================================================= - # Per-KV-tile online softmax loop - # ============================================================= - for kt in range(n_kv_tiles): - si_handle = s_cons.wait_and_advance() - - # Load S from TMEM (FP32, QK C-fragment layout). Because the - # vector buffer reuses the S columns, all softmax threads must - # finish this load before any thread writes vector data. - tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS) - cute.arch.fence_view_async_tmem_load() - vec_handoff_bar.arrive_and_wait() - - frg_cnt = 4 - frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt - tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile)) - - # --- C4: Compute tile_max independently for each local QK row --- - old_row_max0 = row_max0 - old_row_max1 = row_max1 - old_row_max2 = row_max2 - old_row_max3 = row_max3 - - row_max0 = tTMEM_LOADrS_frg[None, 0].load().reduce(cute.ReductionOp.MAX, row_max0, 0) - row_max1 = tTMEM_LOADrS_frg[None, 1].load().reduce(cute.ReductionOp.MAX, row_max1, 0) - row_max2 = tTMEM_LOADrS_frg[None, 2].load().reduce(cute.ReductionOp.MAX, row_max2, 0) - row_max3 = tTMEM_LOADrS_frg[None, 3].load().reduce(cute.ReductionOp.MAX, row_max3, 0) - - row_max0_safe = row_max0 - row_max1_safe = row_max1 - row_max2_safe = row_max2 - row_max3_safe = row_max3 - if row_max0 == -cutlass.Float32.inf: - row_max0_safe = cutlass.Float32(0.0) - if row_max1 == -cutlass.Float32.inf: - row_max1_safe = cutlass.Float32(0.0) - if row_max2 == -cutlass.Float32.inf: - row_max2_safe = cutlass.Float32(0.0) - if row_max3 == -cutlass.Float32.inf: - row_max3_safe = cutlass.Float32(0.0) - - # --- C5: Per-row O-rescale factors for the already-accumulated O --- - acc_scale0 = cute.math.exp2(scale * (old_row_max0 - row_max0_safe), fastmath=True) - acc_scale1 = cute.math.exp2(scale * (old_row_max1 - row_max1_safe), fastmath=True) - acc_scale2 = cute.math.exp2(scale * (old_row_max2 - row_max2_safe), fastmath=True) - acc_scale3 = cute.math.exp2(scale * (old_row_max3 - row_max3_safe), fastmath=True) - - # --- C6: Rescale O in TMEM using a row-indexed vector handoff --- - # Store per-QK-row acc_scale into vec[row, 0], then read vec[pv_row, 0] - # from the PV/O partition. This is the CUTLASS-style vector bridge, - # but folded into the same four softmax warps, so it needs an - # explicit warpgroup barrier between store and load. - if kt > 0: - pv_done_bar.arrive_and_wait() - - thr_vs0 = tiled_tmem_store_vec.get_slice(qk_row0) - tVStore0 = thr_vs0.partition_D(tStS_vec) - tVStoreSrc0 = thr_vs0.partition_S(tScS_vec) - rVec0 = cute.make_rmem_tensor(tVStoreSrc0.shape, self.qk_acc_dtype) - rVec0[0] = acc_scale0 - rVec0[1] = row_max0_safe - cute.copy(tiled_tmem_store_vec, rVec0, tVStore0) - - thr_vs1 = tiled_tmem_store_vec.get_slice(qk_row1) - tVStore1 = thr_vs1.partition_D(tStS_vec) - tVStoreSrc1 = thr_vs1.partition_S(tScS_vec) - rVec1 = cute.make_rmem_tensor(tVStoreSrc1.shape, self.qk_acc_dtype) - rVec1[0] = acc_scale1 - rVec1[1] = row_max1_safe - cute.copy(tiled_tmem_store_vec, rVec1, tVStore1) - - thr_vs2 = tiled_tmem_store_vec.get_slice(qk_row2) - tVStore2 = thr_vs2.partition_D(tStS_vec) - tVStoreSrc2 = thr_vs2.partition_S(tScS_vec) - rVec2 = cute.make_rmem_tensor(tVStoreSrc2.shape, self.qk_acc_dtype) - rVec2[0] = acc_scale2 - rVec2[1] = row_max2_safe - cute.copy(tiled_tmem_store_vec, rVec2, tVStore2) - - thr_vs3 = tiled_tmem_store_vec.get_slice(qk_row3) - tVStore3 = thr_vs3.partition_D(tStS_vec) - tVStoreSrc3 = thr_vs3.partition_S(tScS_vec) - rVec3 = cute.make_rmem_tensor(tVStoreSrc3.shape, self.qk_acc_dtype) - rVec3[0] = acc_scale3 - rVec3[1] = row_max3_safe - cute.copy(tiled_tmem_store_vec, rVec3, tVStore3) - - cute.arch.fence_view_async_tmem_store() - vec_handoff_bar.arrive_and_wait() - - pv_row = tTMEM_LOADcO[0][0] - thr_vl = tiled_tmem_load_vec.get_slice(pv_row) - tVLoad = thr_vl.partition_S(tStS_vec) - tVLoadDst = thr_vl.partition_D(tScS_vec) - rVecPV = cute.make_rmem_tensor(tVLoadDst.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load_vec, tVLoad, rVecPV) - cute.arch.fence_view_async_tmem_load() - acc_scale_pv = rVecPV[0] - - tTMrO = cute.make_rmem_tensor((tTMEM_LOADcO.shape, o_col_tiles), self.qk_acc_dtype) - for i in range(o_col_tiles): - tTMrO_i_ = tTMrO[None, i] - tTMrO_i_layout = cute.composition(tTMrO_i_.layout, cute.make_layout(tTMrO.shape[0])) - tTMrO_i = cute.make_tensor(tTMrO_i_.iterator, tTMrO_i_layout) - tTMEM_LOADtO_i = cute.make_tensor(tTMEM_LOADtO.iterator + i * corr_tile_size, tTMEM_LOADtO.layout) - tTMEM_STOREtO_i = cute.make_tensor(tTMEM_STOREtO.iterator + i * corr_tile_size, tTMEM_STOREtO.layout) - cute.copy(o_tiled_tmem_load, tTMEM_LOADtO_i, tTMrO_i) - for j in cutlass.range(cute.size(tTMrO_i), vectorize=True): - tTMrO_i[j] = tTMrO_i[j] * acc_scale_pv - cute.copy(o_tiled_tmem_store, tTMrO_i, tTMEM_STOREtO_i) - cute.arch.fence_view_async_tmem_store() - - # Rescale the four online row sums. - row_sum0 = row_sum0 * acc_scale0 - row_sum1 = row_sum1 * acc_scale1 - row_sum2 = row_sum2 * acc_scale2 - row_sum3 = row_sum3 * acc_scale3 - - # --- C7: Compute P = exp2((S - row_max[row]) * scale), per row --- - minus_row_max_scale0 = (cutlass.Float32(0.0) - row_max0_safe) * scale - minus_row_max_scale1 = (cutlass.Float32(0.0) - row_max1_safe) * scale - minus_row_max_scale2 = (cutlass.Float32(0.0) - row_max2_safe) * scale - minus_row_max_scale3 = (cutlass.Float32(0.0) - row_max3_safe) * scale - - rP_words = cute.make_rmem_tensor(tTMEM_STOREcP.shape, self.qk_acc_dtype) - rP_bf16 = cute.make_tensor(cute.recast_ptr(rP_words.iterator, dtype=self.q_dtype), tTMEM_LOADrS.layout) - rP_bf16_frg = cute.logical_divide(rP_bf16, cute.make_layout(frg_tile)) - - for k in cutlass.range(cute.size(tTMEM_LOADrS_frg, mode=[0]), vectorize=True): - tTMEM_LOADrS_frg[k, 0] = tTMEM_LOADrS_frg[k, 0] * scale + minus_row_max_scale0 - tTMEM_LOADrS_frg[k, 0] = cute.math.exp2(tTMEM_LOADrS_frg[k, 0], fastmath=True) - s_vec0 = tTMEM_LOADrS_frg[None, 0].load() - rP_bf16_frg[None, 0].store(s_vec0.to(self.q_dtype)) - - for k in cutlass.range(cute.size(tTMEM_LOADrS_frg, mode=[0]), vectorize=True): - tTMEM_LOADrS_frg[k, 1] = tTMEM_LOADrS_frg[k, 1] * scale + minus_row_max_scale1 - tTMEM_LOADrS_frg[k, 1] = cute.math.exp2(tTMEM_LOADrS_frg[k, 1], fastmath=True) - s_vec1 = tTMEM_LOADrS_frg[None, 1].load() - rP_bf16_frg[None, 1].store(s_vec1.to(self.q_dtype)) - - for k in cutlass.range(cute.size(tTMEM_LOADrS_frg, mode=[0]), vectorize=True): - tTMEM_LOADrS_frg[k, 2] = tTMEM_LOADrS_frg[k, 2] * scale + minus_row_max_scale2 - tTMEM_LOADrS_frg[k, 2] = cute.math.exp2(tTMEM_LOADrS_frg[k, 2], fastmath=True) - s_vec2 = tTMEM_LOADrS_frg[None, 2].load() - rP_bf16_frg[None, 2].store(s_vec2.to(self.q_dtype)) - - for k in cutlass.range(cute.size(tTMEM_LOADrS_frg, mode=[0]), vectorize=True): - tTMEM_LOADrS_frg[k, 3] = tTMEM_LOADrS_frg[k, 3] * scale + minus_row_max_scale3 - tTMEM_LOADrS_frg[k, 3] = cute.math.exp2(tTMEM_LOADrS_frg[k, 3], fastmath=True) - s_vec3 = tTMEM_LOADrS_frg[None, 3].load() - rP_bf16_frg[None, 3].store(s_vec3.to(self.q_dtype)) - - # Store P to TMEM. - cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP) - cute.arch.fence_view_async_tmem_store() - si_handle.release() - softmax_done_bar.arrive() - - # --- C8: Row sum accumulation, independently for each local QK row --- - tile_sum0 = tTMEM_LOADrS_frg[None, 0].load().reduce(cute.ReductionOp.ADD, cutlass.Float32(0.0), 0) - tile_sum1 = tTMEM_LOADrS_frg[None, 1].load().reduce(cute.ReductionOp.ADD, cutlass.Float32(0.0), 0) - tile_sum2 = tTMEM_LOADrS_frg[None, 2].load().reduce(cute.ReductionOp.ADD, cutlass.Float32(0.0), 0) - tile_sum3 = tTMEM_LOADrS_frg[None, 3].load().reduce(cute.ReductionOp.ADD, cutlass.Float32(0.0), 0) - - row_sum0 = row_sum0 + tile_sum0 - row_sum1 = row_sum1 + tile_sum1 - row_sum2 = row_sum2 + tile_sum2 - row_sum3 = row_sum3 + tile_sum3 - - # --- C9: Final normalization via row-indexed TMEM vector --- - # Wait for the final PV MMA to finish producing O. - pv_done_bar.arrive_and_wait() - - # Publish final row_sum per QK row into vec[row, 0]. - thr_vs0 = tiled_tmem_store_vec.get_slice(qk_row0) - tVStore0 = thr_vs0.partition_D(tStS_vec) - tVStoreSrc0 = thr_vs0.partition_S(tScS_vec) - rVec0 = cute.make_rmem_tensor(tVStoreSrc0.shape, self.qk_acc_dtype) - rVec0[0] = row_sum0 - rVec0[1] = row_max0 - cute.copy(tiled_tmem_store_vec, rVec0, tVStore0) - - thr_vs1 = tiled_tmem_store_vec.get_slice(qk_row1) - tVStore1 = thr_vs1.partition_D(tStS_vec) - tVStoreSrc1 = thr_vs1.partition_S(tScS_vec) - rVec1 = cute.make_rmem_tensor(tVStoreSrc1.shape, self.qk_acc_dtype) - rVec1[0] = row_sum1 - rVec1[1] = row_max1 - cute.copy(tiled_tmem_store_vec, rVec1, tVStore1) - - thr_vs2 = tiled_tmem_store_vec.get_slice(qk_row2) - tVStore2 = thr_vs2.partition_D(tStS_vec) - tVStoreSrc2 = thr_vs2.partition_S(tScS_vec) - rVec2 = cute.make_rmem_tensor(tVStoreSrc2.shape, self.qk_acc_dtype) - rVec2[0] = row_sum2 - rVec2[1] = row_max2 - cute.copy(tiled_tmem_store_vec, rVec2, tVStore2) - - thr_vs3 = tiled_tmem_store_vec.get_slice(qk_row3) - tVStore3 = thr_vs3.partition_D(tStS_vec) - tVStoreSrc3 = thr_vs3.partition_S(tScS_vec) - rVec3 = cute.make_rmem_tensor(tVStoreSrc3.shape, self.qk_acc_dtype) - rVec3[0] = row_sum3 - rVec3[1] = row_max3 - cute.copy(tiled_tmem_store_vec, rVec3, tVStore3) - - cute.arch.fence_view_async_tmem_store() - vec_handoff_bar.arrive_and_wait() - - # Read the correct row_sum for this PV/O row and normalize O. - pv_row_final = tTMEM_LOADcO[0][0] - thr_vl_final = tiled_tmem_load_vec.get_slice(pv_row_final) - tVLoadFinal = thr_vl_final.partition_S(tStS_vec) - tVLoadFinalDst = thr_vl_final.partition_D(tScS_vec) - rVecFinal = cute.make_rmem_tensor(tVLoadFinalDst.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load_vec, tVLoadFinal, rVecFinal) - cute.arch.fence_view_async_tmem_load() - - inv_row_sum = cutlass.Float32(1.0) / rVecFinal[0] - - tTMrO_final = cute.make_rmem_tensor((tTMEM_LOADcO.shape, o_col_tiles), self.qk_acc_dtype) - for i in range(o_col_tiles): - tTMrO_i_ = tTMrO_final[None, i] - tTMrO_i_layout = cute.composition(tTMrO_i_.layout, cute.make_layout(tTMrO_final.shape[0])) - tTMrO_i = cute.make_tensor(tTMrO_i_.iterator, tTMrO_i_layout) - tTMEM_LOADtO_i = cute.make_tensor( - tTMEM_LOADtO.iterator + i * corr_tile_size, tTMEM_LOADtO.layout) - tTMEM_STOREtO_i = cute.make_tensor( - tTMEM_STOREtO.iterator + i * corr_tile_size, tTMEM_STOREtO.layout) - cute.copy(o_tiled_tmem_load, tTMEM_LOADtO_i, tTMrO_i) - for j in cutlass.range(cute.size(tTMrO_i), vectorize=True): - tTMrO_i[j] = tTMrO_i[j] * inv_row_sum - cute.copy(o_tiled_tmem_store, tTMrO_i, tTMEM_STOREtO_i) - cute.arch.fence_view_async_tmem_store() - - # Now O in TMEM is normalized. Use standard epilogue_tma_store with identity. - tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout) - acc_cons_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage) - c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) - c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp) - acc_cons_st = utils.gemm.sm100.epilogue_tma_store( - self, tidx, warp_idx, tma_c, tCtO_base, sC, tCgC, epi_tile, 0, - const_expr(lambda x: x), - (0,0,0), acc_cons_st, acc_pipe, c_pipe) - c_pipe.producer_tail() - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def test(): - import math - torch.manual_seed(42) - for n in [128, 256, 384]: - m, hd = 128, HEAD_DIM - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device="cuda") - k = torch.randn(n, hd, 1, dtype=torch.bfloat16, device="cuda") - v = torch.randn(n, hd, dtype=torch.bfloat16, device="cuda") - v_kernel = v.unsqueeze(-1) - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device="cuda") - qf = q[:,:,0].float(); kf = k[:,:,0].float() - attn = qf @ kf.T / math.sqrt(hd) - ref = torch.softmax(attn, dim=-1) @ v.float() - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - kernel = FmhaV3Softmax(s_k=n) - print(f"n={n}: Compiling...", flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - print(f"n={n}: tmem: s0={kernel.tmem_s0_offset} p0={kernel.tmem_p0_offset} o0={kernel.tmem_o0_offset} vec={kernel.tmem_vec_offset} alloc={kernel.num_tmem_alloc_cols}", flush=True) - print(f"n={n}: Running...", flush=True) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - out = c[:,:,0].float() - cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - max_err = (out - ref).abs().max().item() - print(f"FMHA softmax n={n}: cosine {cos:.6f} max_err {max_err:.6f} {'PASS' if cos >= 0.999 else 'FAIL'}", flush=True) - -if __name__ == "__main__": - test() - diff --git a/tests/archive/unit_test_fmha_v3_per_row_min.py b/tests/archive/unit_test_fmha_v3_per_row_min.py deleted file mode 100644 index 168260b7..00000000 --- a/tests/archive/unit_test_fmha_v3_per_row_min.py +++ /dev/null @@ -1,465 +0,0 @@ -""" -FMHA v3 + Stage C: QK -> online softmax -> PV with KV-tile interleaving. -Stage C: row_max, exp2, O rescale, row_sum, final normalization. -FMHA pattern P store preserved from Stage B. -""" -import math -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - -HEAD_DIM = 64 - -class FmhaV3Softmax: - def __init__(self, s_k: int = 128): - self.s_k = s_k - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0,1,2,3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.num_c_stage = 2 - self.kv_stage = 2; self.q_stage = 1; self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (128, 128, qk_ik * 4) - pv_ik = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (128, HEAD_DIM, pv_ik * (128 // pv_ik)) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0]//cute.size(qk_mma.thr_id.shape), HEAD_DIM, self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - self.q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, self.q_dtype, self.q_stage) - self.k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, self.q_dtype, self.kv_stage) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, self.kv_stage) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - qk_thr = qk_mma.get_slice(0); qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - pv_thr = pv_mma.get_slice(0); pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - self.tmem_s0_offset = 0; self.tmem_p0_offset = 32 - # P occupies [tmem_p0_offset, tmem_p0_offset + p_cols_fp32) - # S occupies [0, qk_mma_tiler[1]) = [0, 128) - # O must NOT overlap P. Place O after max(S end, P end), aligned to 32. - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - p_end = self.tmem_p0_offset + p_cols_fp32 # 32 + 64 = 96 - s_cols = self.qk_mma_tiler[1] # 128 - o_after = max(s_cols, p_end) # 128 - self.tmem_o0_offset = ((o_after + 31) // 32) * 32 - self.tmem_vec_offset = 0 # Reuse S region for per-row inv_row_sum vector # align to 32 = 128 - self.tmem_vec_offset = 0 # Reuse S region (free after softmax loop) - o_cols = find_tmem_tensor_col_offset(tOtO) # footprint of O - total = self.tmem_o0_offset + o_cols - # Must be multiple of 32 AND power of 2 - self.num_tmem_alloc_cols = 1 - while self.num_tmem_alloc_cols < total: - self.num_tmem_alloc_cols *= 2 - cta = cute.size(qk_mma.thr_id.shape) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)) - self.q_tx_bytes = cute.size_in_bytes(self.q_dtype, q_s) * cta - self.kv_tx_bytes = cute.size_in_bytes(self.q_dtype, k_s) * cta - self.scale_softmax_log2 = Float32(1.0 / math.sqrt(HEAD_DIM) * math.log2(math.e)) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - # # s_k hardcoded # BROKEN in @cute.jit - # FMHA-style V: reconstruct as (HEAD_DIM, s_k, 1) MN-major - v_fmha = cute.make_tensor( - v.iterator, - cute.make_layout( - (HEAD_DIM, self.s_k, 1), - stride=(1, HEAD_DIM, HEAD_DIM * self.s_k), - ), - ) - self.v_major = LayoutEnum.from_tensor(v_fmha).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, self.a_major, self.b_major, self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, self.qk_acc_dtype, self.cta_group, (128,HEAD_DIM), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)); v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id),q,q_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id),k,k_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v_fmha,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape) - epi_s = cute.select(self.c_smem_s,mode=[0,1]) - tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile) - self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.v_smem_s,self.p_tmem_s,self.c_smem_s,self.epi_tile).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, q_smem_s, k_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx,_,_ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k); cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - q_bar: cute.struct.MemRange[cutlass.Int64, self.q_stage*2] - kv_bar: cute.struct.MemRange[cutlass.Int64, self.kv_stage*2] - s_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage*2] - tmem_dealloc: cutlass.Int64; holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - qp,qc = pipeline.PipelineTmaUmma.create(barrier_storage=st.q_bar.data_ptr(),num_stages=self.q_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.q_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - kvp,kvc = pipeline.PipelineTmaUmma.create(barrier_storage=st.kv_bar.data_ptr(),num_stages=self.kv_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.kv_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - s_prod,s_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.s_bar.data_ptr(),num_stages=1,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.epilogue_warp_id))).make_participants() - softmax_done_bar = pipeline.NamedBarrier(barrier_id=3, num_threads=32 + 32*len(self.epilogue_warp_id)) - pv_done_bar = pipeline.NamedBarrier(barrier_id=4, num_threads=32 + 32*len(self.epilogue_warp_id)) - vec_handoff_bar = pipeline.NamedBarrier(barrier_id=5, num_threads=32*len(self.epilogue_warp_id)) - acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(),num_stages=self.num_acc_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.epilogue_warp_id)),cta_layout_vmnk=cl_vmnk,defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2,num_threads=32*len((self.mma_warp_id,*self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr,barrier_for_retrieve=tmem_bar,allocator_warp_id=self.epilogue_warp_id[0],is_two_cta=cute.size(qk_mma.thr_id.shape)==2,two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk,is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype,layout=q_smem_s.outer,byte_alignment=128,swizzle=q_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype,layout=k_smem_s.outer,byte_alignment=128,swizzle=k_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype,layout=v_smem_s.outer,byte_alignment=128,swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype,layout=c_smem_s.outer,byte_alignment=128,swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None)) - gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None)) - gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None)) - gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None)) - n_kv_tiles = cute.size(gK, mode=[3]) - - qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - tCgV = pv_thr.partition_B(gV); tCgC = pv_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,0,None,0)).shape) - tAsQ,tAgQ = cpasync.tma_partition(tma_q,0,a_lay,cute.group_modes(sQ,0,3),cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape) - tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3)) - tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)]; tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - # --- PV read view (for MMA only, NOT for softmax store) --- - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None,None,None,0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # TMA LOAD - if warp_idx == self.tma_warp_id: - qp.reset(); qh = qp.acquire_and_advance() - cute.copy(tma_q,tAgQ[(None,qh.count)],tAsQ[(None,qh.index)],tma_bar_ptr=qh.barrier) - qp.tail() - kvp.reset(); pk = kvp.try_acquire() - for kt in cutlass.range(n_kv_tiles,unroll=1): - kh = kvp.acquire_and_advance(pk) - cute.copy(tma_k,tBgK[(None,kh.count)],tBsK[(None,kh.index)],tma_bar_ptr=kh.barrier) - pk = cutlass.Boolean(1) - vh = kvp.acquire_and_advance(pk) - cute.copy(tma_v,tVgV[(None,vh.count)],tVsV[(None,vh.index)],tma_bar_ptr=vh.barrier) - pk = cutlass.Boolean(1) - kvp.tail() - - # MMA - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - qc.reset(); qh = qc.wait_and_advance(); qh.release() - kvc.reset(); pk = kvc.try_wait() - acc_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_stage) - acc_pipe.producer_acquire(acc_st) - for kt in range(n_kv_tiles): - kh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - sh = s_prod.acquire_and_advance() - qk_mma.set(tcgen05.Field.ACCUMULATE, False) - for kb in cutlass.range(cute.size(tCrQ,mode=[2]), unroll_full=True): - cute.gemm(qk_mma, tStS0, tCrQ[(None,None,kb,0)], tCrK[(None,None,kb,kh.index)], tStS0) - qk_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - sh.commit(); kh.release() - # softmax_done_bar skipped - vh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - pv_mma.set(tcgen05.Field.ACCUMULATE, kt != 0) - for kb in cutlass.range(cute.size(tOrP0,mode=[2]), unroll_full=True): - cute.gemm(pv_mma, tOtO0, tOrP0[(None,None,kb)], tCrV[(None,None,kb,vh.index)], tOtO0) - pv_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - vh.release() - # pv_done_bar skipped - acc_pipe.producer_commit(acc_st); acc_st.advance() - acc_pipe.producer_tail(acc_st) - - # ===================== EPILOGUE WARPS (STAGE C: ONLINE SOFTMAX) ===================== - if warp_idx < self.mma_warp_id: - tmem.allocate(self.num_tmem_alloc_cols) - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) - sfw_idx = tidx % (32 * len(self.epilogue_warp_id)) - - # --- S load (QK C-fragment) --- - tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS0) - thr_load = tiled_tmem_load.get_slice(sfw_idx) - tTMEM_LOADtS = thr_load.partition_S(tStS0) - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tTMEM_LOADcS = thr_load.partition_D(tScS) - - # --- P store (QK C-fragment composition, FMHA pattern) --- - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - tStP_layout = cute.composition(tStS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tStP0 = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStP_layout) - tmem_store_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStP0) - thr_store = tiled_tmem_store.get_slice(sfw_idx) - tTMEM_STOREtP = thr_store.partition_D(tStP0) - tScP_layout = cute.composition(tScS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tScP = cute.make_tensor(tScS.iterator, tScP_layout) - tTMEM_STOREcP = thr_store.partition_S(tScP) - - # --- Vector TMEM (per-row row_sum storage, FMHA pattern) --- - # composition(tStS.layout, (128, 2)) = 2 FP32 columns per logical row - # vec[0] = row_sum (final, after loop), vec[1] = unused - # Reuses S TMEM region (offset 0), free after softmax loop writes - - tStS_vec_layout = cute.composition(tStS.layout, cute.make_layout((128, 2))) - tStS_vec = cute.make_tensor(tStS.iterator + self.tmem_vec_offset, tStS_vec_layout) - tScS_vec_layout = cute.composition(tScS.layout, cute.make_layout((128, 2))) - tScS_vec = cute.make_tensor(tScS.iterator, tScS_vec_layout) - tmem_store_vec_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(2)), self.qk_acc_dtype) - tiled_tmem_store_vec = tcgen05.make_tmem_copy(tmem_store_vec_atom, tStS_vec) - thr_tmem_store_vec = tiled_tmem_store_vec.get_slice(sfw_idx) - tTMEM_STORE_VECtS = thr_tmem_store_vec.partition_D(tStS_vec) - tTMEM_STORE_VECcS = thr_tmem_store_vec.partition_S(tScS_vec) - tmem_load_vec_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(2)), self.qk_acc_dtype) - tiled_tmem_load_vec = tcgen05.make_tmem_copy(tmem_load_vec_atom, tStS_vec) - thr_tmem_load_vec = tiled_tmem_load_vec.get_slice(sfw_idx) - tTMEM_LOAD_VECtS = thr_tmem_load_vec.partition_S(tStS_vec) - tTMEM_LOAD_VECcS = thr_tmem_load_vec.partition_D(tScS_vec) - - # --- C6: O TMEM load/store for rescale (correction_rescale pattern) --- - corr_tile_size = 16 - cO = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1])) - tOcO = pv_thr.partition_C(cO) - o_tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), self.qk_acc_dtype) - o_tmem_store_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), self.qk_acc_dtype) - tOtO_i_layout = cute.composition(tOtO0.layout, cute.make_layout((128, corr_tile_size))) - tOcO_i_layout = cute.composition(tOcO.layout, cute.make_layout((128, corr_tile_size))) - tOtO_i = cute.make_tensor(tOtO0.iterator, tOtO_i_layout) - tOcO_i = cute.make_tensor(tOcO.iterator, tOcO_i_layout) - o_tiled_tmem_load = tcgen05.make_tmem_copy(o_tmem_load_atom, tOtO_i) - o_tiled_tmem_store = tcgen05.make_tmem_copy(o_tmem_store_atom, tOtO_i) - o_thr_load = o_tiled_tmem_load.get_slice(sfw_idx) - o_thr_store = o_tiled_tmem_store.get_slice(sfw_idx) - tTMEM_LOADtO = o_thr_load.partition_S(tOtO_i) - tTMEM_LOADcO = o_thr_load.partition_D(tOcO_i) - tTMEM_STOREtO = o_thr_store.partition_D(tOtO_i) - o_col_tiles = self.pv_mma_tiler[1] // corr_tile_size - - # --- C2: Per-QK-fragment-row state (persist across KV tiles) --- - # The QK TMEM load fragment is logically 4 rows x 32 columns for each - # softmax thread. The old scalar row_max/row_sum reduced across all - # 4 rows and therefore produced a row_sum around 4.0. Keep one - # online-softmax state per local QK row. - qk_frg_cnt = 4 - qk_frg_tile = cute.size(tTMEM_LOADcS) // qk_frg_cnt - tTMEM_LOADcS_frg = cute.logical_divide(tTMEM_LOADcS, cute.make_layout(qk_frg_tile)) - - qk_row0 = tTMEM_LOADcS_frg[0, 0][0] - qk_row1 = tTMEM_LOADcS_frg[0, 1][0] - qk_row2 = tTMEM_LOADcS_frg[0, 2][0] - qk_row3 = tTMEM_LOADcS_frg[0, 3][0] - - row_max0 = -cutlass.Float32.inf - row_max1 = -cutlass.Float32.inf - row_max2 = -cutlass.Float32.inf - row_max3 = -cutlass.Float32.inf - - row_sum0 = cutlass.Float32(0.0) - row_sum1 = cutlass.Float32(0.0) - row_sum2 = cutlass.Float32(0.0) - row_sum3 = cutlass.Float32(0.0) - - # --- C3: QK scale = 1/sqrt(HEAD_DIM) * log2(e) for exp2 --- - scale = self.scale_softmax_log2 - - # ============================================================= - # Per-KV-tile online softmax loop - # ============================================================= - for kt in range(n_kv_tiles): - si_handle = s_cons.wait_and_advance() - - # Load S from TMEM (FP32, QK C-fragment layout). Because the - # vector buffer reuses the S columns, all softmax threads must - # finish this load before any thread writes vector data. - tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS) - cute.arch.fence_view_async_tmem_load() - # vec_handoff_bar skipped - - frg_cnt = 4 - frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt - tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile)) - - # --- C4: Compute tile_max independently for each local QK row --- - old_row_max0 = row_max0 - old_row_max1 = row_max1 - old_row_max2 = row_max2 - old_row_max3 = row_max3 - - # Per-row max: explicit loop (avoid .load().reduce on sliced tensor) - for k in cutlass.range(cute.size(tTMEM_LOADrS_frg, mode=[0])): - row_max0 = cutlass.Float32.max(row_max0, tTMEM_LOADrS_frg[k, 0]) - row_max1 = cutlass.Float32.max(row_max1, tTMEM_LOADrS_frg[k, 1]) - row_max2 = cutlass.Float32.max(row_max2, tTMEM_LOADrS_frg[k, 2]) - row_max3 = cutlass.Float32.max(row_max3, tTMEM_LOADrS_frg[k, 3]) - - row_max0_safe = row_max0 - row_max1_safe = row_max1 - row_max2_safe = row_max2 - row_max3_safe = row_max3 - if row_max0 == -cutlass.Float32.inf: - row_max0_safe = cutlass.Float32(0.0) - if row_max1 == -cutlass.Float32.inf: - row_max1_safe = cutlass.Float32(0.0) - if row_max2 == -cutlass.Float32.inf: - row_max2_safe = cutlass.Float32(0.0) - if row_max3 == -cutlass.Float32.inf: - row_max3_safe = cutlass.Float32(0.0) - - # --- C5: Per-row O-rescale factors for the already-accumulated O --- - acc_scale0 = cute.math.exp2(scale * (old_row_max0 - row_max0_safe), fastmath=True) - acc_scale1 = cute.math.exp2(scale * (old_row_max1 - row_max1_safe), fastmath=True) - acc_scale2 = cute.math.exp2(scale * (old_row_max2 - row_max2_safe), fastmath=True) - acc_scale3 = cute.math.exp2(scale * (old_row_max3 - row_max3_safe), fastmath=True) - - # --- C6: SKIPPED (minimal test) --- - if kt > 0: - pass # pv_done_bar skipped - else: - pass # pv_done_bar skipped - - # --- C7: Compute P = exp2((S - row_max[row]) * scale), per row --- - minus_row_max_scale0 = (cutlass.Float32(0.0) - row_max0_safe) * scale - minus_row_max_scale1 = (cutlass.Float32(0.0) - row_max1_safe) * scale - minus_row_max_scale2 = (cutlass.Float32(0.0) - row_max2_safe) * scale - minus_row_max_scale3 = (cutlass.Float32(0.0) - row_max3_safe) * scale - - rP_words = cute.make_rmem_tensor(tTMEM_STOREcP.shape, self.qk_acc_dtype) - rP_bf16 = cute.make_tensor(cute.recast_ptr(rP_words.iterator, dtype=self.q_dtype), tTMEM_LOADrS.layout) - rP_bf16_frg = cute.logical_divide(rP_bf16, cute.make_layout(frg_tile)) - - for k in cutlass.range(cute.size(tTMEM_LOADrS_frg, mode=[0]), vectorize=True): - tTMEM_LOADrS_frg[k, 0] = tTMEM_LOADrS_frg[k, 0] * scale + minus_row_max_scale0 - tTMEM_LOADrS_frg[k, 0] = cute.math.exp2(tTMEM_LOADrS_frg[k, 0], fastmath=True) - s_vec0 = tTMEM_LOADrS_frg[None, 0].load() - rP_bf16_frg[None, 0].store(s_vec0.to(self.q_dtype)) - - for k in cutlass.range(cute.size(tTMEM_LOADrS_frg, mode=[0]), vectorize=True): - tTMEM_LOADrS_frg[k, 1] = tTMEM_LOADrS_frg[k, 1] * scale + minus_row_max_scale1 - tTMEM_LOADrS_frg[k, 1] = cute.math.exp2(tTMEM_LOADrS_frg[k, 1], fastmath=True) - s_vec1 = tTMEM_LOADrS_frg[None, 1].load() - rP_bf16_frg[None, 1].store(s_vec1.to(self.q_dtype)) - - for k in cutlass.range(cute.size(tTMEM_LOADrS_frg, mode=[0]), vectorize=True): - tTMEM_LOADrS_frg[k, 2] = tTMEM_LOADrS_frg[k, 2] * scale + minus_row_max_scale2 - tTMEM_LOADrS_frg[k, 2] = cute.math.exp2(tTMEM_LOADrS_frg[k, 2], fastmath=True) - s_vec2 = tTMEM_LOADrS_frg[None, 2].load() - rP_bf16_frg[None, 2].store(s_vec2.to(self.q_dtype)) - - for k in cutlass.range(cute.size(tTMEM_LOADrS_frg, mode=[0]), vectorize=True): - tTMEM_LOADrS_frg[k, 3] = tTMEM_LOADrS_frg[k, 3] * scale + minus_row_max_scale3 - tTMEM_LOADrS_frg[k, 3] = cute.math.exp2(tTMEM_LOADrS_frg[k, 3], fastmath=True) - s_vec3 = tTMEM_LOADrS_frg[None, 3].load() - rP_bf16_frg[None, 3].store(s_vec3.to(self.q_dtype)) - - # Store P to TMEM. - cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP) - cute.arch.fence_view_async_tmem_store() - si_handle.release() - # softmax_done_bar skipped - - # --- C8: Row sum accumulation, independently for each local QK row --- - # Per-row sum: explicit loop - tile_sum0 = cutlass.Float32(0.0) - tile_sum1 = cutlass.Float32(0.0) - tile_sum2 = cutlass.Float32(0.0) - tile_sum3 = cutlass.Float32(0.0) - for k in cutlass.range(cute.size(tTMEM_LOADrS_frg, mode=[0])): - tile_sum0 = tile_sum0 + tTMEM_LOADrS_frg[k, 0] - tile_sum1 = tile_sum1 + tTMEM_LOADrS_frg[k, 1] - tile_sum2 = tile_sum2 + tTMEM_LOADrS_frg[k, 2] - tile_sum3 = tile_sum3 + tTMEM_LOADrS_frg[k, 3] - - row_sum0 = row_sum0 + tile_sum0 - row_sum1 = row_sum1 + tile_sum1 - row_sum2 = row_sum2 + tile_sum2 - row_sum3 = row_sum3 + tile_sum3 - - # --- C9: SKIPPED (minimal test) --- - # pv_done_bar skipped - tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout) - acc_cons_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage) - c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) - c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp) - acc_cons_st = utils.gemm.sm100.epilogue_tma_store( - self, tidx, warp_idx, tma_c, tCtO_base, sC, tCgC, epi_tile, 0, - const_expr(lambda x: x), - (0,0,0), acc_cons_st, acc_pipe, c_pipe) - c_pipe.producer_tail() - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def test(): - import math - torch.manual_seed(42) - for n in [128, 256, 384]: - m, hd = 128, HEAD_DIM - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device="cuda") - k = torch.randn(n, hd, 1, dtype=torch.bfloat16, device="cuda") - v = torch.randn(n, hd, dtype=torch.bfloat16, device="cuda") - v_kernel = v.unsqueeze(-1) - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device="cuda") - qf = q[:,:,0].float(); kf = k[:,:,0].float() - attn = qf @ kf.T / math.sqrt(hd) - ref = torch.softmax(attn, dim=-1) @ v.float() - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - kernel = FmhaV3Softmax(s_k=n) - print(f"n={n}: Compiling...", flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - print(f"n={n}: tmem: s0={kernel.tmem_s0_offset} p0={kernel.tmem_p0_offset} o0={kernel.tmem_o0_offset} vec={kernel.tmem_vec_offset} alloc={kernel.num_tmem_alloc_cols}", flush=True) - print(f"n={n}: Running...", flush=True) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - out = c[:,:,0].float() - cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - max_err = (out - ref).abs().max().item() - print(f"FMHA softmax n={n}: cosine {cos:.6f} max_err {max_err:.6f} {'PASS' if cos >= 0.999 else 'FAIL'}", flush=True) - -if __name__ == "__main__": - test() - diff --git a/tests/archive/unit_test_fmha_v3_proper.py b/tests/archive/unit_test_fmha_v3_proper.py deleted file mode 100644 index 3db346e2..00000000 --- a/tests/archive/unit_test_fmha_v3_proper.py +++ /dev/null @@ -1,416 +0,0 @@ -""" -FMHA v3 Proper: 11-warp with correction warp group + epilogue warp. -Warp layout: softmax(0-3), correction(4-7), MMA(8), TMA(9), epilogue(10) -""" -import math, torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda, cutlass.torch as ct - -HEAD_DIM = 64 - -class FmhaV3Proper: - def __init__(self): - self.qk_acc_dtype = Float32; self.pv_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = False; self.cluster_shape_mn = (1, 1) - self.cta_group = tcgen05.CtaGroup.ONE - self.softmax_warp_ids = (0,1,2,3) - self.correction_warp_ids = (4,5,6,7) - self.mma_warp_id = 8; self.tma_warp_id = 9; self.epilogue_warp_id = 10 - self.threads_per_cta = 352 - self.q_stage = 1; self.kv_stage = 2; self.num_acc_stage = 1 - self.mma_softmax_stage = 1; self.softmax_corr_stage = 1 - self.mma_corr_stage = 2; self.epi_stage = 2; self.num_c_stage = 2 - self.scale_softmax_log2 = Float32(1.0 / math.sqrt(HEAD_DIM) * math.log2(math.e)) - - def _setup(self, qk_mma, pv_mma): - qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (128, 128, qk_ik * 4) - pv_ik = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (128, HEAD_DIM, pv_ik * (128 // pv_ik)) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0]//cute.size(qk_mma.thr_id.shape), HEAD_DIM, self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1 - self.q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, self.q_dtype, self.q_stage) - self.k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, self.q_dtype, self.kv_stage) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, self.kv_stage) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - qk_thr = qk_mma.get_slice(0); qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - pv_thr = pv_mma.get_slice(0); pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - self.tmem_s0_offset = 0; self.tmem_p0_offset = 32; self.tmem_vec0_offset = 0 - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - o_after = max(self.qk_mma_tiler[1], self.tmem_p0_offset + p_cols_fp32) - self.tmem_o0_offset = ((o_after + 31) // 32) * 32 - o_cols = find_tmem_tensor_col_offset(tOtO) - total = self.tmem_o0_offset + o_cols - self.num_tmem_alloc_cols = 1 - while self.num_tmem_alloc_cols < total: self.num_tmem_alloc_cols *= 2 - cta = cute.size(qk_mma.thr_id.shape) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)) - self.q_tx_bytes = cute.size_in_bytes(self.q_dtype, q_s) * cta - self.kv_tx_bytes = cute.size_in_bytes(self.q_dtype, k_s) * cta - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - v_fmha = cute.make_tensor(v.iterator, cute.make_layout((HEAD_DIM, 128, 1), stride=(1, HEAD_DIM, HEAD_DIM * 128))) - self.v_major = LayoutEnum.from_tensor(v_fmha).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, self.a_major, self.b_major, self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, self.pv_acc_dtype, self.cta_group, (128,HEAD_DIM), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)); v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id),q,q_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id),k,k_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v_fmha,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape) - epi_s = cute.select(self.c_smem_s,mode=[0,1]) - tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile) - self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.v_smem_s,self.p_tmem_s,self.c_smem_s,self.epi_tile).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, q_smem_s, k_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx,_,_ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k); cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - q_bar: cute.struct.MemRange[cutlass.Int64, self.q_stage*2] - kv_bar: cute.struct.MemRange[cutlass.Int64, self.kv_stage*2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, self.mma_softmax_stage*2] - si_corr_bar: cute.struct.MemRange[cutlass.Int64, self.softmax_corr_stage*2] - mma_corr_bar: cute.struct.MemRange[cutlass.Int64, self.mma_corr_stage*2] - corr_epi_bar: cute.struct.MemRange[cutlass.Int64, self.epi_stage*2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage*2] - tmem_dealloc: cutlass.Int64; holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - qp,qc = pipeline.PipelineTmaUmma.create(barrier_storage=st.q_bar.data_ptr(),num_stages=self.q_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.q_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - kvp,kvc = pipeline.PipelineTmaUmma.create(barrier_storage=st.kv_bar.data_ptr(),num_stages=self.kv_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.kv_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - mma_si_prod,mma_si_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.mma_si_bar.data_ptr(),num_stages=self.mma_softmax_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.softmax_warp_ids))).make_participants() - si_corr_prod,si_corr_cons = pipeline.PipelineAsync.create(barrier_storage=st.si_corr_bar.data_ptr(),num_stages=self.softmax_corr_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.softmax_warp_ids)),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.correction_warp_ids))).make_participants() - mma_corr_prod,mma_corr_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.mma_corr_bar.data_ptr(),num_stages=self.mma_corr_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.correction_warp_ids))).make_participants() - corr_epi_prod,corr_epi_cons = pipeline.PipelineAsync.create(barrier_storage=st.corr_epi_bar.data_ptr(),num_stages=self.epi_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.correction_warp_ids)),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32)).make_participants() - acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(),num_stages=self.num_acc_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,len(self.correction_warp_ids)),cta_layout_vmnk=cl_vmnk,defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2,num_threads=32*len((self.mma_warp_id,*self.softmax_warp_ids))) - tmem = utils.TmemAllocator(st.holding.ptr,barrier_for_retrieve=tmem_bar,allocator_warp_id=self.softmax_warp_ids[0],is_two_cta=cute.size(qk_mma.thr_id.shape)==2,two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk,is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype,layout=q_smem_s.outer,byte_alignment=128,swizzle=q_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype,layout=k_smem_s.outer,byte_alignment=128,swizzle=k_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype,layout=v_smem_s.outer,byte_alignment=128,swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype,layout=c_smem_s.outer,byte_alignment=128,swizzle=c_smem_s.inner) - gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None)) - gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None)) - gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None)) - gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None)) - n_kv_tiles = cute.size(gK, mode=[3]) - qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgV = pv_thr.partition_B(gV); tCgC = pv_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,0,None,0)).shape) - tAsQ,tAgQ = cpasync.tma_partition(tma_q,0,a_lay,cute.group_modes(sQ,0,3),cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape) - tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3)) - tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)]; tVgV = tVgV[(None,0,None,0)] - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK); tCrV = pv_mma.make_fragment_B(sV) - qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]); tStS = qk_thr.make_fragment_C(qk_as) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]); tOtO = pv_thr.make_fragment_C(pv_as) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP); tOrP = tOrP_base[(None,None,None,0)] - tOrP0 = cute.make_tensor(tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, tOrP.layout) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # TMA - if warp_idx == self.tma_warp_id: - qp.reset(); qh = qp.acquire_and_advance() - cute.copy(tma_q,tAgQ[(None,qh.count)],tAsQ[(None,qh.index)],tma_bar_ptr=qh.barrier); qp.tail() - kvp.reset(); pk = kvp.try_acquire() - for kt in cutlass.range(n_kv_tiles,unroll=1): - kh = kvp.acquire_and_advance(pk); cute.copy(tma_k,tBgK[(None,kh.count)],tBsK[(None,kh.index)],tma_bar_ptr=kh.barrier); pk = cutlass.Boolean(1) - vh = kvp.acquire_and_advance(pk); cute.copy(tma_v,tVgV[(None,vh.count)],tVsV[(None,vh.index)],tma_bar_ptr=vh.barrier); pk = cutlass.Boolean(1) - kvp.tail() - - # MMA - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - qc.reset(); qh = qc.wait_and_advance(); qh.release() - kvc.reset(); pk = kvc.try_wait() - acc_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_stage) - acc_pipe.producer_acquire(acc_st) - for kt in range(n_kv_tiles): - kh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - sh = mma_si_prod.acquire_and_advance() - qk_mma.set(tcgen05.Field.ACCUMULATE, False) - for kb in cutlass.range(cute.size(tCrQ,mode=[2]), unroll_full=True): - cute.gemm(qk_mma, tStS0, tCrQ[(None,None,kb,0)], tCrK[(None,None,kb,kh.index)], tStS0) - qk_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store(); sh.commit(); kh.release() - if kt > 0: - o_handle = mma_corr_cons.wait_and_advance(); o_handle.release() - sh2 = mma_si_prod.acquire_and_advance() - vh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - pv_mma.set(tcgen05.Field.ACCUMULATE, kt != 0) - for kb in cutlass.range(cute.size(tOrP0,mode=[2]), unroll_full=True): - cute.gemm(pv_mma, tOtO0, tOrP0[(None,None,kb)], tCrV[(None,None,kb,vh.index)], tOtO0) - pv_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store(); vh.release() - o_prod_h = mma_corr_prod.acquire_and_advance(); o_prod_h.commit() - o_handle = mma_corr_cons.wait_and_advance(); o_handle.release() - acc_pipe.producer_commit(acc_st); acc_st.advance(); acc_pipe.producer_tail(acc_st) - - # SOFTMAX (warps 0-3) - if warp_idx < len(self.softmax_warp_ids): - tmem.allocate(self.num_tmem_alloc_cols); tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) - sfw_idx = tidx % (32 * len(self.softmax_warp_ids)) - scale = self.scale_softmax_log2 - tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS0) - thr_load = tiled_tmem_load.get_slice(sfw_idx) - tTMEM_LOADtS = thr_load.partition_S(tStS0) - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS); tTMEM_LOADcS = thr_load.partition_D(tScS) - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - tStP_layout = cute.composition(tStS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tStP0 = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStP_layout) - tmem_store_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStP0) - thr_store = tiled_tmem_store.get_slice(sfw_idx) - tTMEM_STOREtP = thr_store.partition_D(tStP0) - tScP_layout = cute.composition(tScS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tScP = cute.make_tensor(tScS.iterator, tScP_layout); tTMEM_STOREcP = thr_store.partition_S(tScP) - tStS_vec_layout = cute.composition(tStS.layout, cute.make_layout((128, 2))) - tStS_vec = cute.make_tensor(tStS.iterator + self.tmem_vec0_offset, tStS_vec_layout) - tScS_vec_layout = cute.composition(tScS.layout, cute.make_layout((128, 2))) - tScS_vec = cute.make_tensor(tScS.iterator, tScS_vec_layout) - tmem_store_vec_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(2)), self.qk_acc_dtype) - tiled_tmem_store_vec = tcgen05.make_tmem_copy(tmem_store_vec_atom, tStS_vec) - thr_tmem_store_vec = tiled_tmem_store_vec.get_slice(sfw_idx) - tTMEM_STORE_VECtS = thr_tmem_store_vec.partition_D(tStS_vec) - tTMEM_STORE_VECcS = thr_tmem_store_vec.partition_S(tScS_vec) - row_max = -cutlass.Float32.inf; row_sum = cutlass.Float32(0.0) - for kt in range(n_kv_tiles): - si_handle = mma_si_cons.wait_and_advance() - tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS) - old_row_max = row_max - row_max = tTMEM_LOADrS.load().reduce(cute.ReductionOp.MAX, row_max, 0) - row_max_safe = row_max - if row_max == -cutlass.Float32.inf: row_max_safe = cutlass.Float32(0.0) - vec_handle = si_corr_prod.acquire_and_advance() - tTMEM_STORE_VECrS = cute.make_rmem_tensor(tTMEM_STORE_VECcS.shape, self.qk_acc_dtype) - tTMEM_STORE_VECrS[0] = old_row_max; tTMEM_STORE_VECrS[1] = row_max_safe - cute.copy(tiled_tmem_store_vec, tTMEM_STORE_VECrS, tTMEM_STORE_VECtS) - cute.arch.fence_view_async_tmem_store(); vec_handle.commit() - minus_row_max_scale = (cutlass.Float32(0.0) - row_max_safe) * scale - rP_words = cute.make_rmem_tensor(tTMEM_STOREcP.shape, self.qk_acc_dtype) - rP_bf16 = cute.make_tensor(cute.recast_ptr(rP_words.iterator, dtype=self.q_dtype), tTMEM_LOADrS.layout) - frg_cnt = 4; frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt - tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile)) - rP_bf16_frg = cute.logical_divide(rP_bf16, cute.make_layout(frg_tile)) - for j in range(frg_cnt): - for k in cutlass.range(cute.size(tTMEM_LOADrS_frg, mode=[0]), vectorize=True): - tTMEM_LOADrS_frg[k, j] = tTMEM_LOADrS_frg[k, j] * scale + minus_row_max_scale - tTMEM_LOADrS_frg[k, j] = cute.math.exp2(tTMEM_LOADrS_frg[k, j], fastmath=True) - s_vec = tTMEM_LOADrS_frg[None, j].load(); rP_bf16_frg[None, j].store(s_vec.to(self.q_dtype)) - cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP) - cute.arch.fence_view_async_tmem_store(); si_handle.release() - acc_scale = cute.math.exp2(scale * (old_row_max - row_max_safe), fastmath=True) - row_sum = row_sum * acc_scale - local_row_sum_0 = (cutlass.Float32(0.0), cutlass.Float32(0.0)) - local_row_sum_1 = (cutlass.Float32(0.0), cutlass.Float32(0.0)) - local_row_sum_2 = (cutlass.Float32(0.0), cutlass.Float32(0.0)) - local_row_sum_3 = (cutlass.Float32(0.0), cutlass.Float32(0.0)) - reduction_unroll = 4; rfrg_tile = cute.size(tTMEM_LOADrS) // reduction_unroll - tTMEM_LOADrS_rfrg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(rfrg_tile)) - for j in cutlass.range_constexpr(0, cute.size(tTMEM_LOADrS_rfrg, mode=[0]), 2): - local_row_sum_0 = cute.arch.add_packed_f32x2(local_row_sum_0, (tTMEM_LOADrS_rfrg[j, 0], tTMEM_LOADrS_rfrg[j+1, 0])) - local_row_sum_1 = cute.arch.add_packed_f32x2(local_row_sum_1, (tTMEM_LOADrS_rfrg[j, 1], tTMEM_LOADrS_rfrg[j+1, 1])) - local_row_sum_2 = cute.arch.add_packed_f32x2(local_row_sum_2, (tTMEM_LOADrS_rfrg[j, 2], tTMEM_LOADrS_rfrg[j+1, 2])) - local_row_sum_3 = cute.arch.add_packed_f32x2(local_row_sum_3, (tTMEM_LOADrS_rfrg[j, 3], tTMEM_LOADrS_rfrg[j+1, 3])) - local_row_sum_0 = cute.arch.add_packed_f32x2(local_row_sum_0, local_row_sum_1) - local_row_sum_2 = cute.arch.add_packed_f32x2(local_row_sum_2, local_row_sum_3) - local_row_sum_0 = cute.arch.add_packed_f32x2(local_row_sum_0, local_row_sum_2) - row_sum = row_sum + local_row_sum_0[0] + local_row_sum_0[1] - # Final vector: (row_sum, row_max) - vec_handle = si_corr_prod.acquire_and_advance() - tTMEM_STORE_VECrS = cute.make_rmem_tensor(tTMEM_STORE_VECcS.shape, self.qk_acc_dtype) - tTMEM_STORE_VECrS[0] = row_sum; tTMEM_STORE_VECrS[1] = row_max - cute.copy(tiled_tmem_store_vec, tTMEM_STORE_VECrS, tTMEM_STORE_VECtS) - cute.arch.fence_view_async_tmem_store(); vec_handle.commit() - si_handle = mma_si_cons.wait_and_advance(); si_corr_prod.acquire(); si_handle.release() - tmem.relinquish_alloc_permit() - - # CORRECTION (warps 4-7) - if warp_idx >= len(self.softmax_warp_ids) and warp_idx < len(self.softmax_warp_ids) + len(self.correction_warp_ids): - corr_idx = tidx % (32 * len(self.correction_warp_ids)) - scale = self.scale_softmax_log2 - # Create tScS from common-scope qk_thr (same as softmax section) - cS_corr = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS_corr) - tStS_vec_layout = cute.composition(tStS.layout, cute.make_layout((128, 2))) - tStS_vec = cute.make_tensor(tStS.iterator + self.tmem_vec0_offset, tStS_vec_layout) - tScS_vec_layout = cute.composition(tScS.layout, cute.make_layout((128, 2))) - tScS_vec = cute.make_tensor(tScS.iterator, tScS_vec_layout) - tmem_load_vec_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(2)), self.qk_acc_dtype) - tiled_tmem_load_vec = tcgen05.make_tmem_copy(tmem_load_vec_atom, tStS_vec) - thr_tmem_load_vec = tiled_tmem_load_vec.get_slice(corr_idx) - tTMEM_LOAD_VECtS = thr_tmem_load_vec.partition_S(tStS_vec) - tTMEM_LOAD_VECcS = thr_tmem_load_vec.partition_D(tScS_vec) - corr_tile_size = 16 - cO = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1])) - tOcO = pv_thr.partition_C(cO) - tOtO_i_layout = cute.composition(tOtO0.layout, cute.make_layout((128, corr_tile_size))) - tOcO_i_layout = cute.composition(tOcO.layout, cute.make_layout((128, corr_tile_size))) - tOtO_i = cute.make_tensor(tOtO0.iterator, tOtO_i_layout) - tOcO_i = cute.make_tensor(tOcO.iterator, tOcO_i_layout) - o_tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), self.pv_acc_dtype) - o_tmem_store_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), self.pv_acc_dtype) - o_tiled_tmem_load = tcgen05.make_tmem_copy(o_tmem_load_atom, tOtO_i) - o_tiled_tmem_store = tcgen05.make_tmem_copy(o_tmem_store_atom, tOtO_i) - o_thr_load = o_tiled_tmem_load.get_slice(corr_idx) - o_thr_store = o_tiled_tmem_store.get_slice(corr_idx) - tTMEM_LOADtO = o_thr_load.partition_S(tOtO_i) - tTMEM_LOADcO = o_thr_load.partition_D(tOcO_i) - tTMEM_STOREtO = o_thr_store.partition_D(tOtO_i) - o_col_tiles = self.pv_mma_tiler[1] // corr_tile_size - - # Ignore first vec (no rescale for first PV) - vec_handle = si_corr_cons.wait_and_advance() - vec_handle.release() - - for kt in range(n_kv_tiles): - if kt > 0: - # Wait for vector (old_max, new_max) from softmax - vec_handle = si_corr_cons.wait_and_advance() - tTMEM_LOAD_VECrS = cute.make_rmem_tensor(tTMEM_LOAD_VECcS.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load_vec, tTMEM_LOAD_VECtS, tTMEM_LOAD_VECrS) - corr_scale_ = scale * (tTMEM_LOAD_VECrS[0] - tTMEM_LOAD_VECrS[1]) - corr_scale = cute.math.exp2(corr_scale_, fastmath=True) - - # Wait for O from MMA - o_handle = mma_corr_cons.wait_and_advance() - - # correction_rescale: O *= corr_scale in TMEM - tTMrO = cute.make_rmem_tensor((tTMEM_LOADcO.shape, o_col_tiles), self.pv_acc_dtype) - for i in range(o_col_tiles): - tTMrO_i_ = tTMrO[None, i] - tTMrO_i_layout = cute.composition(tTMrO_i_.layout, cute.make_layout(tTMrO.shape[0])) - tTMrO_i = cute.make_tensor(tTMrO_i_.iterator, tTMrO_i_layout) - tTMEM_LOADtO_i = cute.make_tensor(tTMEM_LOADtO.iterator + i * corr_tile_size, tTMEM_LOADtO.layout) - tTMEM_STOREtO_i = cute.make_tensor(tTMEM_STOREtO.iterator + i * corr_tile_size, tTMEM_STOREtO.layout) - cute.copy(o_tiled_tmem_load, tTMEM_LOADtO_i, tTMrO_i) - for j in cutlass.range(cute.size(tTMrO_i), vectorize=True): - tTMrO_i[j] = tTMrO_i[j] * corr_scale - cute.copy(o_tiled_tmem_store, tTMrO_i, tTMEM_STOREtO_i) - cute.arch.fence_view_async_tmem_store() - vec_handle.release() - o_handle.release() - - # --- correction_epilog: final normalize O /= row_sum --- - # Wait for final vector (row_sum, row_max) from softmax - vec_handle = si_corr_cons.wait_and_advance() - tTMEM_LOAD_VECrS = cute.make_rmem_tensor(tTMEM_LOAD_VECcS.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load_vec, tTMEM_LOAD_VECtS, tTMEM_LOAD_VECrS) - cute.arch.fence_view_async_tmem_load() - vec_handle.release() - inv_row_sum = cutlass.Float32(1.0) / tTMEM_LOAD_VECrS[0] - - # Wait for final O from MMA - o_handle = mma_corr_cons.wait_and_advance() - epi_handle = corr_epi_prod.acquire_and_advance() - - # Final normalize O in TMEM - tTMrO = cute.make_rmem_tensor((tTMEM_LOADcO.shape, o_col_tiles), self.pv_acc_dtype) - for i in range(o_col_tiles): - tTMrO_i_ = tTMrO[None, i] - tTMrO_i_layout = cute.composition(tTMrO_i_.layout, cute.make_layout(tTMrO.shape[0])) - tTMrO_i = cute.make_tensor(tTMrO_i_.iterator, tTMrO_i_layout) - tTMEM_LOADtO_i = cute.make_tensor(tTMEM_LOADtO.iterator + i * corr_tile_size, tTMEM_LOADtO.layout) - tTMEM_STOREtO_i = cute.make_tensor(tTMEM_STOREtO.iterator + i * corr_tile_size, tTMEM_STOREtO.layout) - cute.copy(o_tiled_tmem_load, tTMEM_LOADtO_i, tTMrO_i) - for j in cutlass.range(cute.size(tTMrO_i), vectorize=True): - tTMrO_i[j] = tTMrO_i[j] * inv_row_sum - cute.copy(o_tiled_tmem_store, tTMrO_i, tTMEM_STOREtO_i) - cute.arch.fence_view_async_tmem_store() - o_handle.release() - epi_handle.commit() - - # --- EPILOGUE WARP (warp 10) - TMA store O --- - # After correction normalizes O in TMEM, the epilogue reads O from TMEM, - # writes to SMEM, then TMA stores from SMEM to GMEM. - # For now, the softmax warps (which have tmem_ptr) handle the TMA store - # after correction signals completion. This matches our working 6-warp code's - # epilogue_tma_store pattern. - # The epilogue warp (warp 10) just waits for the signal and does TMA store. - # Since it doesn't have tmem_ptr, we need a different approach. - # Simplest: let the softmax warps also do the TMA store after correction - # signals O is ready. But softmax warps already exited... - # - # Alternative: the epilogue warp uses acc_pipe + epilogue_tma_store - # which reads from TMEM directly. - # For initial test: skip epilogue TMA store, just verify correction works. - # Then add TMA store via a separate mechanism. - # - # Actually, looking at our working 6-warp code, the epilogue_tma_store - # reads from tCtO_base which is a TMEM tensor at tmem_ptr + offset. - # The epilogue warp doesn't have tmem_ptr. BUT it can create the same - # tensor if it knows the address. The MMA warp has it from alloc_tmem. - # - # For the initial version, let the softmax warps do TMA store - # (they have tmem_ptr) after waiting for correction to finish. - # This is a temporary simplification. - - if warp_idx == self.epilogue_warp_id: - epi_handle = corr_epi_cons.wait_and_advance() - epi_handle.release() - - -def test(): - import math - torch.manual_seed(42) - for n in [128, 256, 384]: - m, hd = 128, HEAD_DIM - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device="cuda") - k = torch.randn(n, hd, 1, dtype=torch.bfloat16, device="cuda") - v = torch.randn(n, hd, dtype=torch.bfloat16, device="cuda") - v_kernel = v.unsqueeze(-1) - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device="cuda") - qf = q[:,:,0].float(); kf = k[:,:,0].float() - attn = qf @ kf.T / math.sqrt(hd) - ref = torch.softmax(attn, dim=-1) @ v.float() - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - kernel = FmhaV3Proper() - print(f"n={n}: Compiling...", flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - print(f"n={n}: tmem: s0={kernel.tmem_s0_offset} p0={kernel.tmem_p0_offset} o0={kernel.tmem_o0_offset} vec={kernel.tmem_vec0_offset} alloc={kernel.num_tmem_alloc_cols}", flush=True) - print(f"n={n}: Running...", flush=True) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - out = c[:,:,0].float() - cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - max_err = (out - ref).abs().max().item() - print(f"FMHA proper n={n}: cosine {cos:.6f} max_err {max_err:.6f} {'PASS' if cos >= 0.999 else 'FAIL'}", flush=True) - -if __name__ == "__main__": - test() diff --git a/tests/archive/unit_test_fmha_v3_pva_c9.py b/tests/archive/unit_test_fmha_v3_pva_c9.py deleted file mode 100644 index c4fbcd65..00000000 --- a/tests/archive/unit_test_fmha_v3_pva_c9.py +++ /dev/null @@ -1,484 +0,0 @@ -""" -FMHA v3 + Stage C: QK -> online softmax -> PV with KV-tile interleaving. -Stage C: row_max, exp2, O rescale, row_sum, final normalization. -FMHA pattern P store preserved from Stage B. -""" -import math -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - -HEAD_DIM = 64 - -class FmhaV3Softmax: - def __init__(self): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0,1,2,3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.num_c_stage = 2 - self.kv_stage = 2; self.q_stage = 1; self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (128, 128, qk_ik * 4) - pv_ik = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (128, HEAD_DIM, pv_ik * (128 // pv_ik)) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0]//cute.size(qk_mma.thr_id.shape), HEAD_DIM, self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - self.q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, self.q_dtype, self.q_stage) - self.k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, self.q_dtype, self.kv_stage) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, self.kv_stage) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - qk_thr = qk_mma.get_slice(0); qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - pv_thr = pv_mma.get_slice(0); pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - self.tmem_s0_offset = 0; self.tmem_p0_offset = 32 - # P occupies [tmem_p0_offset, tmem_p0_offset + p_cols_fp32) - # S occupies [0, qk_mma_tiler[1]) = [0, 128) - # O must NOT overlap P. Place O after max(S end, P end), aligned to 32. - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - p_end = self.tmem_p0_offset + p_cols_fp32 # 32 + 64 = 96 - s_cols = self.qk_mma_tiler[1] # 128 - o_after = max(s_cols, p_end) # 128 - self.tmem_o0_offset = ((o_after + 31) // 32) * 32 - self.tmem_vec_offset = 0 # Reuse S region for per-row inv_row_sum vector # align to 32 = 128 - self.tmem_vec_offset = 0 # Reuse S region (free after softmax loop) - o_cols = find_tmem_tensor_col_offset(tOtO) # footprint of O - total = self.tmem_o0_offset + o_cols - # Must be multiple of 32 AND power of 2 - self.num_tmem_alloc_cols = 1 - while self.num_tmem_alloc_cols < total: - self.num_tmem_alloc_cols *= 2 - cta = cute.size(qk_mma.thr_id.shape) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)) - self.q_tx_bytes = cute.size_in_bytes(self.q_dtype, q_s) * cta - self.kv_tx_bytes = cute.size_in_bytes(self.q_dtype, k_s) * cta - self.scale_softmax_log2 = Float32(1.0 / math.sqrt(HEAD_DIM) * math.log2(math.e)) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - # # s_k hardcoded # BROKEN in @cute.jit - # FMHA-style V: reconstruct as (HEAD_DIM, s_k, 1) MN-major - v_fmha = cute.make_tensor( - v.iterator, - cute.make_layout( - (HEAD_DIM, 128, 1), - stride=(1, HEAD_DIM, HEAD_DIM * 128), - ), - ) - self.v_major = LayoutEnum.from_tensor(v_fmha).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, self.a_major, self.b_major, self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, self.qk_acc_dtype, self.cta_group, (128,HEAD_DIM), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)); v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id),q,q_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id),k,k_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v_fmha,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape) - epi_s = cute.select(self.c_smem_s,mode=[0,1]) - tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile) - self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.v_smem_s,self.p_tmem_s,self.c_smem_s,self.epi_tile).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, q_smem_s, k_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx,_,_ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k); cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - q_bar: cute.struct.MemRange[cutlass.Int64, self.q_stage*2] - kv_bar: cute.struct.MemRange[cutlass.Int64, self.kv_stage*2] - s_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage*2] - tmem_dealloc: cutlass.Int64; holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - qp,qc = pipeline.PipelineTmaUmma.create(barrier_storage=st.q_bar.data_ptr(),num_stages=self.q_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.q_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - kvp,kvc = pipeline.PipelineTmaUmma.create(barrier_storage=st.kv_bar.data_ptr(),num_stages=self.kv_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.kv_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - s_prod,s_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.s_bar.data_ptr(),num_stages=1,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.epilogue_warp_id))).make_participants() - softmax_done_bar = pipeline.NamedBarrier(barrier_id=3, num_threads=32 + 32*len(self.epilogue_warp_id)) - pv_done_bar = pipeline.NamedBarrier(barrier_id=4, num_threads=32 + 32*len(self.epilogue_warp_id)) - acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(),num_stages=self.num_acc_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,len(self.epilogue_warp_id)),cta_layout_vmnk=cl_vmnk,defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2,num_threads=32*len((self.mma_warp_id,*self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr,barrier_for_retrieve=tmem_bar,allocator_warp_id=self.epilogue_warp_id[0],is_two_cta=cute.size(qk_mma.thr_id.shape)==2,two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk,is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype,layout=q_smem_s.outer,byte_alignment=128,swizzle=q_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype,layout=k_smem_s.outer,byte_alignment=128,swizzle=k_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype,layout=v_smem_s.outer,byte_alignment=128,swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype,layout=c_smem_s.outer,byte_alignment=128,swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None)) - gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None)) - gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None)) - gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None)) - n_kv_tiles = cute.size(gK, mode=[3]) - - qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - tCgV = pv_thr.partition_B(gV); tCgC = pv_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,0,None,0)).shape) - tAsQ,tAgQ = cpasync.tma_partition(tma_q,0,a_lay,cute.group_modes(sQ,0,3),cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape) - tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3)) - tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)]; tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - # --- PV read view (for MMA only, NOT for softmax store) --- - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None,None,None,0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # TMA LOAD - if warp_idx == self.tma_warp_id: - qp.reset(); qh = qp.acquire_and_advance() - cute.copy(tma_q,tAgQ[(None,qh.count)],tAsQ[(None,qh.index)],tma_bar_ptr=qh.barrier) - qp.tail() - kvp.reset(); pk = kvp.try_acquire() - for kt in cutlass.range(n_kv_tiles,unroll=1): - kh = kvp.acquire_and_advance(pk) - cute.copy(tma_k,tBgK[(None,kh.count)],tBsK[(None,kh.index)],tma_bar_ptr=kh.barrier) - pk = cutlass.Boolean(1) - vh = kvp.acquire_and_advance(pk) - cute.copy(tma_v,tVgV[(None,vh.count)],tVsV[(None,vh.index)],tma_bar_ptr=vh.barrier) - pk = cutlass.Boolean(1) - kvp.tail() - - # MMA - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - qc.reset(); qh = qc.wait_and_advance(); qh.release() - kvc.reset(); pk = kvc.try_wait() - acc_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_stage) - acc_pipe.producer_acquire(acc_st) - for kt in range(n_kv_tiles): - kh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - sh = s_prod.acquire_and_advance() - qk_mma.set(tcgen05.Field.ACCUMULATE, False) - for kb in cutlass.range(cute.size(tCrQ,mode=[2]), unroll_full=True): - cute.gemm(qk_mma, tStS0, tCrQ[(None,None,kb,0)], tCrK[(None,None,kb,kh.index)], tStS0) - qk_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - sh.commit(); kh.release() - softmax_done_bar.arrive_and_wait() - vh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - pv_mma.set(tcgen05.Field.ACCUMULATE, kt != 0) - for kb in cutlass.range(cute.size(tOrP0,mode=[2]), unroll_full=True): - cute.gemm(pv_mma, tOtO0, tOrP0[(None,None,kb)], tCrV[(None,None,kb,vh.index)], tOtO0) - pv_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - vh.release() - pv_done_bar.arrive() - acc_pipe.producer_commit(acc_st); acc_st.advance() - acc_pipe.producer_tail(acc_st) - - # ===================== EPILOGUE WARPS (STAGE C: ONLINE SOFTMAX) ===================== - if warp_idx < self.mma_warp_id: - tmem.allocate(self.num_tmem_alloc_cols) - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) - sfw_idx = tidx % (32 * len(self.epilogue_warp_id)) - - # --- S load (QK C-fragment) --- - tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS0) - thr_load = tiled_tmem_load.get_slice(sfw_idx) - tTMEM_LOADtS = thr_load.partition_S(tStS0) - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tTMEM_LOADcS = thr_load.partition_D(tScS) - - # --- P store (QK C-fragment composition, FMHA pattern) --- - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - tStP_layout = cute.composition(tStS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tStP0 = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStP_layout) - tmem_store_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStP0) - thr_store = tiled_tmem_store.get_slice(sfw_idx) - tTMEM_STOREtP = thr_store.partition_D(tStP0) - tScP_layout = cute.composition(tScS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tScP = cute.make_tensor(tScS.iterator, tScP_layout) - tTMEM_STOREcP = thr_store.partition_S(tScP) - - # --- Vector TMEM (per-row row_sum storage, FMHA pattern) --- - # composition(tStS.layout, (128, 2)) = 2 FP32 columns per logical row - # vec[0] = row_sum (final, after loop), vec[1] = unused - # Reuses S TMEM region (offset 0), free after softmax loop writes - - tStS_vec_layout = cute.composition(tStS.layout, cute.make_layout((128, 2))) - tStS_vec = cute.make_tensor(tStS.iterator + self.tmem_vec_offset, tStS_vec_layout) - tScS_vec_layout = cute.composition(tScS.layout, cute.make_layout((128, 2))) - tScS_vec = cute.make_tensor(tScS.iterator, tScS_vec_layout) - tmem_store_vec_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(2)), self.qk_acc_dtype) - tiled_tmem_store_vec = tcgen05.make_tmem_copy(tmem_store_vec_atom, tStS_vec) - thr_tmem_store_vec = tiled_tmem_store_vec.get_slice(sfw_idx) - tTMEM_STORE_VECtS = thr_tmem_store_vec.partition_D(tStS_vec) - tTMEM_STORE_VECcS = thr_tmem_store_vec.partition_S(tScS_vec) - tmem_load_vec_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(2)), self.qk_acc_dtype) - tiled_tmem_load_vec = tcgen05.make_tmem_copy(tmem_load_vec_atom, tStS_vec) - thr_tmem_load_vec = tiled_tmem_load_vec.get_slice(sfw_idx) - tTMEM_LOAD_VECtS = thr_tmem_load_vec.partition_S(tStS_vec) - tTMEM_LOAD_VECcS = thr_tmem_load_vec.partition_D(tScS_vec) - - # --- C6: O TMEM load/store for rescale (correction_rescale pattern) --- - corr_tile_size = 16 - cO = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1])) - tOcO = pv_thr.partition_C(cO) - o_tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), self.qk_acc_dtype) - o_tmem_store_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), self.qk_acc_dtype) - tOtO_i_layout = cute.composition(tOtO0.layout, cute.make_layout((128, corr_tile_size))) - tOcO_i_layout = cute.composition(tOcO.layout, cute.make_layout((128, corr_tile_size))) - tOtO_i = cute.make_tensor(tOtO0.iterator, tOtO_i_layout) - tOcO_i = cute.make_tensor(tOcO.iterator, tOcO_i_layout) - o_tiled_tmem_load = tcgen05.make_tmem_copy(o_tmem_load_atom, tOtO_i) - o_tiled_tmem_store = tcgen05.make_tmem_copy(o_tmem_store_atom, tOtO_i) - o_thr_load = o_tiled_tmem_load.get_slice(sfw_idx) - o_thr_store = o_tiled_tmem_store.get_slice(sfw_idx) - tTMEM_LOADtO = o_thr_load.partition_S(tOtO_i) - tTMEM_LOADcO = o_thr_load.partition_D(tOcO_i) - tTMEM_STOREtO = o_thr_store.partition_D(tOtO_i) - o_col_tiles = self.pv_mma_tiler[1] // corr_tile_size - - # --- C2: Per-thread row state (persist across KV tiles) --- - row_max = -cutlass.Float32.inf - row_sum = cutlass.Float32(0.0) - - # --- C3: QK scale = 1/sqrt(HEAD_DIM) * log2(e) for exp2 --- - scale = self.scale_softmax_log2 - - # ============================================================= - # Per-KV-tile online softmax loop - # ============================================================= - for kt in range(n_kv_tiles): - si_handle = s_cons.wait_and_advance() - - # Load S from TMEM (FP32, QK C-fragment layout) - tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS) - - # --- C4: Compute tile_max via .reduce(MAX) --- - old_row_max = row_max - row_max = tTMEM_LOADrS.load().reduce(cute.ReductionOp.MAX, row_max, 0) - row_max_safe = row_max - if row_max == -cutlass.Float32.inf: - row_max_safe = cutlass.Float32(0.0) - - # --- C5: Compute rescale factor --- - acc_scale = cute.math.exp2(scale * (old_row_max - row_max_safe), fastmath=True) - - # --- C6: Rescale O in TMEM (load O, multiply by acc_scale, store O) --- - # acc_scale belongs to QK row (N//4), but O rows are in PV partition (N). - # Store acc_scale to vector by QK row, read by PV row. - if kt > 0: - pv_done_bar.arrive_and_wait() - - # Store acc_scale to vector indexed by QK logical row - qk_row_c6 = tTMEM_LOADcS[0][0] - thr_vs_c6 = tiled_tmem_store_vec.get_slice(qk_row_c6) - tVStore_c6 = thr_vs_c6.partition_D(tStS_vec) - tVStoreSrc_c6 = thr_vs_c6.partition_S(tScS_vec) - tVStoreRmem_c6 = cute.make_rmem_tensor(tVStoreSrc_c6.shape, self.qk_acc_dtype) - tVStoreRmem_c6[0] = acc_scale - cute.copy(tiled_tmem_store_vec, tVStoreRmem_c6, tVStore_c6) - cute.arch.fence_view_async_tmem_store() - - # Read acc_scale from vector indexed by PV logical row - pv_row_c6 = tTMEM_LOADcO[0][0] - thr_vl_c6 = tiled_tmem_load_vec.get_slice(pv_row_c6) - tVLoad_c6 = thr_vl_c6.partition_S(tStS_vec) - tVLoadDst_c6 = thr_vl_c6.partition_D(tScS_vec) - tVLoadRmem_c6 = cute.make_rmem_tensor(tVLoadDst_c6.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load_vec, tVLoad_c6, tVLoadRmem_c6) - cute.arch.fence_view_async_tmem_load() - acc_scale_pv = tVLoadRmem_c6[0] - - tTMrO = cute.make_rmem_tensor((tTMEM_LOADcO.shape, o_col_tiles), self.qk_acc_dtype) - for i in range(o_col_tiles): - tTMrO_i_ = tTMrO[None, i] - tTMrO_i_layout = cute.composition(tTMrO_i_.layout, cute.make_layout(tTMrO.shape[0])) - tTMrO_i = cute.make_tensor(tTMrO_i_.iterator, tTMrO_i_layout) - tTMEM_LOADtO_i = cute.make_tensor(tTMEM_LOADtO.iterator + i * corr_tile_size, tTMEM_LOADtO.layout) - tTMEM_STOREtO_i = cute.make_tensor(tTMEM_STOREtO.iterator + i * corr_tile_size, tTMEM_STOREtO.layout) - cute.copy(o_tiled_tmem_load, tTMEM_LOADtO_i, tTMrO_i) - for j in cutlass.range(cute.size(tTMrO_i), vectorize=True): - tTMrO_i[j] = tTMrO_i[j] * acc_scale_pv - cute.copy(o_tiled_tmem_store, tTMrO_i, tTMEM_STOREtO_i) - cute.arch.fence_view_async_tmem_store() - - # Rescale row_sum - row_sum = row_sum * acc_scale - - # --- C7: Compute P = exp2((S - row_max_safe) * scale) --- - minus_row_max_scale = (cutlass.Float32(0.0) - row_max_safe) * scale - - # Register bridge (FMHA pattern: FP32 backing + BF16 view) - rP_words = cute.make_rmem_tensor(tTMEM_STOREcP.shape, self.qk_acc_dtype) - rP_bf16 = cute.make_tensor(cute.recast_ptr(rP_words.iterator, dtype=self.q_dtype), tTMEM_LOADrS.layout) - - frg_cnt = 4 - frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt - tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile)) - rP_bf16_frg = cute.logical_divide(rP_bf16, cute.make_layout(frg_tile)) - - # Scale S, compute exp2, store through register bridge - for j in range(frg_cnt): - for k in cutlass.range(cute.size(tTMEM_LOADrS_frg, mode=[0]), vectorize=True): - tTMEM_LOADrS_frg[k, j] = tTMEM_LOADrS_frg[k, j] * scale + minus_row_max_scale - tTMEM_LOADrS_frg[k, j] = cute.math.exp2(tTMEM_LOADrS_frg[k, j], fastmath=True) - s_vec = tTMEM_LOADrS_frg[None, j].load() - rP_bf16_frg[None, j].store(s_vec.to(self.q_dtype)) - - # Store P to TMEM - cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP) - cute.arch.fence_view_async_tmem_store() - si_handle.release() - softmax_done_bar.arrive() - - # --- C8: Row sum accumulation (CUTLASS FMHA packed f32x2 pattern) --- - # P values still in tTMEM_LOADrS registers. - # 4 accumulators for 4 reduction_unroll columns. - local_row_sum_0 = (cutlass.Float32(0.0), cutlass.Float32(0.0)) - local_row_sum_1 = (cutlass.Float32(0.0), cutlass.Float32(0.0)) - local_row_sum_2 = (cutlass.Float32(0.0), cutlass.Float32(0.0)) - local_row_sum_3 = (cutlass.Float32(0.0), cutlass.Float32(0.0)) - - reduction_unroll = 4 - rfrg_tile = cute.size(tTMEM_LOADrS) // reduction_unroll - tTMEM_LOADrS_rfrg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(rfrg_tile)) - - for j in cutlass.range_constexpr(0, cute.size(tTMEM_LOADrS_rfrg, mode=[0]), 2): - local_row_sum_0 = cute.arch.add_packed_f32x2( - local_row_sum_0, (tTMEM_LOADrS_rfrg[j, 0], tTMEM_LOADrS_rfrg[j + 1, 0])) - local_row_sum_1 = cute.arch.add_packed_f32x2( - local_row_sum_1, (tTMEM_LOADrS_rfrg[j, 1], tTMEM_LOADrS_rfrg[j + 1, 1])) - local_row_sum_2 = cute.arch.add_packed_f32x2( - local_row_sum_2, (tTMEM_LOADrS_rfrg[j, 2], tTMEM_LOADrS_rfrg[j + 1, 2])) - local_row_sum_3 = cute.arch.add_packed_f32x2( - local_row_sum_3, (tTMEM_LOADrS_rfrg[j, 3], tTMEM_LOADrS_rfrg[j + 1, 3])) - - local_row_sum_0 = cute.arch.add_packed_f32x2(local_row_sum_0, local_row_sum_1) - local_row_sum_2 = cute.arch.add_packed_f32x2(local_row_sum_2, local_row_sum_3) - local_row_sum_0 = cute.arch.add_packed_f32x2(local_row_sum_0, local_row_sum_2) - tile_sum = local_row_sum_0[0] + local_row_sum_0[1] - - row_sum = row_sum + tile_sum - - # --- C9: Final normalization via O TMEM rescale --- - pv_done_bar.arrive_and_wait() - - # Compute inv_row_sum by reading P from TMEM using PV A-fragment layout. - # P was stored by softmax into TMEM at tmem_p0_offset. - # The PV A-fragment (tOrP0) correctly maps thread N to PV row N's P values. - # Sum P per PV row to get the unnormalized row_sum. - # Since P = exp2((S - max) * scale) (unnormalized), row_sum = sum(P) per row. - # Then O_normalized = O_unnormalized / row_sum. - - # Read P using PV A-fragment (the same layout PV MMA uses to read P) - # tOrP0 is already set up for PV MMA input. - # We need to sum its values per thread. - pv_row_sum = cutlass.Float32(0.0) - for kb in cutlass.range(cute.size(tOrP0, mode=[2])): - p_frag = tOrP0[(None, None, kb)] - for j in cutlass.range(cute.size(p_frag), vectorize=True): - pv_row_sum = pv_row_sum + p_frag[j] - - inv_row_sum = cutlass.Float32(1.0) / pv_row_sum - - # Normalize O in TMEM using PV-correct inv_row_sum - tTMrO_final = cute.make_rmem_tensor((tTMEM_LOADcO.shape, o_col_tiles), self.qk_acc_dtype) - for i in range(o_col_tiles): - tTMrO_i_ = tTMrO_final[None, i] - tTMrO_i_layout = cute.composition(tTMrO_i_.layout, cute.make_layout(tTMrO_final.shape[0])) - tTMrO_i = cute.make_tensor(tTMrO_i_.iterator, tTMrO_i_layout) - tTMEM_LOADtO_i = cute.make_tensor( - tTMEM_LOADtO.iterator + i * corr_tile_size, tTMEM_LOADtO.layout) - tTMEM_STOREtO_i = cute.make_tensor( - tTMEM_STOREtO.iterator + i * corr_tile_size, tTMEM_STOREtO.layout) - cute.copy(o_tiled_tmem_load, tTMEM_LOADtO_i, tTMrO_i) - for j in cutlass.range(cute.size(tTMrO_i), vectorize=True): - tTMrO_i[j] = tTMrO_i[j] * inv_row_sum - cute.copy(o_tiled_tmem_store, tTMrO_i, tTMEM_STOREtO_i) - cute.arch.fence_view_async_tmem_store() - - # Now O in TMEM is normalized. Use standard epilogue_tma_store with identity. - tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout) - acc_cons_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage) - c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) - c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp) - acc_cons_st = utils.gemm.sm100.epilogue_tma_store( - self, tidx, warp_idx, tma_c, tCtO_base, sC, tCgC, epi_tile, 0, - const_expr(lambda x: x), - (0,0,0), acc_cons_st, acc_pipe, c_pipe) - c_pipe.producer_tail() - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def test(): - import math - torch.manual_seed(42) - for n in [128, 256, 384]: - m, hd = 128, HEAD_DIM - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device="cuda") - k = torch.randn(n, hd, 1, dtype=torch.bfloat16, device="cuda") - v = torch.randn(n, hd, dtype=torch.bfloat16, device="cuda") - v_kernel = v.unsqueeze(-1) - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device="cuda") - qf = q[:,:,0].float(); kf = k[:,:,0].float() - attn = qf @ kf.T / math.sqrt(hd) - ref = torch.softmax(attn, dim=-1) @ v.float() - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - kernel = FmhaV3Softmax() - print(f"n={n}: Compiling...", flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - print(f"n={n}: tmem: s0={kernel.tmem_s0_offset} p0={kernel.tmem_p0_offset} o0={kernel.tmem_o0_offset} vec={kernel.tmem_vec_offset} alloc={kernel.num_tmem_alloc_cols}", flush=True) - print(f"n={n}: Running...", flush=True) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - out = c[:,:,0].float() - cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - max_err = (out - ref).abs().max().item() - print(f"FMHA softmax n={n}: cosine {cos:.6f} max_err {max_err:.6f} {'PASS' if cos >= 0.999 else 'FAIL'}", flush=True) - -if __name__ == "__main__": - test() - - diff --git a/tests/archive/unit_test_fmha_v3_scalar.py b/tests/archive/unit_test_fmha_v3_scalar.py deleted file mode 100644 index 113073f1..00000000 --- a/tests/archive/unit_test_fmha_v3_scalar.py +++ /dev/null @@ -1,493 +0,0 @@ -""" -FMHA v3 + Stage C: QK -> online softmax -> PV with KV-tile interleaving. -Stage C: row_max, exp2, O rescale, row_sum, final normalization. -FMHA pattern P store preserved from Stage B. -""" -import math -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - -HEAD_DIM = 64 - -class FmhaV3Softmax: - def __init__(self): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0,1,2,3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.num_c_stage = 2 - self.kv_stage = 2; self.q_stage = 1; self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (128, 128, qk_ik * 4) - pv_ik = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (128, HEAD_DIM, pv_ik * (128 // pv_ik)) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0]//cute.size(qk_mma.thr_id.shape), HEAD_DIM, self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - self.q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, self.q_dtype, self.q_stage) - self.k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, self.q_dtype, self.kv_stage) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, self.kv_stage) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - qk_thr = qk_mma.get_slice(0); qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - pv_thr = pv_mma.get_slice(0); pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - self.tmem_s0_offset = 0; self.tmem_p0_offset = 32 - # P occupies [tmem_p0_offset, tmem_p0_offset + p_cols_fp32) - # S occupies [0, qk_mma_tiler[1]) = [0, 128) - # O must NOT overlap P. Place O after max(S end, P end), aligned to 32. - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - p_end = self.tmem_p0_offset + p_cols_fp32 # 32 + 64 = 96 - s_cols = self.qk_mma_tiler[1] # 128 - o_after = max(s_cols, p_end) # 128 - self.tmem_o0_offset = ((o_after + 31) // 32) * 32 - self.tmem_vec_offset = 0 # Reuse S region for per-row inv_row_sum vector # align to 32 = 128 - self.tmem_vec_offset = 0 # Reuse S region (free after softmax loop) - o_cols = find_tmem_tensor_col_offset(tOtO) # footprint of O - total = self.tmem_o0_offset + o_cols - # Must be multiple of 32 AND power of 2 - self.num_tmem_alloc_cols = 1 - while self.num_tmem_alloc_cols < total: - self.num_tmem_alloc_cols *= 2 - cta = cute.size(qk_mma.thr_id.shape) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)) - self.q_tx_bytes = cute.size_in_bytes(self.q_dtype, q_s) * cta - self.kv_tx_bytes = cute.size_in_bytes(self.q_dtype, k_s) * cta - self.scale_softmax_log2 = Float32(1.0 / math.sqrt(HEAD_DIM) * math.log2(math.e)) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - # # s_k hardcoded # BROKEN in @cute.jit - # FMHA-style V: reconstruct as (HEAD_DIM, s_k, 1) MN-major - v_fmha = cute.make_tensor( - v.iterator, - cute.make_layout( - (HEAD_DIM, 128, 1), - stride=(1, HEAD_DIM, HEAD_DIM * 128), - ), - ) - self.v_major = LayoutEnum.from_tensor(v_fmha).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, self.a_major, self.b_major, self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, self.qk_acc_dtype, self.cta_group, (128,HEAD_DIM), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)); v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id),q,q_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id),k,k_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v_fmha,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape) - epi_s = cute.select(self.c_smem_s,mode=[0,1]) - tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile) - self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.v_smem_s,self.p_tmem_s,self.c_smem_s,self.epi_tile).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, q_smem_s, k_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx,_,_ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k); cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - q_bar: cute.struct.MemRange[cutlass.Int64, self.q_stage*2] - kv_bar: cute.struct.MemRange[cutlass.Int64, self.kv_stage*2] - s_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage*2] - tmem_dealloc: cutlass.Int64; holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - qp,qc = pipeline.PipelineTmaUmma.create(barrier_storage=st.q_bar.data_ptr(),num_stages=self.q_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.q_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - kvp,kvc = pipeline.PipelineTmaUmma.create(barrier_storage=st.kv_bar.data_ptr(),num_stages=self.kv_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.kv_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - s_prod,s_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.s_bar.data_ptr(),num_stages=1,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.epilogue_warp_id))).make_participants() - softmax_done_bar = pipeline.NamedBarrier(barrier_id=3, num_threads=32 + 32*len(self.epilogue_warp_id)) - pv_done_bar = pipeline.NamedBarrier(barrier_id=4, num_threads=32 + 32*len(self.epilogue_warp_id)) - acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(),num_stages=self.num_acc_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,len(self.epilogue_warp_id)),cta_layout_vmnk=cl_vmnk,defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2,num_threads=32*len((self.mma_warp_id,*self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr,barrier_for_retrieve=tmem_bar,allocator_warp_id=self.epilogue_warp_id[0],is_two_cta=cute.size(qk_mma.thr_id.shape)==2,two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk,is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype,layout=q_smem_s.outer,byte_alignment=128,swizzle=q_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype,layout=k_smem_s.outer,byte_alignment=128,swizzle=k_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype,layout=v_smem_s.outer,byte_alignment=128,swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype,layout=c_smem_s.outer,byte_alignment=128,swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None)) - gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None)) - gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None)) - gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None)) - n_kv_tiles = cute.size(gK, mode=[3]) - - qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - tCgV = pv_thr.partition_B(gV); tCgC = pv_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,0,None,0)).shape) - tAsQ,tAgQ = cpasync.tma_partition(tma_q,0,a_lay,cute.group_modes(sQ,0,3),cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape) - tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3)) - tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)]; tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - # --- PV read view (for MMA only, NOT for softmax store) --- - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None,None,None,0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # TMA LOAD - if warp_idx == self.tma_warp_id: - qp.reset(); qh = qp.acquire_and_advance() - cute.copy(tma_q,tAgQ[(None,qh.count)],tAsQ[(None,qh.index)],tma_bar_ptr=qh.barrier) - qp.tail() - kvp.reset(); pk = kvp.try_acquire() - for kt in cutlass.range(n_kv_tiles,unroll=1): - kh = kvp.acquire_and_advance(pk) - cute.copy(tma_k,tBgK[(None,kh.count)],tBsK[(None,kh.index)],tma_bar_ptr=kh.barrier) - pk = cutlass.Boolean(1) - vh = kvp.acquire_and_advance(pk) - cute.copy(tma_v,tVgV[(None,vh.count)],tVsV[(None,vh.index)],tma_bar_ptr=vh.barrier) - pk = cutlass.Boolean(1) - kvp.tail() - - # MMA - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - qc.reset(); qh = qc.wait_and_advance(); qh.release() - kvc.reset(); pk = kvc.try_wait() - acc_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_stage) - acc_pipe.producer_acquire(acc_st) - for kt in range(n_kv_tiles): - kh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - sh = s_prod.acquire_and_advance() - qk_mma.set(tcgen05.Field.ACCUMULATE, False) - for kb in cutlass.range(cute.size(tCrQ,mode=[2]), unroll_full=True): - cute.gemm(qk_mma, tStS0, tCrQ[(None,None,kb,0)], tCrK[(None,None,kb,kh.index)], tStS0) - qk_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - sh.commit(); kh.release() - softmax_done_bar.arrive_and_wait() - vh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - pv_mma.set(tcgen05.Field.ACCUMULATE, kt != 0) - for kb in cutlass.range(cute.size(tOrP0,mode=[2]), unroll_full=True): - cute.gemm(pv_mma, tOtO0, tOrP0[(None,None,kb)], tCrV[(None,None,kb,vh.index)], tOtO0) - pv_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - vh.release() - pv_done_bar.arrive() - acc_pipe.producer_commit(acc_st); acc_st.advance() - acc_pipe.producer_tail(acc_st) - - # ===================== EPILOGUE WARPS (STAGE C: ONLINE SOFTMAX) ===================== - if warp_idx < self.mma_warp_id: - tmem.allocate(self.num_tmem_alloc_cols) - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) - sfw_idx = tidx % (32 * len(self.epilogue_warp_id)) - - # --- S load (QK C-fragment) --- - tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS0) - thr_load = tiled_tmem_load.get_slice(sfw_idx) - tTMEM_LOADtS = thr_load.partition_S(tStS0) - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tTMEM_LOADcS = thr_load.partition_D(tScS) - - # --- P store (QK C-fragment composition, FMHA pattern) --- - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - tStP_layout = cute.composition(tStS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tStP0 = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStP_layout) - tmem_store_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStP0) - thr_store = tiled_tmem_store.get_slice(sfw_idx) - tTMEM_STOREtP = thr_store.partition_D(tStP0) - tScP_layout = cute.composition(tScS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tScP = cute.make_tensor(tScS.iterator, tScP_layout) - tTMEM_STOREcP = thr_store.partition_S(tScP) - - # --- Vector TMEM (per-row row_sum storage, FMHA pattern) --- - # composition(tStS.layout, (128, 2)) = 2 FP32 columns per logical row - # vec[0] = row_sum (final, after loop), vec[1] = unused - # Reuses S TMEM region (offset 0), free after softmax loop writes - - tStS_vec_layout = cute.composition(tStS.layout, cute.make_layout((128, 2))) - tStS_vec = cute.make_tensor(tStS.iterator + self.tmem_vec_offset, tStS_vec_layout) - tScS_vec_layout = cute.composition(tScS.layout, cute.make_layout((128, 2))) - tScS_vec = cute.make_tensor(tScS.iterator, tScS_vec_layout) - tmem_store_vec_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(2)), self.qk_acc_dtype) - tiled_tmem_store_vec = tcgen05.make_tmem_copy(tmem_store_vec_atom, tStS_vec) - thr_tmem_store_vec = tiled_tmem_store_vec.get_slice(sfw_idx) - tTMEM_STORE_VECtS = thr_tmem_store_vec.partition_D(tStS_vec) - tTMEM_STORE_VECcS = thr_tmem_store_vec.partition_S(tScS_vec) - tmem_load_vec_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(2)), self.qk_acc_dtype) - tiled_tmem_load_vec = tcgen05.make_tmem_copy(tmem_load_vec_atom, tStS_vec) - thr_tmem_load_vec = tiled_tmem_load_vec.get_slice(sfw_idx) - tTMEM_LOAD_VECtS = thr_tmem_load_vec.partition_S(tStS_vec) - tTMEM_LOAD_VECcS = thr_tmem_load_vec.partition_D(tScS_vec) - - # --- C6: O TMEM load/store for rescale (correction_rescale pattern) --- - corr_tile_size = 16 - cO = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1])) - tOcO = pv_thr.partition_C(cO) - o_tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), self.qk_acc_dtype) - o_tmem_store_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), self.qk_acc_dtype) - tOtO_i_layout = cute.composition(tOtO0.layout, cute.make_layout((128, corr_tile_size))) - tOcO_i_layout = cute.composition(tOcO.layout, cute.make_layout((128, corr_tile_size))) - tOtO_i = cute.make_tensor(tOtO0.iterator, tOtO_i_layout) - tOcO_i = cute.make_tensor(tOcO.iterator, tOcO_i_layout) - o_tiled_tmem_load = tcgen05.make_tmem_copy(o_tmem_load_atom, tOtO_i) - o_tiled_tmem_store = tcgen05.make_tmem_copy(o_tmem_store_atom, tOtO_i) - o_thr_load = o_tiled_tmem_load.get_slice(sfw_idx) - o_thr_store = o_tiled_tmem_store.get_slice(sfw_idx) - tTMEM_LOADtO = o_thr_load.partition_S(tOtO_i) - tTMEM_LOADcO = o_thr_load.partition_D(tOcO_i) - tTMEM_STOREtO = o_thr_store.partition_D(tOtO_i) - o_col_tiles = self.pv_mma_tiler[1] // corr_tile_size - - # --- C2: Per-thread row state (persist across KV tiles) --- - row_max = -cutlass.Float32.inf - row_sum = cutlass.Float32(0.0) - - # --- C3: QK scale = 1/sqrt(HEAD_DIM) * log2(e) for exp2 --- - scale = self.scale_softmax_log2 - - # ============================================================= - # Per-KV-tile online softmax loop - # ============================================================= - for kt in range(n_kv_tiles): - si_handle = s_cons.wait_and_advance() - - # Load S from TMEM (FP32, QK C-fragment layout) - tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS) - - # --- C4: Compute tile_max via .reduce(MAX) --- - old_row_max = row_max - row_max = tTMEM_LOADrS.load().reduce(cute.ReductionOp.MAX, row_max, 0) - row_max_safe = row_max - if row_max == -cutlass.Float32.inf: - row_max_safe = cutlass.Float32(0.0) - - # --- C5: Compute rescale factor --- - acc_scale = cute.math.exp2(scale * (old_row_max - row_max_safe), fastmath=True) - - # --- C6: Rescale O in TMEM (direct scalar approach) --- - # Each softmax thread computes acc_scale from its QK row_max. - # In the QK C-fragment with 128 threads and 128 rows, thread N = row N. - # In the PV C-fragment with 128 threads and 128 rows, thread N = row N. - # So acc_scale for thread N's QK row = acc_scale for thread N's PV row. - # Use acc_scale directly (no vector indirection needed). - if kt > 0: - pv_done_bar.arrive_and_wait() - - acc_scale_pv = acc_scale # Direct scalar - - tTMrO = cute.make_rmem_tensor((tTMEM_LOADcO.shape, o_col_tiles), self.qk_acc_dtype) - for i in range(o_col_tiles): - tTMrO_i_ = tTMrO[None, i] - tTMrO_i_layout = cute.composition(tTMrO_i_.layout, cute.make_layout(tTMrO.shape[0])) - tTMrO_i = cute.make_tensor(tTMrO_i_.iterator, tTMrO_i_layout) - tTMEM_LOADtO_i = cute.make_tensor(tTMEM_LOADtO.iterator + i * corr_tile_size, tTMEM_LOADtO.layout) - tTMEM_STOREtO_i = cute.make_tensor(tTMEM_STOREtO.iterator + i * corr_tile_size, tTMEM_STOREtO.layout) - cute.copy(o_tiled_tmem_load, tTMEM_LOADtO_i, tTMrO_i) - for j in cutlass.range(cute.size(tTMrO_i), vectorize=True): - tTMrO_i[j] = tTMrO_i[j] * acc_scale_pv - cute.copy(o_tiled_tmem_store, tTMrO_i, tTMEM_STOREtO_i) - cute.arch.fence_view_async_tmem_store() - - # # --- C7: Compute P = exp2((S - row_max_safe) * scale) --- - minus_row_max_scale = (cutlass.Float32(0.0) - row_max_safe) * scale - - # Register bridge (FMHA pattern: FP32 backing + BF16 view) - rP_words = cute.make_rmem_tensor(tTMEM_STOREcP.shape, self.qk_acc_dtype) - rP_bf16 = cute.make_tensor(cute.recast_ptr(rP_words.iterator, dtype=self.q_dtype), tTMEM_LOADrS.layout) - - frg_cnt = 4 - frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt - tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile)) - rP_bf16_frg = cute.logical_divide(rP_bf16, cute.make_layout(frg_tile)) - - # Scale S, compute exp2, store through register bridge - for j in range(frg_cnt): - for k in cutlass.range(cute.size(tTMEM_LOADrS_frg, mode=[0]), vectorize=True): - tTMEM_LOADrS_frg[k, j] = tTMEM_LOADrS_frg[k, j] * scale + minus_row_max_scale - tTMEM_LOADrS_frg[k, j] = cute.math.exp2(tTMEM_LOADrS_frg[k, j], fastmath=True) - s_vec = tTMEM_LOADrS_frg[None, j].load() - rP_bf16_frg[None, j].store(s_vec.to(self.q_dtype)) - - # Store P to TMEM - cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP) - cute.arch.fence_view_async_tmem_store() - si_handle.release() - softmax_done_bar.arrive() - - # --- C8: Row sum accumulation (CUTLASS FMHA packed f32x2 pattern) --- - # P values still in tTMEM_LOADrS registers. - # 4 accumulators for 4 reduction_unroll columns. - local_row_sum_0 = (cutlass.Float32(0.0), cutlass.Float32(0.0)) - local_row_sum_1 = (cutlass.Float32(0.0), cutlass.Float32(0.0)) - local_row_sum_2 = (cutlass.Float32(0.0), cutlass.Float32(0.0)) - local_row_sum_3 = (cutlass.Float32(0.0), cutlass.Float32(0.0)) - - reduction_unroll = 4 - rfrg_tile = cute.size(tTMEM_LOADrS) // reduction_unroll - tTMEM_LOADrS_rfrg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(rfrg_tile)) - - for j in cutlass.range_constexpr(0, cute.size(tTMEM_LOADrS_rfrg, mode=[0]), 2): - local_row_sum_0 = cute.arch.add_packed_f32x2( - local_row_sum_0, (tTMEM_LOADrS_rfrg[j, 0], tTMEM_LOADrS_rfrg[j + 1, 0])) - local_row_sum_1 = cute.arch.add_packed_f32x2( - local_row_sum_1, (tTMEM_LOADrS_rfrg[j, 1], tTMEM_LOADrS_rfrg[j + 1, 1])) - local_row_sum_2 = cute.arch.add_packed_f32x2( - local_row_sum_2, (tTMEM_LOADrS_rfrg[j, 2], tTMEM_LOADrS_rfrg[j + 1, 2])) - local_row_sum_3 = cute.arch.add_packed_f32x2( - local_row_sum_3, (tTMEM_LOADrS_rfrg[j, 3], tTMEM_LOADrS_rfrg[j + 1, 3])) - - local_row_sum_0 = cute.arch.add_packed_f32x2(local_row_sum_0, local_row_sum_1) - local_row_sum_2 = cute.arch.add_packed_f32x2(local_row_sum_2, local_row_sum_3) - local_row_sum_0 = cute.arch.add_packed_f32x2(local_row_sum_0, local_row_sum_2) - tile_sum = local_row_sum_0[0] + local_row_sum_0[1] - - row_sum = row_sum + tile_sum - - # --- C9: Final normalization via O TMEM rescale --- - pv_done_bar.arrive_and_wait() - - # Compute inv_row_sum from P in TMEM using PV partition. - # P was stored by softmax loop into TMEM at offset tmem_p0_offset. - # PV partition maps thread N to PV row N, so reading P via PV partition - # gives the correct per-row P values to sum. - # This avoids the QK→PV row mapping mismatch (QK: N->N//4, PV: N->N). - - # P is stored as BF16 in TMEM at tmem_p0_offset. - # We need to read it via PV TMEM load and sum the values. - # P has shape (128, HEAD_DIM//2) in FP32 columns (64 BF16 = 32 FP32 cols). - # Use the P TMEM load partition (PV A-fragment read). - - # Actually, P was stored via QK C-fragment store (St32x32bOp Repetition(32)). - # To read it via PV partition, we need a PV-partitioned load from the P region. - # Let's use the same o_tiled_tmem_load but pointed at P's TMEM offset. - - # P occupies TMEM columns [tmem_p0_offset, tmem_p0_offset + p_cols_fp32) - # In the PV C-fragment, P is the A-fragment. We can use tOrP0's layout. - # tOrP0 was set up with offset for PV MMA read. - - # Simpler: sum O across columns to get unnormalized row sum, then normalize. - # For V=identity, O = P@V = sum(P per row). So O.sum(dim=-1) = row_sum. - # For arbitrary V, O = P@V. O.sum(dim=-1) = sum_j(P@V)[j] = sum_j(sum_i P[i]*V[i,j]) - # This is NOT sum(P). So this trick only works for V=identity. - - # Correct approach: read P from TMEM, sum it per PV row. - # P is at TMEM offset tmem_p0_offset, stored as BF16 with St32x32bOp. - # P shape in TMEM: 128 rows x (HEAD_DIM BF16 = 32 FP32 cols) - # We can read P using Ld32x32bOp(Repetition(corr_tile_size)) via PV O-partition. - - # Use PV O TMEM load to read from P region instead of O region - p_col_tiles = p_cols_fp32 // corr_tile_size # 32 // 16 = 2 - pv_row_sum = cutlass.Float32(0.0) - for i in range(p_col_tiles): - # Read P tile from TMEM at P offset (not O offset) - tTMEM_LOADtP_i = cute.make_tensor( - tTMEM_LOADtO.iterator + (self.tmem_p0_offset - self.tmem_o0_offset) + i * corr_tile_size, - tTMEM_LOADtO.layout) - tTMrP_i = cute.make_rmem_tensor(tTMEM_LOADcO.shape, self.qk_acc_dtype) - cute.copy(o_tiled_tmem_load, tTMEM_LOADtP_i, tTMrP_i) - # Use .reduce(SUM) instead of scalar accumulation (vectorizer can't handle scalar in vectorized loop) - tile_p_sum = tTMrP_i.load().reduce(cute.ReductionOp.ADD, cutlass.Float32(0.0), 0) - pv_row_sum = pv_row_sum + tile_p_sum - - inv_row_sum = cutlass.Float32(1.0) / pv_row_sum - - # Normalize O in TMEM using PV-correct inv_row_sum - tTMrO_final = cute.make_rmem_tensor((tTMEM_LOADcO.shape, o_col_tiles), self.qk_acc_dtype) - for i in range(o_col_tiles): - tTMrO_i_ = tTMrO_final[None, i] - tTMrO_i_layout = cute.composition(tTMrO_i_.layout, cute.make_layout(tTMrO_final.shape[0])) - tTMrO_i = cute.make_tensor(tTMrO_i_.iterator, tTMrO_i_layout) - tTMEM_LOADtO_i = cute.make_tensor( - tTMEM_LOADtO.iterator + i * corr_tile_size, tTMEM_LOADtO.layout) - tTMEM_STOREtO_i = cute.make_tensor( - tTMEM_STOREtO.iterator + i * corr_tile_size, tTMEM_STOREtO.layout) - cute.copy(o_tiled_tmem_load, tTMEM_LOADtO_i, tTMrO_i) - for j in cutlass.range(cute.size(tTMrO_i), vectorize=True): - tTMrO_i[j] = tTMrO_i[j] * inv_row_sum - cute.copy(o_tiled_tmem_store, tTMrO_i, tTMEM_STOREtO_i) - cute.arch.fence_view_async_tmem_store() - - # Now O in TMEM is normalized. Use standard epilogue_tma_store with identity. - tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout) - acc_cons_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage) - c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) - c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp) - acc_cons_st = utils.gemm.sm100.epilogue_tma_store( - self, tidx, warp_idx, tma_c, tCtO_base, sC, tCgC, epi_tile, 0, - const_expr(lambda x: x), - (0,0,0), acc_cons_st, acc_pipe, c_pipe) - c_pipe.producer_tail() - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def test(): - import math - torch.manual_seed(42) - for n in [128, 256, 384]: - m, hd = 128, HEAD_DIM - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device="cuda") - k = torch.randn(n, hd, 1, dtype=torch.bfloat16, device="cuda") - v = torch.randn(n, hd, dtype=torch.bfloat16, device="cuda") - v_kernel = v.unsqueeze(-1) - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device="cuda") - qf = q[:,:,0].float(); kf = k[:,:,0].float() - attn = qf @ kf.T / math.sqrt(hd) - ref = torch.softmax(attn, dim=-1) @ v.float() - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - kernel = FmhaV3Softmax() - print(f"n={n}: Compiling...", flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - print(f"n={n}: tmem: s0={kernel.tmem_s0_offset} p0={kernel.tmem_p0_offset} o0={kernel.tmem_o0_offset} vec={kernel.tmem_vec_offset} alloc={kernel.num_tmem_alloc_cols}", flush=True) - print(f"n={n}: Running...", flush=True) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - out = c[:,:,0].float() - cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - max_err = (out - ref).abs().max().item() - print(f"FMHA softmax n={n}: cosine {cos:.6f} max_err {max_err:.6f} {'PASS' if cos >= 0.999 else 'FAIL'}", flush=True) - -if __name__ == "__main__": - test() - - diff --git a/tests/archive/unit_test_fmha_v3_shapes.py b/tests/archive/unit_test_fmha_v3_shapes.py deleted file mode 100644 index c6405285..00000000 --- a/tests/archive/unit_test_fmha_v3_shapes.py +++ /dev/null @@ -1,518 +0,0 @@ -""" -FMHA v3 + Stage C: QK -> online softmax -> PV with KV-tile interleaving. -Stage C: row_max, exp2, O rescale, row_sum, final normalization. -FMHA pattern P store preserved from Stage B. -""" -import math -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - -HEAD_DIM = 64 - -class FmhaV3Softmax: - def __init__(self): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0,1,2,3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.num_c_stage = 2 - self.kv_stage = 2; self.q_stage = 1; self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (128, 128, qk_ik * 4) - pv_ik = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (128, HEAD_DIM, pv_ik * (128 // pv_ik)) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0]//cute.size(qk_mma.thr_id.shape), HEAD_DIM, self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - self.q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, self.q_dtype, self.q_stage) - self.k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, self.q_dtype, self.kv_stage) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, self.kv_stage) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - qk_thr = qk_mma.get_slice(0); qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - pv_thr = pv_mma.get_slice(0); pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - self.tmem_s0_offset = 0; self.tmem_p0_offset = 32 - # P occupies [tmem_p0_offset, tmem_p0_offset + p_cols_fp32) - # S occupies [0, qk_mma_tiler[1]) = [0, 128) - # O must NOT overlap P. Place O after max(S end, P end), aligned to 32. - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - p_end = self.tmem_p0_offset + p_cols_fp32 # 32 + 64 = 96 - s_cols = self.qk_mma_tiler[1] # 128 - o_after = max(s_cols, p_end) # 128 - self.tmem_o0_offset = ((o_after + 31) // 32) * 32 - self.tmem_vec_offset = 0 # Reuse S region for per-row inv_row_sum vector # align to 32 = 128 - self.tmem_vec_offset = 0 # Reuse S region (free after softmax loop) - o_cols = find_tmem_tensor_col_offset(tOtO) # footprint of O - total = self.tmem_o0_offset + o_cols - # Must be multiple of 32 AND power of 2 - self.num_tmem_alloc_cols = 1 - while self.num_tmem_alloc_cols < total: - self.num_tmem_alloc_cols *= 2 - cta = cute.size(qk_mma.thr_id.shape) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)) - self.q_tx_bytes = cute.size_in_bytes(self.q_dtype, q_s) * cta - self.kv_tx_bytes = cute.size_in_bytes(self.q_dtype, k_s) * cta - self.scale_softmax_log2 = Float32(1.0 / math.sqrt(HEAD_DIM) * math.log2(math.e)) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - # # s_k hardcoded # BROKEN in @cute.jit - # FMHA-style V: reconstruct as (HEAD_DIM, s_k, 1) MN-major - v_fmha = cute.make_tensor( - v.iterator, - cute.make_layout( - (HEAD_DIM, 128, 1), - stride=(1, HEAD_DIM, HEAD_DIM * 128), - ), - ) - self.v_major = LayoutEnum.from_tensor(v_fmha).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, self.a_major, self.b_major, self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, self.qk_acc_dtype, self.cta_group, (128,HEAD_DIM), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)); v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id),q,q_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id),k,k_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v_fmha,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape) - epi_s = cute.select(self.c_smem_s,mode=[0,1]) - tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile) - self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.v_smem_s,self.p_tmem_s,self.c_smem_s,self.epi_tile).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, q_smem_s, k_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx,_,_ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k); cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - q_bar: cute.struct.MemRange[cutlass.Int64, self.q_stage*2] - kv_bar: cute.struct.MemRange[cutlass.Int64, self.kv_stage*2] - s_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage*2] - tmem_dealloc: cutlass.Int64; holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - qp,qc = pipeline.PipelineTmaUmma.create(barrier_storage=st.q_bar.data_ptr(),num_stages=self.q_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.q_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - kvp,kvc = pipeline.PipelineTmaUmma.create(barrier_storage=st.kv_bar.data_ptr(),num_stages=self.kv_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.kv_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - s_prod,s_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.s_bar.data_ptr(),num_stages=1,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.epilogue_warp_id))).make_participants() - softmax_done_bar = pipeline.NamedBarrier(barrier_id=3, num_threads=32 + 32*len(self.epilogue_warp_id)) - pv_done_bar = pipeline.NamedBarrier(barrier_id=4, num_threads=32 + 32*len(self.epilogue_warp_id)) - acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(),num_stages=self.num_acc_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,len(self.epilogue_warp_id)),cta_layout_vmnk=cl_vmnk,defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2,num_threads=32*len((self.mma_warp_id,*self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr,barrier_for_retrieve=tmem_bar,allocator_warp_id=self.epilogue_warp_id[0],is_two_cta=cute.size(qk_mma.thr_id.shape)==2,two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk,is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype,layout=q_smem_s.outer,byte_alignment=128,swizzle=q_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype,layout=k_smem_s.outer,byte_alignment=128,swizzle=k_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype,layout=v_smem_s.outer,byte_alignment=128,swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype,layout=c_smem_s.outer,byte_alignment=128,swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None)) - gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None)) - gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None)) - gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None)) - n_kv_tiles = cute.size(gK, mode=[3]) - - qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - tCgV = pv_thr.partition_B(gV); tCgC = pv_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,0,None,0)).shape) - tAsQ,tAgQ = cpasync.tma_partition(tma_q,0,a_lay,cute.group_modes(sQ,0,3),cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape) - tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3)) - tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)]; tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - # --- PV read view (for MMA only, NOT for softmax store) --- - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None,None,None,0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # TMA LOAD - if warp_idx == self.tma_warp_id: - qp.reset(); qh = qp.acquire_and_advance() - cute.copy(tma_q,tAgQ[(None,qh.count)],tAsQ[(None,qh.index)],tma_bar_ptr=qh.barrier) - qp.tail() - kvp.reset(); pk = kvp.try_acquire() - for kt in cutlass.range(n_kv_tiles,unroll=1): - kh = kvp.acquire_and_advance(pk) - cute.copy(tma_k,tBgK[(None,kh.count)],tBsK[(None,kh.index)],tma_bar_ptr=kh.barrier) - pk = cutlass.Boolean(1) - vh = kvp.acquire_and_advance(pk) - cute.copy(tma_v,tVgV[(None,vh.count)],tVsV[(None,vh.index)],tma_bar_ptr=vh.barrier) - pk = cutlass.Boolean(1) - kvp.tail() - - # MMA - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - qc.reset(); qh = qc.wait_and_advance(); qh.release() - kvc.reset(); pk = kvc.try_wait() - acc_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_stage) - acc_pipe.producer_acquire(acc_st) - for kt in range(n_kv_tiles): - kh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - sh = s_prod.acquire_and_advance() - qk_mma.set(tcgen05.Field.ACCUMULATE, False) - for kb in cutlass.range(cute.size(tCrQ,mode=[2]), unroll_full=True): - cute.gemm(qk_mma, tStS0, tCrQ[(None,None,kb,0)], tCrK[(None,None,kb,kh.index)], tStS0) - qk_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - sh.commit(); kh.release() - softmax_done_bar.arrive_and_wait() - vh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - pv_mma.set(tcgen05.Field.ACCUMULATE, kt != 0) - for kb in cutlass.range(cute.size(tOrP0,mode=[2]), unroll_full=True): - cute.gemm(pv_mma, tOtO0, tOrP0[(None,None,kb)], tCrV[(None,None,kb,vh.index)], tOtO0) - pv_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - vh.release() - pv_done_bar.arrive() - acc_pipe.producer_commit(acc_st); acc_st.advance() - acc_pipe.producer_tail(acc_st) - - # ===================== EPILOGUE WARPS (STAGE C: ONLINE SOFTMAX) ===================== - if warp_idx < self.mma_warp_id: - tmem.allocate(self.num_tmem_alloc_cols) - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) - sfw_idx = tidx % (32 * len(self.epilogue_warp_id)) # DEBUG: print fragment shapes (only from thread 0) - if sfw_idx == 0: - print(f"DEBUG sfw_idx=0: tTMEM_LOADcS shape={tTMEM_LOADcS.shape} size={cute.size(tTMEM_LOADcS)}") - print(f"DEBUG sfw_idx=0: tScS shape={tScS.shape} size={cute.size(tScS)}") - # Check which rows thread 0 handles - for i in range(min(4, cute.size(tScS, mode=[0]))): - row_col = tScS[i][0] - print(f" tScS[{i}][0] = {row_col}") - - # --- S load (QK C-fragment) --- - tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS0) - thr_load = tiled_tmem_load.get_slice(sfw_idx) - tTMEM_LOADtS = thr_load.partition_S(tStS0) - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tTMEM_LOADcS = thr_load.partition_D(tScS) - - # --- P store (QK C-fragment composition, FMHA pattern) --- - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - tStP_layout = cute.composition(tStS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tStP0 = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStP_layout) - tmem_store_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStP0) - thr_store = tiled_tmem_store.get_slice(sfw_idx) - tTMEM_STOREtP = thr_store.partition_D(tStP0) - tScP_layout = cute.composition(tScS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tScP = cute.make_tensor(tScS.iterator, tScP_layout) - tTMEM_STOREcP = thr_store.partition_S(tScP) - - # --- Vector TMEM (per-row row_sum storage, FMHA pattern) --- - # composition(tStS.layout, (128, 2)) = 2 FP32 columns per logical row - # vec[0] = row_sum (final, after loop), vec[1] = unused - # Reuses S TMEM region (offset 0), free after softmax loop writes - - tStS_vec_layout = cute.composition(tStS.layout, cute.make_layout((128, 2))) - tStS_vec = cute.make_tensor(tStS.iterator + self.tmem_vec_offset, tStS_vec_layout) - tScS_vec_layout = cute.composition(tScS.layout, cute.make_layout((128, 2))) - tScS_vec = cute.make_tensor(tScS.iterator, tScS_vec_layout) - tmem_store_vec_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(2)), self.qk_acc_dtype) - tiled_tmem_store_vec = tcgen05.make_tmem_copy(tmem_store_vec_atom, tStS_vec) - thr_tmem_store_vec = tiled_tmem_store_vec.get_slice(sfw_idx) - tTMEM_STORE_VECtS = thr_tmem_store_vec.partition_D(tStS_vec) - tTMEM_STORE_VECcS = thr_tmem_store_vec.partition_S(tScS_vec) - tmem_load_vec_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(2)), self.qk_acc_dtype) - tiled_tmem_load_vec = tcgen05.make_tmem_copy(tmem_load_vec_atom, tStS_vec) - thr_tmem_load_vec = tiled_tmem_load_vec.get_slice(sfw_idx) - tTMEM_LOAD_VECtS = thr_tmem_load_vec.partition_S(tStS_vec) - tTMEM_LOAD_VECcS = thr_tmem_load_vec.partition_D(tScS_vec) - - # --- C6: O TMEM load/store for rescale (correction_rescale pattern) --- - corr_tile_size = 16 - cO = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1])) - tOcO = pv_thr.partition_C(cO) - o_tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), self.qk_acc_dtype) - o_tmem_store_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), self.qk_acc_dtype) - tOtO_i_layout = cute.composition(tOtO0.layout, cute.make_layout((128, corr_tile_size))) - tOcO_i_layout = cute.composition(tOcO.layout, cute.make_layout((128, corr_tile_size))) - tOtO_i = cute.make_tensor(tOtO0.iterator, tOtO_i_layout) - tOcO_i = cute.make_tensor(tOcO.iterator, tOcO_i_layout) - o_tiled_tmem_load = tcgen05.make_tmem_copy(o_tmem_load_atom, tOtO_i) - o_tiled_tmem_store = tcgen05.make_tmem_copy(o_tmem_store_atom, tOtO_i) - o_thr_load = o_tiled_tmem_load.get_slice(sfw_idx) - o_thr_store = o_tiled_tmem_store.get_slice(sfw_idx) - tTMEM_LOADtO = o_thr_load.partition_S(tOtO_i) - tTMEM_LOADcO = o_thr_load.partition_D(tOcO_i) - tTMEM_STOREtO = o_thr_store.partition_D(tOtO_i) - o_col_tiles = self.pv_mma_tiler[1] // corr_tile_size - - # --- C2: Per-thread row state (persist across KV tiles) --- - row_max = -cutlass.Float32.inf - row_sum = cutlass.Float32(0.0) - - # --- C3: QK scale = 1/sqrt(HEAD_DIM) * log2(e) for exp2 --- - scale = self.scale_softmax_log2 - - # ============================================================= - # Per-KV-tile online softmax loop - # ============================================================= - for kt in range(n_kv_tiles): - si_handle = s_cons.wait_and_advance() - - # Load S from TMEM (FP32, QK C-fragment layout) - tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS) - - # --- C4: Compute tile_max via .reduce(MAX) --- - old_row_max = row_max - row_max = tTMEM_LOADrS.load().reduce(cute.ReductionOp.MAX, row_max, 0) - row_max_safe = row_max - if row_max == -cutlass.Float32.inf: - row_max_safe = cutlass.Float32(0.0) - - # --- C5: Compute rescale factor --- - acc_scale = cute.math.exp2(scale * (old_row_max - row_max_safe), fastmath=True) - - # --- C6: Rescale O in TMEM (load O, multiply by acc_scale, store O) --- - # acc_scale belongs to QK row (N//4), but O rows are in PV partition (N). - # Store acc_scale to vector by QK row, read by PV row. - if kt > 0: - pv_done_bar.arrive_and_wait() - - # Store acc_scale to vector indexed by QK logical row - qk_row_c6 = tTMEM_LOADcS[0][0] - thr_vs_c6 = tiled_tmem_store_vec.get_slice(qk_row_c6) - tVStore_c6 = thr_vs_c6.partition_D(tStS_vec) - tVStoreSrc_c6 = thr_vs_c6.partition_S(tScS_vec) - tVStoreRmem_c6 = cute.make_rmem_tensor(tVStoreSrc_c6.shape, self.qk_acc_dtype) - tVStoreRmem_c6[0] = acc_scale - cute.copy(tiled_tmem_store_vec, tVStoreRmem_c6, tVStore_c6) - cute.arch.fence_view_async_tmem_store() - - # Read acc_scale from vector indexed by PV logical row - pv_row_c6 = tTMEM_LOADcO[0][0] - thr_vl_c6 = tiled_tmem_load_vec.get_slice(pv_row_c6) - tVLoad_c6 = thr_vl_c6.partition_S(tStS_vec) - tVLoadDst_c6 = thr_vl_c6.partition_D(tScS_vec) - tVLoadRmem_c6 = cute.make_rmem_tensor(tVLoadDst_c6.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load_vec, tVLoad_c6, tVLoadRmem_c6) - cute.arch.fence_view_async_tmem_load() - acc_scale_pv = tVLoadRmem_c6[0] - - tTMrO = cute.make_rmem_tensor((tTMEM_LOADcO.shape, o_col_tiles), self.qk_acc_dtype) - for i in range(o_col_tiles): - tTMrO_i_ = tTMrO[None, i] - tTMrO_i_layout = cute.composition(tTMrO_i_.layout, cute.make_layout(tTMrO.shape[0])) - tTMrO_i = cute.make_tensor(tTMrO_i_.iterator, tTMrO_i_layout) - tTMEM_LOADtO_i = cute.make_tensor(tTMEM_LOADtO.iterator + i * corr_tile_size, tTMEM_LOADtO.layout) - tTMEM_STOREtO_i = cute.make_tensor(tTMEM_STOREtO.iterator + i * corr_tile_size, tTMEM_STOREtO.layout) - cute.copy(o_tiled_tmem_load, tTMEM_LOADtO_i, tTMrO_i) - for j in cutlass.range(cute.size(tTMrO_i), vectorize=True): - tTMrO_i[j] = tTMrO_i[j] * acc_scale_pv - cute.copy(o_tiled_tmem_store, tTMrO_i, tTMEM_STOREtO_i) - cute.arch.fence_view_async_tmem_store() - - # Rescale row_sum - row_sum = row_sum * acc_scale - - # --- C7: Compute P = exp2((S - row_max_safe) * scale) --- - minus_row_max_scale = (cutlass.Float32(0.0) - row_max_safe) * scale - - # Register bridge (FMHA pattern: FP32 backing + BF16 view) - rP_words = cute.make_rmem_tensor(tTMEM_STOREcP.shape, self.qk_acc_dtype) - rP_bf16 = cute.make_tensor(cute.recast_ptr(rP_words.iterator, dtype=self.q_dtype), tTMEM_LOADrS.layout) - - frg_cnt = 4 - frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt - tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile)) - rP_bf16_frg = cute.logical_divide(rP_bf16, cute.make_layout(frg_tile)) - - # Scale S, compute exp2, store through register bridge - for j in range(frg_cnt): - for k in cutlass.range(cute.size(tTMEM_LOADrS_frg, mode=[0]), vectorize=True): - tTMEM_LOADrS_frg[k, j] = tTMEM_LOADrS_frg[k, j] * scale + minus_row_max_scale - tTMEM_LOADrS_frg[k, j] = cute.math.exp2(tTMEM_LOADrS_frg[k, j], fastmath=True) - s_vec = tTMEM_LOADrS_frg[None, j].load() - rP_bf16_frg[None, j].store(s_vec.to(self.q_dtype)) - - # Store P to TMEM - cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP) - cute.arch.fence_view_async_tmem_store() - si_handle.release() - softmax_done_bar.arrive() - - # --- C8: Row sum accumulation (CUTLASS FMHA packed f32x2 pattern) --- - # P values still in tTMEM_LOADrS registers. - # 4 accumulators for 4 reduction_unroll columns. - local_row_sum_0 = (cutlass.Float32(0.0), cutlass.Float32(0.0)) - local_row_sum_1 = (cutlass.Float32(0.0), cutlass.Float32(0.0)) - local_row_sum_2 = (cutlass.Float32(0.0), cutlass.Float32(0.0)) - local_row_sum_3 = (cutlass.Float32(0.0), cutlass.Float32(0.0)) - - reduction_unroll = 4 - rfrg_tile = cute.size(tTMEM_LOADrS) // reduction_unroll - tTMEM_LOADrS_rfrg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(rfrg_tile)) - - for j in cutlass.range_constexpr(0, cute.size(tTMEM_LOADrS_rfrg, mode=[0]), 2): - local_row_sum_0 = cute.arch.add_packed_f32x2( - local_row_sum_0, (tTMEM_LOADrS_rfrg[j, 0], tTMEM_LOADrS_rfrg[j + 1, 0])) - local_row_sum_1 = cute.arch.add_packed_f32x2( - local_row_sum_1, (tTMEM_LOADrS_rfrg[j, 1], tTMEM_LOADrS_rfrg[j + 1, 1])) - local_row_sum_2 = cute.arch.add_packed_f32x2( - local_row_sum_2, (tTMEM_LOADrS_rfrg[j, 2], tTMEM_LOADrS_rfrg[j + 1, 2])) - local_row_sum_3 = cute.arch.add_packed_f32x2( - local_row_sum_3, (tTMEM_LOADrS_rfrg[j, 3], tTMEM_LOADrS_rfrg[j + 1, 3])) - - local_row_sum_0 = cute.arch.add_packed_f32x2(local_row_sum_0, local_row_sum_1) - local_row_sum_2 = cute.arch.add_packed_f32x2(local_row_sum_2, local_row_sum_3) - local_row_sum_0 = cute.arch.add_packed_f32x2(local_row_sum_0, local_row_sum_2) - tile_sum = local_row_sum_0[0] + local_row_sum_0[1] - - row_sum = row_sum + tile_sum - - # --- C9: Final normalization via O TMEM rescale --- - pv_done_bar.arrive_and_wait() - - # Compute inv_row_sum from P in TMEM using PV partition. - # P was stored by softmax loop into TMEM at offset tmem_p0_offset. - # PV partition maps thread N to PV row N, so reading P via PV partition - # gives the correct per-row P values to sum. - # This avoids the QK→PV row mapping mismatch (QK: N->N//4, PV: N->N). - - # P is stored as BF16 in TMEM at tmem_p0_offset. - # We need to read it via PV TMEM load and sum the values. - # P has shape (128, HEAD_DIM//2) in FP32 columns (64 BF16 = 32 FP32 cols). - # Use the P TMEM load partition (PV A-fragment read). - - # Actually, P was stored via QK C-fragment store (St32x32bOp Repetition(32)). - # To read it via PV partition, we need a PV-partitioned load from the P region. - # Let's use the same o_tiled_tmem_load but pointed at P's TMEM offset. - - # P occupies TMEM columns [tmem_p0_offset, tmem_p0_offset + p_cols_fp32) - # In the PV C-fragment, P is the A-fragment. We can use tOrP0's layout. - # tOrP0 was set up with offset for PV MMA read. - - # Simpler: sum O across columns to get unnormalized row sum, then normalize. - # For V=identity, O = P@V = sum(P per row). So O.sum(dim=-1) = row_sum. - # For arbitrary V, O = P@V. O.sum(dim=-1) = sum_j(P@V)[j] = sum_j(sum_i P[i]*V[i,j]) - # This is NOT sum(P). So this trick only works for V=identity. - - # Correct approach: read P from TMEM, sum it per PV row. - # P is at TMEM offset tmem_p0_offset, stored as BF16 with St32x32bOp. - # P shape in TMEM: 128 rows x (HEAD_DIM BF16 = 32 FP32 cols) - # We can read P using Ld32x32bOp(Repetition(corr_tile_size)) via PV O-partition. - - # Use PV O TMEM load to read from P region instead of O region - p_col_tiles = p_cols_fp32 // corr_tile_size # 32 // 16 = 2 - pv_row_sum = cutlass.Float32(0.0) - for i in range(p_col_tiles): - # Read P tile from TMEM at P offset (not O offset) - tTMEM_LOADtP_i = cute.make_tensor( - tTMEM_LOADtO.iterator + (self.tmem_p0_offset - self.tmem_o0_offset) + i * corr_tile_size, - tTMEM_LOADtO.layout) - tTMrP_i = cute.make_rmem_tensor(tTMEM_LOADcO.shape, self.qk_acc_dtype) - cute.copy(o_tiled_tmem_load, tTMEM_LOADtP_i, tTMrP_i) - # Use .reduce(SUM) instead of scalar accumulation (vectorizer can't handle scalar in vectorized loop) - tile_p_sum = tTMrP_i.load().reduce(cute.ReductionOp.ADD, cutlass.Float32(0.0), 0) - pv_row_sum = pv_row_sum + tile_p_sum - - inv_row_sum = cutlass.Float32(1.0) / pv_row_sum - - # Normalize O in TMEM using PV-correct inv_row_sum - tTMrO_final = cute.make_rmem_tensor((tTMEM_LOADcO.shape, o_col_tiles), self.qk_acc_dtype) - for i in range(o_col_tiles): - tTMrO_i_ = tTMrO_final[None, i] - tTMrO_i_layout = cute.composition(tTMrO_i_.layout, cute.make_layout(tTMrO_final.shape[0])) - tTMrO_i = cute.make_tensor(tTMrO_i_.iterator, tTMrO_i_layout) - tTMEM_LOADtO_i = cute.make_tensor( - tTMEM_LOADtO.iterator + i * corr_tile_size, tTMEM_LOADtO.layout) - tTMEM_STOREtO_i = cute.make_tensor( - tTMEM_STOREtO.iterator + i * corr_tile_size, tTMEM_STOREtO.layout) - cute.copy(o_tiled_tmem_load, tTMEM_LOADtO_i, tTMrO_i) - for j in cutlass.range(cute.size(tTMrO_i), vectorize=True): - tTMrO_i[j] = tTMrO_i[j] * inv_row_sum - cute.copy(o_tiled_tmem_store, tTMrO_i, tTMEM_STOREtO_i) - cute.arch.fence_view_async_tmem_store() - - # Now O in TMEM is normalized. Use standard epilogue_tma_store with identity. - tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout) - acc_cons_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage) - c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) - c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp) - acc_cons_st = utils.gemm.sm100.epilogue_tma_store( - self, tidx, warp_idx, tma_c, tCtO_base, sC, tCgC, epi_tile, 0, - const_expr(lambda x: x), - (0,0,0), acc_cons_st, acc_pipe, c_pipe) - c_pipe.producer_tail() - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def test(): - import math - torch.manual_seed(42) - for n in [128, 256, 384]: - m, hd = 128, HEAD_DIM - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device="cuda") - k = torch.randn(n, hd, 1, dtype=torch.bfloat16, device="cuda") - v = torch.randn(n, hd, dtype=torch.bfloat16, device="cuda") - v_kernel = v.unsqueeze(-1) - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device="cuda") - qf = q[:,:,0].float(); kf = k[:,:,0].float() - attn = qf @ kf.T / math.sqrt(hd) - ref = torch.softmax(attn, dim=-1) @ v.float() - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - kernel = FmhaV3Softmax() - print(f"n={n}: Compiling...", flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - print(f"n={n}: tmem: s0={kernel.tmem_s0_offset} p0={kernel.tmem_p0_offset} o0={kernel.tmem_o0_offset} vec={kernel.tmem_vec_offset} alloc={kernel.num_tmem_alloc_cols}", flush=True) - print(f"n={n}: Running...", flush=True) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - out = c[:,:,0].float() - cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - max_err = (out - ref).abs().max().item() - print(f"FMHA softmax n={n}: cosine {cos:.6f} max_err {max_err:.6f} {'PASS' if cos >= 0.999 else 'FAIL'}", flush=True) - -if __name__ == "__main__": - test() - - diff --git a/tests/archive/unit_test_fmha_v3_softmax.py b/tests/archive/unit_test_fmha_v3_softmax.py deleted file mode 100644 index fb3b98ae..00000000 --- a/tests/archive/unit_test_fmha_v3_softmax.py +++ /dev/null @@ -1,511 +0,0 @@ -""" -FMHA v3 + Stage C: QK -> online softmax -> PV with KV-tile interleaving. -Stage C: row_max, exp2, O rescale, row_sum, final normalization. -FMHA pattern P store preserved from Stage B. -""" -import math -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - -HEAD_DIM = 64 - -class FmhaV3Softmax: - def __init__(self): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0,1,2,3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.num_c_stage = 2 - self.kv_stage = 2; self.q_stage = 1; self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (128, 128, qk_ik * 4) - pv_ik = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (128, HEAD_DIM, pv_ik * (128 // pv_ik)) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0]//cute.size(qk_mma.thr_id.shape), HEAD_DIM, self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - self.q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, self.q_dtype, self.q_stage) - self.k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, self.q_dtype, self.kv_stage) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, self.kv_stage) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - qk_thr = qk_mma.get_slice(0); qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - pv_thr = pv_mma.get_slice(0); pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - self.tmem_s0_offset = 0; self.tmem_p0_offset = 32 - # P occupies [tmem_p0_offset, tmem_p0_offset + p_cols_fp32) - # S occupies [0, qk_mma_tiler[1]) = [0, 128) - # O must NOT overlap P. Place O after max(S end, P end), aligned to 32. - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - p_end = self.tmem_p0_offset + p_cols_fp32 # 32 + 64 = 96 - s_cols = self.qk_mma_tiler[1] # 128 - o_after = max(s_cols, p_end) # 128 - self.tmem_o0_offset = ((o_after + 31) // 32) * 32 - self.tmem_vec_offset = 0 # Reuse S region for per-row inv_row_sum vector # align to 32 = 128 - self.tmem_vec_offset = 0 # Reuse S region (free after softmax loop) - o_cols = find_tmem_tensor_col_offset(tOtO) # footprint of O - total = self.tmem_o0_offset + o_cols - # Must be multiple of 32 AND power of 2 - self.num_tmem_alloc_cols = 1 - while self.num_tmem_alloc_cols < total: - self.num_tmem_alloc_cols *= 2 - cta = cute.size(qk_mma.thr_id.shape) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)) - self.q_tx_bytes = cute.size_in_bytes(self.q_dtype, q_s) * cta - self.kv_tx_bytes = cute.size_in_bytes(self.q_dtype, k_s) * cta - self.scale_softmax_log2 = Float32(1.0 / math.sqrt(HEAD_DIM) * math.log2(math.e)) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - # # s_k hardcoded # BROKEN in @cute.jit - # FMHA-style V: reconstruct as (HEAD_DIM, s_k, 1) MN-major - v_fmha = cute.make_tensor( - v.iterator, - cute.make_layout( - (HEAD_DIM, 128, 1), - stride=(1, HEAD_DIM, HEAD_DIM * 128), - ), - ) - self.v_major = LayoutEnum.from_tensor(v_fmha).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, self.a_major, self.b_major, self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, self.qk_acc_dtype, self.cta_group, (128,HEAD_DIM), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)); v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id),q,q_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id),k,k_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v_fmha,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape) - epi_s = cute.select(self.c_smem_s,mode=[0,1]) - tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile) - self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.v_smem_s,self.p_tmem_s,self.c_smem_s,self.epi_tile).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, q_smem_s, k_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx,_,_ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k); cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - q_bar: cute.struct.MemRange[cutlass.Int64, self.q_stage*2] - kv_bar: cute.struct.MemRange[cutlass.Int64, self.kv_stage*2] - s_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage*2] - tmem_dealloc: cutlass.Int64; holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - qp,qc = pipeline.PipelineTmaUmma.create(barrier_storage=st.q_bar.data_ptr(),num_stages=self.q_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.q_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - kvp,kvc = pipeline.PipelineTmaUmma.create(barrier_storage=st.kv_bar.data_ptr(),num_stages=self.kv_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.kv_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - s_prod,s_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.s_bar.data_ptr(),num_stages=1,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.epilogue_warp_id))).make_participants() - softmax_done_bar = pipeline.NamedBarrier(barrier_id=3, num_threads=32 + 32*len(self.epilogue_warp_id)) - pv_done_bar = pipeline.NamedBarrier(barrier_id=4, num_threads=32 + 32*len(self.epilogue_warp_id)) - acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(),num_stages=self.num_acc_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,len(self.epilogue_warp_id)),cta_layout_vmnk=cl_vmnk,defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2,num_threads=32*len((self.mma_warp_id,*self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr,barrier_for_retrieve=tmem_bar,allocator_warp_id=self.epilogue_warp_id[0],is_two_cta=cute.size(qk_mma.thr_id.shape)==2,two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk,is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype,layout=q_smem_s.outer,byte_alignment=128,swizzle=q_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype,layout=k_smem_s.outer,byte_alignment=128,swizzle=k_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype,layout=v_smem_s.outer,byte_alignment=128,swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype,layout=c_smem_s.outer,byte_alignment=128,swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None)) - gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None)) - gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None)) - gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None)) - n_kv_tiles = cute.size(gK, mode=[3]) - - qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - tCgV = pv_thr.partition_B(gV); tCgC = pv_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,0,None,0)).shape) - tAsQ,tAgQ = cpasync.tma_partition(tma_q,0,a_lay,cute.group_modes(sQ,0,3),cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape) - tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3)) - tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)]; tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - # --- PV read view (for MMA only, NOT for softmax store) --- - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None,None,None,0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # TMA LOAD - if warp_idx == self.tma_warp_id: - qp.reset(); qh = qp.acquire_and_advance() - cute.copy(tma_q,tAgQ[(None,qh.count)],tAsQ[(None,qh.index)],tma_bar_ptr=qh.barrier) - qp.tail() - kvp.reset(); pk = kvp.try_acquire() - for kt in cutlass.range(n_kv_tiles,unroll=1): - kh = kvp.acquire_and_advance(pk) - cute.copy(tma_k,tBgK[(None,kh.count)],tBsK[(None,kh.index)],tma_bar_ptr=kh.barrier) - pk = cutlass.Boolean(1) - vh = kvp.acquire_and_advance(pk) - cute.copy(tma_v,tVgV[(None,vh.count)],tVsV[(None,vh.index)],tma_bar_ptr=vh.barrier) - pk = cutlass.Boolean(1) - kvp.tail() - - # MMA - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - qc.reset(); qh = qc.wait_and_advance(); qh.release() - kvc.reset(); pk = kvc.try_wait() - acc_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_stage) - acc_pipe.producer_acquire(acc_st) - for kt in range(n_kv_tiles): - kh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - sh = s_prod.acquire_and_advance() - qk_mma.set(tcgen05.Field.ACCUMULATE, False) - for kb in cutlass.range(cute.size(tCrQ,mode=[2]), unroll_full=True): - cute.gemm(qk_mma, tStS0, tCrQ[(None,None,kb,0)], tCrK[(None,None,kb,kh.index)], tStS0) - qk_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - sh.commit(); kh.release() - softmax_done_bar.arrive_and_wait() - vh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - pv_mma.set(tcgen05.Field.ACCUMULATE, kt != 0) - for kb in cutlass.range(cute.size(tOrP0,mode=[2]), unroll_full=True): - cute.gemm(pv_mma, tOtO0, tOrP0[(None,None,kb)], tCrV[(None,None,kb,vh.index)], tOtO0) - pv_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - vh.release() - pv_done_bar.arrive() - acc_pipe.producer_commit(acc_st); acc_st.advance() - acc_pipe.producer_tail(acc_st) - - # ===================== EPILOGUE WARPS (STAGE C: ONLINE SOFTMAX) ===================== - if warp_idx < self.mma_warp_id: - tmem.allocate(self.num_tmem_alloc_cols) - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) - sfw_idx = tidx % (32 * len(self.epilogue_warp_id)) - - # --- S load (QK C-fragment) --- - tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS0) - thr_load = tiled_tmem_load.get_slice(sfw_idx) - tTMEM_LOADtS = thr_load.partition_S(tStS0) - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tTMEM_LOADcS = thr_load.partition_D(tScS) - - # --- P store (QK C-fragment composition, FMHA pattern) --- - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - tStP_layout = cute.composition(tStS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tStP0 = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStP_layout) - tmem_store_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStP0) - thr_store = tiled_tmem_store.get_slice(sfw_idx) - tTMEM_STOREtP = thr_store.partition_D(tStP0) - tScP_layout = cute.composition(tScS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tScP = cute.make_tensor(tScS.iterator, tScP_layout) - tTMEM_STOREcP = thr_store.partition_S(tScP) - - # --- Vector TMEM (per-row row_sum storage, FMHA pattern) --- - # composition(tStS.layout, (128, 2)) = 2 FP32 columns per logical row - # vec[0] = row_sum (final, after loop), vec[1] = unused - # Reuses S TMEM region (offset 0), free after softmax loop writes - - tStS_vec_layout = cute.composition(tStS.layout, cute.make_layout((128, 2))) - tStS_vec = cute.make_tensor(tStS.iterator + self.tmem_vec_offset, tStS_vec_layout) - tScS_vec_layout = cute.composition(tScS.layout, cute.make_layout((128, 2))) - tScS_vec = cute.make_tensor(tScS.iterator, tScS_vec_layout) - tmem_store_vec_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(2)), self.qk_acc_dtype) - tiled_tmem_store_vec = tcgen05.make_tmem_copy(tmem_store_vec_atom, tStS_vec) - thr_tmem_store_vec = tiled_tmem_store_vec.get_slice(sfw_idx) - tTMEM_STORE_VECtS = thr_tmem_store_vec.partition_D(tStS_vec) - tTMEM_STORE_VECcS = thr_tmem_store_vec.partition_S(tScS_vec) - tmem_load_vec_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(2)), self.qk_acc_dtype) - tiled_tmem_load_vec = tcgen05.make_tmem_copy(tmem_load_vec_atom, tStS_vec) - thr_tmem_load_vec = tiled_tmem_load_vec.get_slice(sfw_idx) - tTMEM_LOAD_VECtS = thr_tmem_load_vec.partition_S(tStS_vec) - tTMEM_LOAD_VECcS = thr_tmem_load_vec.partition_D(tScS_vec) - - # --- C6: O TMEM load/store for rescale (correction_rescale pattern) --- - corr_tile_size = 16 - cO = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1])) - tOcO = pv_thr.partition_C(cO) - o_tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), self.qk_acc_dtype) - o_tmem_store_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), self.qk_acc_dtype) - tOtO_i_layout = cute.composition(tOtO0.layout, cute.make_layout((128, corr_tile_size))) - tOcO_i_layout = cute.composition(tOcO.layout, cute.make_layout((128, corr_tile_size))) - tOtO_i = cute.make_tensor(tOtO0.iterator, tOtO_i_layout) - tOcO_i = cute.make_tensor(tOcO.iterator, tOcO_i_layout) - o_tiled_tmem_load = tcgen05.make_tmem_copy(o_tmem_load_atom, tOtO_i) - o_tiled_tmem_store = tcgen05.make_tmem_copy(o_tmem_store_atom, tOtO_i) - o_thr_load = o_tiled_tmem_load.get_slice(sfw_idx) - o_thr_store = o_tiled_tmem_store.get_slice(sfw_idx) - tTMEM_LOADtO = o_thr_load.partition_S(tOtO_i) - tTMEM_LOADcO = o_thr_load.partition_D(tOcO_i) - tTMEM_STOREtO = o_thr_store.partition_D(tOtO_i) - o_col_tiles = self.pv_mma_tiler[1] // corr_tile_size - - # --- C2: Per-thread row state (persist across KV tiles) --- - row_max = -cutlass.Float32.inf - row_sum = cutlass.Float32(0.0) - - # --- C3: QK scale = 1/sqrt(HEAD_DIM) * log2(e) for exp2 --- - scale = self.scale_softmax_log2 - - # ============================================================= - # Per-KV-tile online softmax loop - # ============================================================= - for kt in range(n_kv_tiles): - si_handle = s_cons.wait_and_advance() - - # Load S from TMEM (FP32, QK C-fragment layout) - tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS) - - # --- C4: Compute tile_max via .reduce(MAX) --- - old_row_max = row_max - row_max = tTMEM_LOADrS.load().reduce(cute.ReductionOp.MAX, row_max, 0) - row_max_safe = row_max - if row_max == -cutlass.Float32.inf: - row_max_safe = cutlass.Float32(0.0) - - # --- C5: Compute rescale factor --- - acc_scale = cute.math.exp2(scale * (old_row_max - row_max_safe), fastmath=True) - - # --- C6: Rescale O in TMEM (load O, multiply by acc_scale, store O) --- - # acc_scale belongs to QK row (N//4), but O rows are in PV partition (N). - # Store acc_scale to vector by QK row, read by PV row. - if kt > 0: - pv_done_bar.arrive_and_wait() - - # Store acc_scale to vector indexed by QK logical row - qk_row_c6 = tTMEM_LOADcS[0][0] - thr_vs_c6 = tiled_tmem_store_vec.get_slice(qk_row_c6) - tVStore_c6 = thr_vs_c6.partition_D(tStS_vec) - tVStoreSrc_c6 = thr_vs_c6.partition_S(tScS_vec) - tVStoreRmem_c6 = cute.make_rmem_tensor(tVStoreSrc_c6.shape, self.qk_acc_dtype) - tVStoreRmem_c6[0] = acc_scale - cute.copy(tiled_tmem_store_vec, tVStoreRmem_c6, tVStore_c6) - cute.arch.fence_view_async_tmem_store() - - # Read acc_scale from vector indexed by PV logical row - pv_row_c6 = tTMEM_LOADcO[0][0] - thr_vl_c6 = tiled_tmem_load_vec.get_slice(pv_row_c6) - tVLoad_c6 = thr_vl_c6.partition_S(tStS_vec) - tVLoadDst_c6 = thr_vl_c6.partition_D(tScS_vec) - tVLoadRmem_c6 = cute.make_rmem_tensor(tVLoadDst_c6.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load_vec, tVLoad_c6, tVLoadRmem_c6) - cute.arch.fence_view_async_tmem_load() - acc_scale_pv = tVLoadRmem_c6[0] - - tTMrO = cute.make_rmem_tensor((tTMEM_LOADcO.shape, o_col_tiles), self.qk_acc_dtype) - for i in range(o_col_tiles): - tTMrO_i_ = tTMrO[None, i] - tTMrO_i_layout = cute.composition(tTMrO_i_.layout, cute.make_layout(tTMrO.shape[0])) - tTMrO_i = cute.make_tensor(tTMrO_i_.iterator, tTMrO_i_layout) - tTMEM_LOADtO_i = cute.make_tensor(tTMEM_LOADtO.iterator + i * corr_tile_size, tTMEM_LOADtO.layout) - tTMEM_STOREtO_i = cute.make_tensor(tTMEM_STOREtO.iterator + i * corr_tile_size, tTMEM_STOREtO.layout) - cute.copy(o_tiled_tmem_load, tTMEM_LOADtO_i, tTMrO_i) - for j in cutlass.range(cute.size(tTMrO_i), vectorize=True): - tTMrO_i[j] = tTMrO_i[j] * acc_scale_pv - cute.copy(o_tiled_tmem_store, tTMrO_i, tTMEM_STOREtO_i) - cute.arch.fence_view_async_tmem_store() - - # Rescale row_sum - row_sum = row_sum * acc_scale - - # --- C7: Compute P = exp2((S - row_max_safe) * scale) --- - minus_row_max_scale = (cutlass.Float32(0.0) - row_max_safe) * scale - - # Register bridge (FMHA pattern: FP32 backing + BF16 view) - rP_words = cute.make_rmem_tensor(tTMEM_STOREcP.shape, self.qk_acc_dtype) - rP_bf16 = cute.make_tensor(cute.recast_ptr(rP_words.iterator, dtype=self.q_dtype), tTMEM_LOADrS.layout) - - frg_cnt = 4 - frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt - tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile)) - rP_bf16_frg = cute.logical_divide(rP_bf16, cute.make_layout(frg_tile)) - - # Scale S, compute exp2, store through register bridge - for j in range(frg_cnt): - for k in cutlass.range(cute.size(tTMEM_LOADrS_frg, mode=[0]), vectorize=True): - tTMEM_LOADrS_frg[k, j] = tTMEM_LOADrS_frg[k, j] * scale + minus_row_max_scale - tTMEM_LOADrS_frg[k, j] = cute.math.exp2(tTMEM_LOADrS_frg[k, j], fastmath=True) - s_vec = tTMEM_LOADrS_frg[None, j].load() - rP_bf16_frg[None, j].store(s_vec.to(self.q_dtype)) - - # Store P to TMEM - cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP) - cute.arch.fence_view_async_tmem_store() - si_handle.release() - softmax_done_bar.arrive() - - # --- C8: Row sum accumulation (CUTLASS FMHA packed f32x2 pattern) --- - # P values still in tTMEM_LOADrS registers. - # 4 accumulators for 4 reduction_unroll columns. - local_row_sum_0 = (cutlass.Float32(0.0), cutlass.Float32(0.0)) - local_row_sum_1 = (cutlass.Float32(0.0), cutlass.Float32(0.0)) - local_row_sum_2 = (cutlass.Float32(0.0), cutlass.Float32(0.0)) - local_row_sum_3 = (cutlass.Float32(0.0), cutlass.Float32(0.0)) - - reduction_unroll = 4 - rfrg_tile = cute.size(tTMEM_LOADrS) // reduction_unroll - tTMEM_LOADrS_rfrg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(rfrg_tile)) - - for j in cutlass.range_constexpr(0, cute.size(tTMEM_LOADrS_rfrg, mode=[0]), 2): - local_row_sum_0 = cute.arch.add_packed_f32x2( - local_row_sum_0, (tTMEM_LOADrS_rfrg[j, 0], tTMEM_LOADrS_rfrg[j + 1, 0])) - local_row_sum_1 = cute.arch.add_packed_f32x2( - local_row_sum_1, (tTMEM_LOADrS_rfrg[j, 1], tTMEM_LOADrS_rfrg[j + 1, 1])) - local_row_sum_2 = cute.arch.add_packed_f32x2( - local_row_sum_2, (tTMEM_LOADrS_rfrg[j, 2], tTMEM_LOADrS_rfrg[j + 1, 2])) - local_row_sum_3 = cute.arch.add_packed_f32x2( - local_row_sum_3, (tTMEM_LOADrS_rfrg[j, 3], tTMEM_LOADrS_rfrg[j + 1, 3])) - - local_row_sum_0 = cute.arch.add_packed_f32x2(local_row_sum_0, local_row_sum_1) - local_row_sum_2 = cute.arch.add_packed_f32x2(local_row_sum_2, local_row_sum_3) - local_row_sum_0 = cute.arch.add_packed_f32x2(local_row_sum_0, local_row_sum_2) - tile_sum = local_row_sum_0[0] + local_row_sum_0[1] - - row_sum = row_sum + tile_sum - - # --- C9: Final normalization via O TMEM rescale --- - pv_done_bar.arrive_and_wait() - - # Compute inv_row_sum from P in TMEM using PV partition. - # P was stored by softmax loop into TMEM at offset tmem_p0_offset. - # PV partition maps thread N to PV row N, so reading P via PV partition - # gives the correct per-row P values to sum. - # This avoids the QK→PV row mapping mismatch (QK: N->N//4, PV: N->N). - - # P is stored as BF16 in TMEM at tmem_p0_offset. - # We need to read it via PV TMEM load and sum the values. - # P has shape (128, HEAD_DIM//2) in FP32 columns (64 BF16 = 32 FP32 cols). - # Use the P TMEM load partition (PV A-fragment read). - - # Actually, P was stored via QK C-fragment store (St32x32bOp Repetition(32)). - # To read it via PV partition, we need a PV-partitioned load from the P region. - # Let's use the same o_tiled_tmem_load but pointed at P's TMEM offset. - - # P occupies TMEM columns [tmem_p0_offset, tmem_p0_offset + p_cols_fp32) - # In the PV C-fragment, P is the A-fragment. We can use tOrP0's layout. - # tOrP0 was set up with offset for PV MMA read. - - # Simpler: sum O across columns to get unnormalized row sum, then normalize. - # For V=identity, O = P@V = sum(P per row). So O.sum(dim=-1) = row_sum. - # For arbitrary V, O = P@V. O.sum(dim=-1) = sum_j(P@V)[j] = sum_j(sum_i P[i]*V[i,j]) - # This is NOT sum(P). So this trick only works for V=identity. - - # Correct approach: read P from TMEM, sum it per PV row. - # P is at TMEM offset tmem_p0_offset, stored as BF16 with St32x32bOp. - # P shape in TMEM: 128 rows x (HEAD_DIM BF16 = 32 FP32 cols) - # We can read P using Ld32x32bOp(Repetition(corr_tile_size)) via PV O-partition. - - # Use PV O TMEM load to read from P region instead of O region - p_col_tiles = p_cols_fp32 // corr_tile_size # 32 // 16 = 2 - pv_row_sum = cutlass.Float32(0.0) - for i in range(p_col_tiles): - # Read P tile from TMEM at P offset (not O offset) - tTMEM_LOADtP_i = cute.make_tensor( - tTMEM_LOADtO.iterator + (self.tmem_p0_offset - self.tmem_o0_offset) + i * corr_tile_size, - tTMEM_LOADtO.layout) - tTMrP_i = cute.make_rmem_tensor(tTMEM_LOADcO.shape, self.qk_acc_dtype) - cute.copy(o_tiled_tmem_load, tTMEM_LOADtP_i, tTMrP_i) - # Use .reduce(SUM) instead of scalar accumulation (vectorizer can't handle scalar in vectorized loop) - tile_p_sum = tTMrP_i.load().reduce(cute.ReductionOp.ADD, cutlass.Float32(0.0), 0) - pv_row_sum = pv_row_sum + tile_p_sum - - inv_row_sum = cutlass.Float32(1.0) / pv_row_sum - - # Normalize O in TMEM using PV-correct inv_row_sum - tTMrO_final = cute.make_rmem_tensor((tTMEM_LOADcO.shape, o_col_tiles), self.qk_acc_dtype) - for i in range(o_col_tiles): - tTMrO_i_ = tTMrO_final[None, i] - tTMrO_i_layout = cute.composition(tTMrO_i_.layout, cute.make_layout(tTMrO_final.shape[0])) - tTMrO_i = cute.make_tensor(tTMrO_i_.iterator, tTMrO_i_layout) - tTMEM_LOADtO_i = cute.make_tensor( - tTMEM_LOADtO.iterator + i * corr_tile_size, tTMEM_LOADtO.layout) - tTMEM_STOREtO_i = cute.make_tensor( - tTMEM_STOREtO.iterator + i * corr_tile_size, tTMEM_STOREtO.layout) - cute.copy(o_tiled_tmem_load, tTMEM_LOADtO_i, tTMrO_i) - for j in cutlass.range(cute.size(tTMrO_i), vectorize=True): - tTMrO_i[j] = tTMrO_i[j] * inv_row_sum - cute.copy(o_tiled_tmem_store, tTMrO_i, tTMEM_STOREtO_i) - cute.arch.fence_view_async_tmem_store() - - # Now O in TMEM is normalized. Use standard epilogue_tma_store with identity. - tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout) - acc_cons_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage) - c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) - c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp) - acc_cons_st = utils.gemm.sm100.epilogue_tma_store( - self, tidx, warp_idx, tma_c, tCtO_base, sC, tCgC, epi_tile, 0, - const_expr(lambda x: x), - (0,0,0), acc_cons_st, acc_pipe, c_pipe) - c_pipe.producer_tail() - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def test(): - import math - torch.manual_seed(42) - for n in [128, 256, 384]: - m, hd = 128, HEAD_DIM - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device="cuda") - k = torch.randn(n, hd, 1, dtype=torch.bfloat16, device="cuda") - v = torch.randn(n, hd, dtype=torch.bfloat16, device="cuda") - v_kernel = v.unsqueeze(-1) - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device="cuda") - qf = q[:,:,0].float(); kf = k[:,:,0].float() - attn = qf @ kf.T / math.sqrt(hd) - ref = torch.softmax(attn, dim=-1) @ v.float() - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - kernel = FmhaV3Softmax() - print(f"n={n}: Compiling...", flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - print(f"n={n}: tmem: s0={kernel.tmem_s0_offset} p0={kernel.tmem_p0_offset} o0={kernel.tmem_o0_offset} vec={kernel.tmem_vec_offset} alloc={kernel.num_tmem_alloc_cols}", flush=True) - print(f"n={n}: Running...", flush=True) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - out = c[:,:,0].float() - cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - max_err = (out - ref).abs().max().item() - print(f"FMHA softmax n={n}: cosine {cos:.6f} max_err {max_err:.6f} {'PASS' if cos >= 0.999 else 'FAIL'}", flush=True) - -if __name__ == "__main__": - test() - - diff --git a/tests/archive/unit_test_fmha_v3_stage_c.py b/tests/archive/unit_test_fmha_v3_stage_c.py deleted file mode 100644 index 69a7ab33..00000000 --- a/tests/archive/unit_test_fmha_v3_stage_c.py +++ /dev/null @@ -1,327 +0,0 @@ -""" -FMHA v3: QK -> softmax -> PV with KV-tile interleaving. -Bug 4b fix (FMHA pattern): P store uses QK C-fragment layout composition, -NOT PV A-fragment layout. Register bridge: FP32 backing (store partition shape) -recast to BF16 view (QK-load layout). -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - -HEAD_DIM = 64 - -class FmhaV3: - def __init__(self): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0,1,2,3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.num_c_stage = 2 - self.kv_stage = 2; self.q_stage = 1; self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (128, 128, qk_ik * 4) - pv_ik = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (128, HEAD_DIM, pv_ik * (128 // pv_ik)) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0]//cute.size(qk_mma.thr_id.shape), HEAD_DIM, self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - self.q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, self.q_dtype, self.q_stage) - self.k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, self.q_dtype, self.kv_stage) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, self.kv_stage) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - qk_thr = qk_mma.get_slice(0); qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - pv_thr = pv_mma.get_slice(0); pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - self.tmem_s0_offset = 0; self.tmem_p0_offset = 32 - # P occupies [tmem_p0_offset, tmem_p0_offset + p_cols_fp32) - # S occupies [0, qk_mma_tiler[1]) = [0, 128) - # O must NOT overlap P. Place O after max(S end, P end), aligned to 32. - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - p_end = self.tmem_p0_offset + p_cols_fp32 # 32 + 64 = 96 - s_cols = self.qk_mma_tiler[1] # 128 - o_after = max(s_cols, p_end) # 128 - self.tmem_o0_offset = ((o_after + 31) // 32) * 32 # align to 32 = 128 - o_cols = find_tmem_tensor_col_offset(tOtO) # footprint of O - total = self.tmem_o0_offset + o_cols - # Must be multiple of 32 AND power of 2 - self.num_tmem_alloc_cols = 1 - while self.num_tmem_alloc_cols < total: - self.num_tmem_alloc_cols *= 2 - cta = cute.size(qk_mma.thr_id.shape) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)) - self.q_tx_bytes = cute.size_in_bytes(self.q_dtype, q_s) * cta - self.kv_tx_bytes = cute.size_in_bytes(self.q_dtype, k_s) * cta - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - # FMHA-style V: reconstruct as (HEAD_DIM, s_k, 1) MN-major - v_fmha = cute.make_tensor( - v.iterator, - cute.make_layout( - (HEAD_DIM, 128, 1), - stride=(1, HEAD_DIM, HEAD_DIM * 128), - ), - ) - self.v_major = LayoutEnum.from_tensor(v_fmha).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, self.a_major, self.b_major, self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, self.qk_acc_dtype, self.cta_group, (128,HEAD_DIM), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)); v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id),q,q_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id),k,k_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v_fmha,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape) - epi_s = cute.select(self.c_smem_s,mode=[0,1]) - tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile) - self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.v_smem_s,self.p_tmem_s,self.c_smem_s,self.epi_tile).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, q_smem_s, k_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx,_,_ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k); cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - q_bar: cute.struct.MemRange[cutlass.Int64, self.q_stage*2] - kv_bar: cute.struct.MemRange[cutlass.Int64, self.kv_stage*2] - s_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage*2] - tmem_dealloc: cutlass.Int64; holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - qp,qc = pipeline.PipelineTmaUmma.create(barrier_storage=st.q_bar.data_ptr(),num_stages=self.q_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.q_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - kvp,kvc = pipeline.PipelineTmaUmma.create(barrier_storage=st.kv_bar.data_ptr(),num_stages=self.kv_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.kv_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - s_prod,s_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.s_bar.data_ptr(),num_stages=1,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.epilogue_warp_id))).make_participants() - softmax_done_bar = pipeline.NamedBarrier(barrier_id=3, num_threads=32 + 32*len(self.epilogue_warp_id)) - acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(),num_stages=self.num_acc_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,len(self.epilogue_warp_id)),cta_layout_vmnk=cl_vmnk,defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2,num_threads=32*len((self.mma_warp_id,*self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr,barrier_for_retrieve=tmem_bar,allocator_warp_id=self.epilogue_warp_id[0],is_two_cta=cute.size(qk_mma.thr_id.shape)==2,two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk,is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype,layout=q_smem_s.outer,byte_alignment=128,swizzle=q_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype,layout=k_smem_s.outer,byte_alignment=128,swizzle=k_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype,layout=v_smem_s.outer,byte_alignment=128,swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype,layout=c_smem_s.outer,byte_alignment=128,swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None)) - gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None)) - gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None)) - gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None)) - n_kv_tiles = cute.size(gK, mode=[3]) - - qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - tCgV = pv_thr.partition_B(gV); tCgC = pv_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,0,None,0)).shape) - tAsQ,tAgQ = cpasync.tma_partition(tma_q,0,a_lay,cute.group_modes(sQ,0,3),cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape) - tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3)) - tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)]; tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - # --- PV read view (for MMA only, NOT for softmax store) --- - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None,None,None,0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # TMA LOAD - if warp_idx == self.tma_warp_id: - qp.reset(); qh = qp.acquire_and_advance() - cute.copy(tma_q,tAgQ[(None,qh.count)],tAsQ[(None,qh.index)],tma_bar_ptr=qh.barrier) - qp.tail() - kvp.reset(); pk = kvp.try_acquire() - for kt in cutlass.range(n_kv_tiles,unroll=1): - kh = kvp.acquire_and_advance(pk) - cute.copy(tma_k,tBgK[(None,kh.count)],tBsK[(None,kh.index)],tma_bar_ptr=kh.barrier) - pk = cutlass.Boolean(1) - vh = kvp.acquire_and_advance(pk) - cute.copy(tma_v,tVgV[(None,vh.count)],tVsV[(None,vh.index)],tma_bar_ptr=vh.barrier) - pk = cutlass.Boolean(1) - kvp.tail() - - # MMA - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - qc.reset(); qh = qc.wait_and_advance(); qh.release() - kvc.reset(); pk = kvc.try_wait() - acc_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_stage) - acc_pipe.producer_acquire(acc_st) - for kt in range(n_kv_tiles): - kh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - sh = s_prod.acquire_and_advance() - qk_mma.set(tcgen05.Field.ACCUMULATE, False) - for kb in cutlass.range(cute.size(tCrQ,mode=[2]), unroll_full=True): - cute.gemm(qk_mma, tStS0, tCrQ[(None,None,kb,0)], tCrK[(None,None,kb,kh.index)], tStS0) - qk_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - sh.commit(); kh.release() - softmax_done_bar.arrive_and_wait() - vh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - pv_mma.set(tcgen05.Field.ACCUMULATE, kt != 0) - for kb in cutlass.range(cute.size(tOrP0,mode=[2]), unroll_full=True): - cute.gemm(pv_mma, tOtO0, tOrP0[(None,None,kb)], tCrV[(None,None,kb,vh.index)], tOtO0) - pv_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - vh.release() - acc_pipe.producer_commit(acc_st); acc_st.advance() - acc_pipe.producer_tail(acc_st) - - # EPILOGUE - if warp_idx < self.mma_warp_id: - tmem.allocate(self.num_tmem_alloc_cols) - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) - sfw_idx = tidx % (32 * len(self.epilogue_warp_id)) - - # --- S load (QK C-fragment layout) --- - tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS0) - thr_load = tiled_tmem_load.get_slice(sfw_idx) - tTMEM_LOADtS = thr_load.partition_S(tStS0) - - # S coordinate tensor (QK C-fragment) - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tTMEM_LOADcS = thr_load.partition_D(tScS) - - # --- P store (QK C-fragment layout composition, FMHA pattern) --- - # P logical columns = PV K = QK N = pv_mma_tiler[2] - # Packed FP32 columns: BF16 pairs packed into FP32 words - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - # BF16: 128 * 16 / 32 = 64 - - # P TMEM destination: QK C-fragment layout composed with P sub-tile - tStP_layout = cute.composition( - tStS.layout, - cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32)), - ) - tStP0 = cute.make_tensor( - tStS.iterator + self.tmem_p0_offset, - tStP_layout, - ) - - # P TMEM store atom and tiled copy - tmem_store_atom = cute.make_copy_atom( - tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), - self.qk_acc_dtype, - ) - tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStP0) - thr_store = tiled_tmem_store.get_slice(sfw_idx) - tTMEM_STOREtP = thr_store.partition_D(tStP0) - - # P coordinate tensor: QK C-fragment coordinate composed with P sub-tile - tScP_layout = cute.composition( - tScS.layout, - cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32)), - ) - tScP = cute.make_tensor(tScS.iterator, tScP_layout) - tTMEM_STOREcP = thr_store.partition_S(tScP) - - for kt in range(n_kv_tiles): - si_handle = s_cons.wait_and_advance() - - # Load S from TMEM (FP32, QK C-fragment layout) - tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS) - - # Register bridge (FMHA pattern): - # rP_words: FP32 backing store with store-partition shape - # rP_bf16: BF16 view over same registers using QK-load layout - rP_words = cute.make_rmem_tensor(tTMEM_STOREcP.shape, self.qk_acc_dtype) - rP_bf16 = cute.make_tensor( - cute.recast_ptr(rP_words.iterator, dtype=self.q_dtype), - tTMEM_LOADrS.layout, - ) - - # Fragmented load→convert→store: - # Load S as FP32, convert to BF16, store through rP_bf16 view - frg_cnt = 4 - frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt - tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile)) - rP_bf16_frg = cute.logical_divide(rP_bf16, cute.make_layout(frg_tile)) - for j in range(frg_cnt): - s_vec = tTMEM_LOADrS_frg[None, j].load() - rP_bf16_frg[None, j].store(s_vec.to(self.q_dtype)) - - # Copy packed FP32 backing registers to TMEM - cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP) - cute.arch.fence_view_async_tmem_store() - si_handle.release() - softmax_done_bar.arrive() - - tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout) - acc_cons_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage) - c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) - c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp) - acc_cons_st = utils.gemm.sm100.epilogue_tma_store(self, tidx, warp_idx, tma_c, tCtO_base, sC, tCgC, epi_tile, 0, const_expr(lambda x: x), (0,0,0), acc_cons_st, acc_pipe, c_pipe) - c_pipe.producer_tail() - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def test(): - torch.manual_seed(42) - for n in [128]: - m, hd = 128, HEAD_DIM - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(n, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(n, hd, dtype=torch.bfloat16, device='cuda') - # V passed as (n, hd) row-major — FMHA-style reconstruction inside kernel - v_kernel = v.unsqueeze(-1) - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - qf = q[:,:,0].float(); kf = k[:,:,0].float() - ref = (qf @ kf.T).bfloat16().float() @ v.float() - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - kernel = FmhaV3() - print(f'n={n}: Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - print(f'n={n}: tmem_offsets: s0={kernel.tmem_s0_offset} p0={kernel.tmem_p0_offset} o0={kernel.tmem_o0_offset} alloc={kernel.num_tmem_alloc_cols}', flush=True) - print(f'n={n}: Running...', flush=True) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - out = c[:,:,0].float() - cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - print(f'FMHA v3 n={n} V=ones: cosine {cos:.6f} {"PASS" if cos >= 0.99 else "FAIL"}') - if cos < 0.99: - print(f' out[0,:4]={out[0,:4].tolist()} ref[0,:4]={ref[0,:4].tolist()}') - -if __name__ == '__main__': - test() diff --git a/tests/archive/unit_test_fmha_v3_stage_c2.py b/tests/archive/unit_test_fmha_v3_stage_c2.py deleted file mode 100644 index 168170e4..00000000 --- a/tests/archive/unit_test_fmha_v3_stage_c2.py +++ /dev/null @@ -1,466 +0,0 @@ -""" -FMHA v3 Stage-C: Real softmax + O normalization. -Builds on the 12w identity-softmax test by replacing identity softmax with -online softmax (row_max, exp2 scaling, P store) and adding O normalization -by row_sum before the epilogue writes to GMEM. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct -import math - -HEAD_DIM = 64 - -class FmhaV3StageC2: - def __init__(self, s_k=128, scale_softmax=None): - self.s_k = s_k - self.acc_dtype = Float32; self.qk_acc_dtype = Float32; self.pv_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - # 12-warp layout - self.softmax_warp_ids = (0, 1, 2, 3) - self.correction_warp_ids = (4, 5, 6, 7) - self.mma_warp_id = 8; self.tma_warp_id = 9 - self.epilogue_warp_id = (10,); self.empty_warp_id = 11 - self.threads_per_cta = 32 * 12 - # Pipeline stages - self.mma_softmax_stage = 1; self.softmax_corr_stage = 1 - self.mma_corr_stage = 2; self.epi_stage = 2 - # TMA stages - self.kv_stage = 2; self.q_stage = 1; self.num_c_stage = 2 - # Softmax - self.scale_softmax = scale_softmax if scale_softmax is not None else 1.0 / math.sqrt(HEAD_DIM) - self.scale_softmax_log2 = self.scale_softmax * math.log2(math.e) - - def _setup(self, qk_mma, pv_mma): - qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (128, 128, qk_ik * 4) - pv_ik = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (128, HEAD_DIM, pv_ik * (128 // pv_ik)) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0]//cute.size(qk_mma.thr_id.shape), HEAD_DIM, self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - self.q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, self.q_dtype, self.q_stage) - self.k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, self.q_dtype, self.kv_stage) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, self.kv_stage) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - qk_thr = qk_mma.get_slice(0); qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - pv_thr = pv_mma.get_slice(0); pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - self.tmem_s0_offset = 0; self.tmem_vec0_offset = 0; self.tmem_p0_offset = 32 - # P occupies [tmem_p0_offset, tmem_p0_offset + p_cols_fp32) - # S occupies [0, qk_mma_tiler[1]) = [0, 128) - # O must NOT overlap P. Place O after max(S end, P end), aligned to 32. - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - p_end = self.tmem_p0_offset + p_cols_fp32 # 32 + 64 = 96 - s_cols = self.qk_mma_tiler[1] # 128 - o_after = max(s_cols, p_end) # 128 - self.tmem_o0_offset = ((o_after + 31) // 32) * 32 # align to 32 = 128 - o_cols = find_tmem_tensor_col_offset(tOtO) # footprint of O - total = self.tmem_o0_offset + o_cols - # Must be multiple of 32 AND power of 2 - self.num_tmem_alloc_cols = 1 - while self.num_tmem_alloc_cols < total: - self.num_tmem_alloc_cols *= 2 - cta = cute.size(qk_mma.thr_id.shape) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)) - self.q_tx_bytes = cute.size_in_bytes(self.q_dtype, q_s) * cta - self.kv_tx_bytes = cute.size_in_bytes(self.q_dtype, k_s) * cta - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - # FMHA-style V: reconstruct as (HEAD_DIM, s_k, 1) MN-major - v_fmha = cute.make_tensor( - v.iterator, - cute.make_layout( - (HEAD_DIM, self.s_k, 1), - stride=(1, HEAD_DIM, HEAD_DIM * self.s_k), - ), - ) - self.v_major = LayoutEnum.from_tensor(v_fmha).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, self.a_major, self.b_major, self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, self.qk_acc_dtype, self.cta_group, (128,HEAD_DIM), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)); v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id),q,q_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id),k,k_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v_fmha,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape) - epi_s = cute.select(self.c_smem_s,mode=[0,1]) - tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile) - self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.v_smem_s,self.p_tmem_s,self.c_smem_s,self.epi_tile).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, q_smem_s, k_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx, _, _ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k); cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - q_bar: cute.struct.MemRange[cutlass.Int64, self.q_stage * 2] - kv_bar: cute.struct.MemRange[cutlass.Int64, self.kv_stage * 2] - mma_s_bar: cute.struct.MemRange[cutlass.Int64, self.mma_softmax_stage * 2] - s_corr_bar: cute.struct.MemRange[cutlass.Int64, self.softmax_corr_stage * 2] - mma_corr_bar: cute.struct.MemRange[cutlass.Int64, self.mma_corr_stage * 2] - corr_epi_bar: cute.struct.MemRange[cutlass.Int64, self.epi_stage * 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, 2] - tmem_dealloc: cutlass.Int64; holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - def cg(n): return pipeline.CooperativeGroup(pipeline.Agent.Thread, n) - qp, qc = pipeline.PipelineTmaUmma.create(barrier_storage=st.q_bar.data_ptr(), num_stages=self.q_stage, producer_group=cg(1), consumer_group=cg(1), tx_count=self.q_tx_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True).make_participants() - kvp, kvc = pipeline.PipelineTmaUmma.create(barrier_storage=st.kv_bar.data_ptr(), num_stages=self.kv_stage, producer_group=cg(1), consumer_group=cg(1), tx_count=self.kv_tx_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True).make_participants() - # MMA → Softmax: S ready - mma_s_prod, mma_s_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.mma_s_bar.data_ptr(), num_stages=self.mma_softmax_stage, producer_group=cg(1), consumer_group=cg(32 * len(self.softmax_warp_ids)), cta_layout_vmnk=cl_vmnk, defer_sync=True).make_participants() - # Softmax → Correction: vec ready - s_corr_prod, s_corr_cons = pipeline.PipelineAsync.create(barrier_storage=st.s_corr_bar.data_ptr(), num_stages=self.softmax_corr_stage, producer_group=cg(32 * len(self.softmax_warp_ids)), consumer_group=cg(32 * len(self.correction_warp_ids))).make_participants() - # MMA → Correction: O ready - mma_corr_prod, mma_corr_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.mma_corr_bar.data_ptr(), num_stages=self.mma_corr_stage, producer_group=cg(1), consumer_group=cg(32 * len(self.correction_warp_ids)), cta_layout_vmnk=cl_vmnk, defer_sync=True).make_participants() - # Correction → Epilogue: O in SMEM ready - corr_epi_prod, corr_epi_cons = pipeline.PipelineAsync.create(barrier_storage=st.corr_epi_bar.data_ptr(), num_stages=self.epi_stage, producer_group=cg(32 * len(self.correction_warp_ids)), consumer_group=cg(32)).make_participants() - # Accumulator pipeline for epilogue (full pipeline, not participants) - acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(), num_stages=1, producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), cta_layout_vmnk=cl_vmnk, defer_sync=True) - # TMEM alloc barrier: softmax + correction + MMA + epilogue - tmem_bar = pipeline.NamedBarrier(barrier_id=2, num_threads=32 * len((*self.softmax_warp_ids, *self.correction_warp_ids, self.mma_warp_id, self.epilogue_warp_id))) - # Softmax done barrier: MMA waits for softmax to produce P before starting PV - softmax_done_bar = pipeline.NamedBarrier(barrier_id=3, num_threads=32 * len(self.softmax_warp_ids) + 32) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, allocator_warp_id=self.softmax_warp_ids[0], is_two_cta=cute.size(qk_mma.thr_id.shape) == 2, two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - if warp_idx == self.empty_warp_id: - cute.arch.mbarrier_init(st.tmem_dealloc, 32 * len((*self.softmax_warp_ids, *self.correction_warp_ids))) - cute.arch.mbarrier_init_fence() - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=q_smem_s.outer, byte_alignment=128, swizzle=q_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=k_smem_s.outer, byte_alignment=128, swizzle=k_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None, 0, None)), (None, None, None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0, None, None)), (None, None, None)) - gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0, None, None)), (None, None, None)) - gC = cute.local_tile(mC, cute.slice_(self.pv_mma_tiler, (None, None, 0)), (None, None, None)) - n_kv_tiles = cute.size(gK, mode=[3]) - - qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - tCgV = pv_thr.partition_B(gV); tCgC = pv_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0, 0, None, 0)).shape) - tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ, 0, 3), cute.group_modes(tCgQ, 0, 3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0, None, 0, 0)).shape) - tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK, 0, 3), cute.group_modes(tCgK, 0, 3)) - tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV, 0, 3), cute.group_modes(tCgV, 0, 3)) - tAgQ = tAgQ[(None, 0, None, 0)]; tBgK = tBgK[(None, 0, None, 0)]; tVgV = tVgV[(None, 0, None, 0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK); tCrV = pv_mma.make_fragment_B(sV) - qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None, None, None, 0)] - tOrP0 = cute.make_tensor(tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, tOrP.layout) - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_as, 1)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_as, 1)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ==================== TMA WARP (9) ==================== - if warp_idx == self.tma_warp_id: - qp.reset(); qh = qp.acquire_and_advance() - cute.copy(tma_q, tAgQ[(None, qh.count)], tAsQ[(None, qh.index)], tma_bar_ptr=qh.barrier) - qp.tail() - kvp.reset(); pk = kvp.try_acquire() - for kt in cutlass.range(n_kv_tiles, unroll=1): - kh = kvp.acquire_and_advance(pk) - cute.copy(tma_k, tBgK[(None, kh.count)], tBsK[(None, kh.index)], tma_bar_ptr=kh.barrier) - pk = cutlass.Boolean(1) - vh = kvp.acquire_and_advance(pk) - cute.copy(tma_v, tVgV[(None, vh.count)], tVsV[(None, vh.index)], tma_bar_ptr=vh.barrier) - pk = cutlass.Boolean(1) - kvp.tail() - - # ==================== MMA WARP (8) ==================== - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - qc.reset(); qh = qc.wait_and_advance(); qh.release() - kvc.reset(); pk = kvc.try_wait() - for kt in range(n_kv_tiles): - # QK -> S - kh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - sh = mma_s_prod.acquire_and_advance() - qk_mma.set(tcgen05.Field.ACCUMULATE, False) - for kb in cutlass.range(cute.size(tCrQ, mode=[2]), unroll_full=True): - cute.gemm(qk_mma, tStS0, tCrQ[(None, None, kb, 0)], tCrK[(None, None, kb, kh.index)], tStS0) - qk_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store(); sh.commit(); kh.release() - # PV -> O (wait for softmax to produce P) - softmax_done_bar.arrive_and_wait() - vh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - oh = mma_corr_prod.acquire_and_advance() - pv_mma.set(tcgen05.Field.ACCUMULATE, kt != 0) - for kb in cutlass.range(cute.size(tOrP0, mode=[2]), unroll_full=True): - cute.gemm(pv_mma, tOtO0, tOrP0[(None, None, kb)], tCrV[(None, None, kb, vh.index)], tOtO0) - pv_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store(); oh.commit(); vh.release() - mma_s_prod.tail(); mma_corr_prod.tail() - cute.arch.relinquish_tmem_alloc_permit() - cute.arch.mbarrier_wait(st.tmem_dealloc, 0) - tmem_ptr = cute.arch.retrieve_tmem_ptr(self.qk_acc_dtype, alignment=16, ptr_to_buffer_holding_addr=st.holding) - cute.arch.dealloc_tmem(tmem_ptr, Int32(self.num_tmem_alloc_cols)) - - # ==================== SOFTMAX WARPS (0-3) ==================== - if warp_idx < len(self.softmax_warp_ids): - tmem.allocate(self.num_tmem_alloc_cols); tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) - sfw_idx = tidx % (32 * len(self.softmax_warp_ids)) - - # S load setup - tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS0) - thr_load = tiled_tmem_load.get_slice(sfw_idx) - tTMEM_LOADtS = thr_load.partition_S(tStS0) - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tTMEM_LOADcS = thr_load.partition_D(tScS) - - # P store setup (QK C-fragment composition) - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - tStP_layout = cute.composition(tStS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tStP0 = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStP_layout) - tmem_store_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStP0) - thr_store = tiled_tmem_store.get_slice(sfw_idx) - tTMEM_STOREtP = thr_store.partition_D(tStP0) - tScP_layout = cute.composition(tScS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tTMEM_STOREcP = thr_store.partition_S(cute.make_tensor(tScS.iterator, tScP_layout)) - - # Vec store setup ([old_max, new_max] per iteration, [row_sum, row_max] at end) - tStS_vec_layout = cute.composition(tStS.layout, cute.make_layout((128, 2))) - tStS_vec = cute.make_tensor(tStS.iterator + self.tmem_vec0_offset, tStS_vec_layout) - tmem_store_vec_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(2)), self.qk_acc_dtype) - tiled_tmem_store_vec = tcgen05.make_tmem_copy(tmem_store_vec_atom, tStS_vec) - thr_store_vec = tiled_tmem_store_vec.get_slice(sfw_idx) - tTMEM_STORE_VECtS = thr_store_vec.partition_D(tStS_vec) - tScS_vec_layout = cute.composition(tScS.layout, cute.make_layout((128, 2))) - tScS_vec = cute.make_tensor(tScS.iterator, tScS_vec_layout) - tTMEM_STORE_VECcS = thr_store_vec.partition_S(tScS_vec) - - row_max = -Float32.inf; row_sum = Float32(0.0) - vec_handle = s_corr_prod.acquire_and_advance() - scale_log2 = Float32(self.scale_softmax_log2) - - for kt in range(n_kv_tiles): - si_handle = mma_s_cons.wait_and_advance() - tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS) - cute.arch.fence_view_async_tmem_load() - - # Row max (element-wise fmax) - old_row_max = row_max - frg_cnt = 4 - frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt - tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile)) - for j in range(frg_cnt): - for k in range(cute.size(tTMEM_LOADrS_frg, mode=[0])): - row_max = cute.arch.fmax(row_max, tTMEM_LOADrS_frg[k, j] * scale_log2) - - row_max_safe = row_max - if row_max == -cutlass.Float32.inf: row_max_safe = Float32(0.0) - - # Vec = [old_max, new_max] for correction - tTMEM_STORE_VECrS = cute.make_rmem_tensor(tTMEM_STORE_VECcS.shape, self.qk_acc_dtype) - tTMEM_STORE_VECrS[0] = old_row_max; tTMEM_STORE_VECrS[1] = row_max_safe - cute.copy(tiled_tmem_store_vec, tTMEM_STORE_VECrS, tTMEM_STORE_VECtS) - cute.arch.fence_view_async_tmem_store() - vec_handle.commit() - - # Scale row_sum and compute P - acc_scale_ = scale_log2 * (old_row_max - row_max_safe) - acc_scale = cute.math.exp2(acc_scale_, fastmath=True) - if old_row_max == -cutlass.Float32.inf: acc_scale = Float32(0.0) - row_sum *= acc_scale - rP_words = cute.make_rmem_tensor(tTMEM_STOREcP.shape, self.qk_acc_dtype) - rP_bf16 = cute.make_tensor(cute.recast_ptr(rP_words.iterator, dtype=self.q_dtype), tTMEM_LOADrS.layout) - minus_row_max_scale = (Float32(0.0) - row_max_safe) * scale_log2 - rP_bf16_frg = cute.logical_divide(rP_bf16, cute.make_layout(frg_tile)) - for j in range(frg_cnt): - for k in range(cute.size(tTMEM_LOADrS_frg, mode=[0])): - tTMEM_LOADrS_frg[k, j] = tTMEM_LOADrS_frg[k, j] * scale_log2 + minus_row_max_scale - tTMEM_LOADrS_frg[k, j] = cute.math.exp2(tTMEM_LOADrS_frg[k, j], fastmath=True) - row_sum = row_sum + tTMEM_LOADrS_frg[k, j] - s_vec = tTMEM_LOADrS_frg[None, j].load() - rP_bf16_frg[None, j].store(s_vec.to(self.q_dtype)) - - cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP) - cute.arch.fence_view_async_tmem_store() - si_handle.release() - softmax_done_bar.arrive() - vec_handle = s_corr_prod.acquire_and_advance() - - # Final vec = [row_sum, row_max] - tTMEM_STORE_VECrS = cute.make_rmem_tensor(tTMEM_STORE_VECcS.shape, self.qk_acc_dtype) - tTMEM_STORE_VECrS[0] = row_sum; tTMEM_STORE_VECrS[1] = row_max - cute.copy(tiled_tmem_store_vec, tTMEM_STORE_VECrS, tTMEM_STORE_VECtS) - cute.arch.fence_view_async_tmem_store() - vec_handle.commit() - s_corr_prod.acquire() # balance final pipe step - s_corr_prod.tail() - cute.arch.mbarrier_arrive(st.tmem_dealloc) - tmem.relinquish_alloc_permit() - - # ==================== CORRECTION WARPS (4-7) ==================== - if warp_idx >= len(self.softmax_warp_ids) and warp_idx < len(self.softmax_warp_ids) + len(self.correction_warp_ids): - tmem.wait_for_alloc() - corr_idx = tidx % (32 * len(self.correction_warp_ids)) - # Vec load setup (compute cS from scratch for correction warps) - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tStS_vec_layout = cute.composition(tStS.layout, cute.make_layout((128, 2))) - tStS_vec = cute.make_tensor(tStS.iterator + self.tmem_vec0_offset, tStS_vec_layout) - tmem_load_vec_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(2)), self.qk_acc_dtype) - tiled_tmem_load_vec = tcgen05.make_tmem_copy(tmem_load_vec_atom, tStS_vec) - thr_load_vec = tiled_tmem_load_vec.get_slice(corr_idx) - tTMEM_LOAD_VECtS = thr_load_vec.partition_S(tStS_vec) - tScS_vec = cute.make_tensor(tScS.iterator, cute.composition(tScS.layout, cute.make_layout((128, 2)))) - tTMEM_LOAD_VECcS = thr_load_vec.partition_D(tScS_vec) - # O rescale setup (matching CUTLASS correction_rescale) - corr_tile_size = 16 - cO = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1])) - tOcO = pv_thr.partition_C(cO) - tOtO_i_layout = cute.composition(tOtO.layout, cute.make_layout((128, corr_tile_size))) - tOcO_i_layout = cute.composition(tOcO.layout, cute.make_layout((128, corr_tile_size))) - tOtO_i = cute.make_tensor(tOtO.iterator, tOtO_i_layout) - tOcO_i = cute.make_tensor(tOcO.iterator, tOcO_i_layout) - tmem_load_o_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), self.pv_acc_dtype) - tmem_store_o_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), self.pv_acc_dtype) - tiled_tmem_load_o = tcgen05.make_tmem_copy(tmem_load_o_atom, tOtO_i) - tiled_tmem_store_o = tcgen05.make_tmem_copy(tmem_store_o_atom, tOtO_i) - thr_load_o = tiled_tmem_load_o.get_slice(corr_idx) - thr_store_o = tiled_tmem_store_o.get_slice(corr_idx) - tTMEM_LOAD_OtO = thr_load_o.partition_S(tOtO_i) - tTMEM_LOAD_OcO = thr_load_o.partition_D(tOcO_i) - tTMEM_STORE_OtO = thr_store_o.partition_D(tOtO_i) - scale_log2 = Float32(self.scale_softmax_log2) - - # Correction rescale loop: for each KV tile (except first), rescale O - first_vec = s_corr_cons.wait_and_advance(); first_vec.release() - for kt in range(n_kv_tiles - 1): - vec = s_corr_cons.wait_and_advance() - # Read vec = [old_max, new_max] - tTMEM_LOAD_VECrS = cute.make_rmem_tensor(tTMEM_LOAD_VECcS.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load_vec, tTMEM_LOAD_VECtS, tTMEM_LOAD_VECrS) - cute.arch.fence_view_async_tmem_load() - old_max = tTMEM_LOAD_VECrS[0]; new_max = tTMEM_LOAD_VECrS[1] - corr_scale = cute.math.exp2(scale_log2 * (old_max - new_max), fastmath=True) - # Wait for O from MMA, rescale O in TMEM - o_handle = mma_corr_cons.wait_and_advance() - o_col_tiles = self.pv_mma_tiler[1] // corr_tile_size - for i in range(o_col_tiles): - tTMEM_LOAD_O_i = cute.make_tensor(tTMEM_LOAD_OtO.iterator + i * corr_tile_size, tTMEM_LOAD_OtO.layout) - tTMEM_STORE_O_i = cute.make_tensor(tTMEM_STORE_OtO.iterator + i * corr_tile_size, tTMEM_STORE_OtO.layout) - tTMrO = cute.make_rmem_tensor(tTMEM_LOAD_OcO.shape, self.pv_acc_dtype) - cute.copy(tiled_tmem_load_o, tTMEM_LOAD_O_i, tTMrO) - for k in cutlass.range(cute.size(tTMrO), vectorize=True): - tTMrO[k] = tTMrO[k] * corr_scale - cute.copy(tiled_tmem_store_o, tTMrO, tTMEM_STORE_O_i) - cute.arch.fence_view_async_tmem_store() - o_handle.release(); vec.release() - - # Final: read [row_sum, row_max], normalize O, write to SMEM via epilogue_tma_store - final_vec = s_corr_cons.wait_and_advance() - tTMEM_LOAD_VECrS = cute.make_rmem_tensor(tTMEM_LOAD_VECcS.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load_vec, tTMEM_LOAD_VECtS, tTMEM_LOAD_VECrS) - cute.arch.fence_view_async_tmem_load() - row_sum = tTMEM_LOAD_VECrS[0]; row_max = tTMEM_LOAD_VECrS[1] - final_vec.release() - - final_o = mma_corr_cons.wait_and_advance() - epi_handle = corr_epi_prod.acquire_and_advance() - - # Normalize O in TMEM by 1/row_sum - inv_row_sum = Float32(1.0) / row_sum - for i in range(self.pv_mma_tiler[1] // corr_tile_size): - tTMEM_LOAD_O_i = cute.make_tensor(tTMEM_LOAD_OtO.iterator + i * corr_tile_size, tTMEM_LOAD_OtO.layout) - tTMEM_STORE_O_i = cute.make_tensor(tTMEM_STORE_OtO.iterator + i * corr_tile_size, tTMEM_STORE_OtO.layout) - tTMrO = cute.make_rmem_tensor(tTMEM_LOAD_OcO.shape, self.pv_acc_dtype) - cute.copy(tiled_tmem_load_o, tTMEM_LOAD_O_i, tTMrO) - for k in cutlass.range(cute.size(tTMrO), vectorize=True): - tTMrO[k] = tTMrO[k] * inv_row_sum - cute.copy(tiled_tmem_store_o, tTMrO, tTMEM_STORE_O_i) - cute.arch.fence_view_async_tmem_store() - final_o.release() - epi_handle.commit() - cute.arch.mbarrier_arrive(st.tmem_dealloc) - - # ==================== EPILOGUE WARP (10) ==================== - if warp_idx == self.epilogue_warp_id[0]: - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) - epi_handle = corr_epi_cons.wait_and_advance() - # Signal acc_pipe that O is ready (correction already normalized in TMEM) - acc_prod_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, 1) - acc_pipe.producer_acquire(acc_prod_st) - acc_pipe.producer_commit(acc_prod_st); acc_prod_st.advance() - acc_pipe.producer_tail(acc_prod_st) - # Write O from TMEM to GMEM via epilogue_tma_store - tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout) - acc_cons_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, 1) - c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32) - c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp) - acc_cons_st = utils.gemm.sm100.epilogue_tma_store(self, tidx, warp_idx, tma_c, tCtO_base, sC, tCgC, epi_tile, 0, const_expr(lambda x: x), (0,0,0), acc_cons_st, acc_pipe, c_pipe) - c_pipe.producer_tail() - epi_handle.release() -def test(): - torch.manual_seed(42) - for n in [128]: - for seed in [42, 123, 999]: - torch.manual_seed(seed) - m, hd = 128, HEAD_DIM - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(n, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(n, hd, dtype=torch.bfloat16, device='cuda') - v_kernel = v.unsqueeze(-1) - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - qf = q[:,:,0].float(); kf = k[:,:,0].float() - scale = 1.0 / math.sqrt(hd) - attn = qf @ kf.T * scale - attn = torch.softmax(attn, dim=-1) - ref = attn @ v.float() - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - kernel = FmhaV3StageC2() - if seed == 42: - print(f'seed={seed}: Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - if seed == 42: - print(f'tmem_offsets: s0={kernel.tmem_s0_offset} p0={kernel.tmem_p0_offset} o0={kernel.tmem_o0_offset} alloc={kernel.num_tmem_alloc_cols}', flush=True) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - out = c[:,:,0].float() - cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - print(f'FMHA Stage-C n={n} seed={seed}: cosine {cos:.6f} {"PASS" if cos >= 0.99 else "FAIL"}') - if cos < 0.99: - print(f' out[0,:4]={out[0,:4].tolist()} ref[0,:4]={ref[0,:4].tolist()}') - -if __name__ == '__main__': - test() diff --git a/tests/archive/unit_test_fmha_v3_stage_c_full.py b/tests/archive/unit_test_fmha_v3_stage_c_full.py deleted file mode 100644 index eec8aa2f..00000000 --- a/tests/archive/unit_test_fmha_v3_stage_c_full.py +++ /dev/null @@ -1,454 +0,0 @@ -""" -FMHA v3 Stage-C: Real online softmax with full multi-tile support. - -Integrated from: -- test_fmha_v3.py (Stage A+B: QK→identity softmax→PV) -- test_fmha_v3_stage_c_full.py (Stage C: single-tile real softmax) -- fmha_v3_stage_c_example3.py (Multi-tile: combined K+V barrier, O rescale, final_o_bar) - -Architecture (6-warp, 192 threads): - Warps 0-3: Softmax + Epilogue (row_max, exp2, P store, O rescale, O normalize, TMA store) - Warp 4: MMA (QK GEMM, PV GEMM) - Warp 5: TMA (Q, K, V load) - -Multi-tile key changes vs single-tile: - 1. Combined K+V barrier: one acquire_and_advance per kt, both K and V share kvh.barrier. - kvh.count == kt naturally — no interleaving, no Python int in TMA coordinates. - 2. kv_tx_bytes covers BOTH K and V transfers per stage. - 3. V FMHA layout uses s_k as the sequence dimension (compile-time, not hardcoded 128). - 4. MMA: same slot index (kvh.index) for both K (QK) and V (PV). Release after PV. - 5. O rescale for kt > 0: exp2((old_max - new_max) * scale_log2) on O in TMEM. - 6. final_o_bar: MMA arrives between producer_commit and producer_tail; - softmax arrives_and_wait before reading O for final normalize. - 7. s_k is a constructor parameter — each sequence length requires its own compiled kernel. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct -import math - -HEAD_DIM = 64 - - -class FmhaV3StageC: - def __init__(self, s_k=128, scale_softmax=None): - # s_k MUST equal actual sequence length n (compile-time constant for V layout). - self.s_k = s_k - self.n_kv_tiles = s_k // 128 # Python int — needed for range() unrolling - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0,1,2,3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.num_c_stage = 2 - self.kv_stage = 2; self.q_stage = 1; self.num_c_stage = 2 - self.scale_softmax = scale_softmax if scale_softmax is not None else 1.0 / math.sqrt(HEAD_DIM) - self.scale_softmax_log2 = self.scale_softmax * math.log2(math.e) - - def _setup(self, qk_mma, pv_mma): - qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (128, 128, qk_ik * 4) - pv_ik = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (128, HEAD_DIM, pv_ik * (128 // pv_ik)) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0]//cute.size(qk_mma.thr_id.shape), HEAD_DIM, self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - self.q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, self.q_dtype, self.q_stage) - self.k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, self.q_dtype, self.kv_stage) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, self.kv_stage) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - qk_thr = qk_mma.get_slice(0); qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - pv_thr = pv_mma.get_slice(0); pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - self.tmem_s0_offset = 0; self.tmem_p0_offset = 32 - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - p_end = self.tmem_p0_offset + p_cols_fp32 - s_cols = self.qk_mma_tiler[1] - o_after = max(s_cols, p_end) - self.tmem_o0_offset = ((o_after + 31) // 32) * 32 - o_cols = find_tmem_tensor_col_offset(tOtO) - total = self.tmem_o0_offset + o_cols - self.num_tmem_alloc_cols = 1 - while self.num_tmem_alloc_cols < total: - self.num_tmem_alloc_cols *= 2 - cta = cute.size(qk_mma.thr_id.shape) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)) - k_s = cute.slice_(self.k_smem_s,(None,None,None,0)) - v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - self.q_tx_bytes = cute.size_in_bytes(self.q_dtype, q_s) * cta - # Combined barrier: tx_count covers BOTH K and V transfers per acquire. - self.kv_tx_bytes = (cute.size_in_bytes(self.q_dtype, k_s) + - cute.size_in_bytes(self.q_dtype, v_s)) * cta - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - # FMHA-style V: reconstruct as (HEAD_DIM, s_k, 1) MN-major - # s_k is compile-time — must match actual sequence length n. - v_fmha = cute.make_tensor( - v.iterator, - cute.make_layout( - (HEAD_DIM, self.s_k, 1), - stride=(1, HEAD_DIM, HEAD_DIM * self.s_k), - ), - ) - self.v_major = LayoutEnum.from_tensor(v_fmha).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, self.a_major, self.b_major, self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, self.qk_acc_dtype, self.cta_group, (128,HEAD_DIM), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)); v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id),q,q_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id),k,k_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v_fmha,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape) - epi_s = cute.select(self.c_smem_s,mode=[0,1]) - tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile) - self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.v_smem_s,self.p_tmem_s,self.c_smem_s,self.epi_tile).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, q_smem_s, k_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx,_,_ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k); cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - q_bar: cute.struct.MemRange[cutlass.Int64, self.q_stage*2] - kv_bar: cute.struct.MemRange[cutlass.Int64, self.kv_stage*2] - s_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage*2] - tmem_dealloc: cutlass.Int64; holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - qp,qc = pipeline.PipelineTmaUmma.create(barrier_storage=st.q_bar.data_ptr(),num_stages=self.q_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.q_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - # Combined K+V pipeline: each stage carries BOTH K and V loaded together. - # One acquire per kt → kvh.count == kt (pipeline state value, accepted by cute.copy). - # No interleaving problem. No Python int in TMA coordinates. - kvp,kvc = pipeline.PipelineTmaUmma.create(barrier_storage=st.kv_bar.data_ptr(),num_stages=self.kv_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.kv_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - s_prod,s_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.s_bar.data_ptr(),num_stages=1,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.epilogue_warp_id))).make_participants() - softmax_done_bar = pipeline.NamedBarrier(barrier_id=3, num_threads=32 + 32*len(self.epilogue_warp_id)) - # Final-O sync: MMA arrives between producer_commit and producer_tail; - # softmax arrives_and_waits before reading O for the final normalize. - # This prevents softmax from racing MMA's PV[N-1] and dividing a - # partially-accumulated O by row_sum. - final_o_bar = pipeline.NamedBarrier(barrier_id=4, num_threads=32 + 32*len(self.epilogue_warp_id)) - acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(),num_stages=self.num_acc_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,len(self.epilogue_warp_id)),cta_layout_vmnk=cl_vmnk,defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2,num_threads=32*len((self.mma_warp_id,*self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr,barrier_for_retrieve=tmem_bar,allocator_warp_id=self.epilogue_warp_id[0],is_two_cta=cute.size(qk_mma.thr_id.shape)==2,two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk,is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype,layout=q_smem_s.outer,byte_alignment=128,swizzle=q_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype,layout=k_smem_s.outer,byte_alignment=128,swizzle=k_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype,layout=v_smem_s.outer,byte_alignment=128,swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype,layout=c_smem_s.outer,byte_alignment=128,swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None)) - gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None)) - gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None)) - gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None)) - n_kv_tiles = cute.size(gK, mode=[3]) - - qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - tCgV = pv_thr.partition_B(gV); tCgC = pv_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,0,None,0)).shape) - tAsQ,tAgQ = cpasync.tma_partition(tma_q,0,a_lay,cute.group_modes(sQ,0,3),cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape) - tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3)) - tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3)) - # GMEM slices: K uses mode 1 for GMEM iter → (None,None,0,0) keeps it free - # V uses mode 2 for GMEM iter → (None,0,None,0) keeps it free - # Q has 1 tile → (None,0,None,0) hardcode is fine - # CRITICAL: K from QK MMA B-partition has GMEM iter at mode 1, NOT mode 2! - # (None,0,None,0) for K hardcodes mode 1 to 0 → always loads tile 0. - # (None,None,0,0) for K keeps mode 1 free → correct multi-tile loading. - # Proven by diag test: (None,0,None,0) gives cos 0.711, (None,None,0,0) gives 0.999999. - tAgQ = tAgQ[(None,0,None,0)] # Q: 1 tile - tBgK = tBgK[(None,None,0,0)] # K: keep mode 1 (GMEM iter) free - tVgV = tVgV[(None,0,None,0)] # V: keep mode 2 (GMEM iter) free - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - # --- PV read view (for MMA only, NOT for softmax store) --- - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None,None,None,0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ===== TMA LOAD warp ===== - # GMEM tile coordinate: manual Int32 counter (kv_coord). SMEM slot: kvh.index. - # Pipeline handle .count is NOT a usable GMEM coordinate. - if warp_idx == self.tma_warp_id: - qp.reset(); qh = qp.acquire_and_advance() - cute.copy(tma_q, tAgQ[(None, Int32(0))], tAsQ[(None, qh.index)], tma_bar_ptr=qh.barrier) - qp.tail() - kvp.reset(); pk = kvp.try_acquire() - # Use cutlass.range with Python int n_kv_tiles for proper pipeline - # semantics (acquire/release). Wrap kt in Int32() for TMA coordinate. - for kt in cutlass.range(self.n_kv_tiles, unroll=1): - coord = Int32(kt) - kvh = kvp.acquire_and_advance(pk) - cute.copy(tma_k, tBgK[(None, coord)], tBsK[(None, kvh.index)], tma_bar_ptr=kvh.barrier) - cute.copy(tma_v, tVgV[(None, coord)], tVsV[(None, kvh.index)], tma_bar_ptr=kvh.barrier) - pk = cutlass.Boolean(1) - kvp.tail() - - # ===== MMA warp ===== - # One wait per kt; same slot index used for both K (QK) and V (PV). - # Release happens AFTER PV — combined slot stays held across QK+PV. - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - qc.reset(); qh = qc.wait_and_advance(); qh.release() - kvc.reset(); pk = kvc.try_wait() - acc_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_stage) - acc_pipe.producer_acquire(acc_st) - for kt in range(self.n_kv_tiles): - kvh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - sh = s_prod.acquire_and_advance() - qk_mma.set(tcgen05.Field.ACCUMULATE, False) - for kb in cutlass.range(cute.size(tCrQ, mode=[2]), unroll_full=True): - cute.gemm(qk_mma, tStS0, tCrQ[(None,None,kb,0)], tCrK[(None,None,kb,kvh.index)], tStS0) - qk_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - sh.commit() - softmax_done_bar.arrive_and_wait() - pv_mma.set(tcgen05.Field.ACCUMULATE, kt != 0) - for kb in cutlass.range(cute.size(tOrP0, mode=[2]), unroll_full=True): - cute.gemm(pv_mma, tOtO0, tOrP0[(None,None,kb)], tCrV[(None,None,kb,kvh.index)], tOtO0) - pv_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - kvh.release() - acc_pipe.producer_commit(acc_st); acc_st.advance() - # Signal softmax FIRST so it can run normalize + epilogue. Then - # wait for the epilogue's consumer-release in producer_tail. - # Reverse order deadlocks: producer_tail blocks waiting for - # consumer release; softmax blocks at final_o_bar waiting for - # MMA arrive; the epilogue (which does the release) is gated - # behind softmax's final_o_bar wait. Cycle. - final_o_bar.arrive() - acc_pipe.producer_tail(acc_st) - - # ===== SOFTMAX + EPILOGUE warps ===== - if warp_idx < self.mma_warp_id: - tmem.allocate(self.num_tmem_alloc_cols) - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) - sfw_idx = tidx % (32 * len(self.epilogue_warp_id)) - - # S load (QK C-fragment layout) - tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS0) - thr_load = tiled_tmem_load.get_slice(sfw_idx) - tTMEM_LOADtS = thr_load.partition_S(tStS0) - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tTMEM_LOADcS = thr_load.partition_D(tScS) - - # P store (QK C-fragment layout composition, FMHA pattern) - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - tStP_layout = cute.composition(tStS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tStP0 = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStP_layout) - tmem_store_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStP0) - thr_store = tiled_tmem_store.get_slice(sfw_idx) - tTMEM_STOREtP = thr_store.partition_D(tStP0) - tScP_layout = cute.composition(tScS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tScP = cute.make_tensor(tScS.iterator, tScP_layout) - tTMEM_STOREcP = thr_store.partition_S(tScP) - - # O rescale / normalize path - cO = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1])) - tOcO = pv_thr.partition_C(cO) - corr_tile_size = 16 - tOtO_i_layout = cute.composition(tOtO.layout, cute.make_layout((128, corr_tile_size))) - tOcO_i_layout = cute.composition(tOcO.layout, cute.make_layout((128, corr_tile_size))) - tOtO_i = cute.make_tensor(tOtO.iterator, tOtO_i_layout) - tOcO_i = cute.make_tensor(tOcO.iterator, tOcO_i_layout) - tmem_load_o_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), self.acc_dtype) - tmem_store_o_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), self.acc_dtype) - tiled_tmem_load_o = tcgen05.make_tmem_copy(tmem_load_o_atom, tOtO_i) - tiled_tmem_store_o = tcgen05.make_tmem_copy(tmem_store_o_atom, tOtO_i) - thr_load_o = tiled_tmem_load_o.get_slice(sfw_idx) - thr_store_o = tiled_tmem_store_o.get_slice(sfw_idx) - tTMEM_LOAD_OtO = thr_load_o.partition_S(tOtO_i) - tTMEM_LOAD_OcO = thr_load_o.partition_D(tOcO_i) - tTMEM_STORE_OtO = thr_store_o.partition_D(tOtO_i) - - o_col_tiles = self.pv_mma_tiler[1] // corr_tile_size - - row_max = -Float32.inf - row_sum = Float32(0.0) - scale_log2 = Float32(self.scale_softmax_log2) - - for kt in range(self.n_kv_tiles): - si_handle = s_cons.wait_and_advance() - - # Load S[kt] - tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS) - cute.arch.fence_view_async_tmem_load() - - # Pass 1: update row_max - old_row_max = row_max - frg_cnt = 4 - frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt - tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile)) - for j in range(frg_cnt): - for k in range(cute.size(tTMEM_LOADrS_frg, mode=[0])): - row_max = cute.arch.fmax(row_max, tTMEM_LOADrS_frg[k, j] * scale_log2) - - row_max_safe = row_max - if row_max == -cutlass.Float32.inf: - row_max_safe = Float32(0.0) - - # acc_scale for both row_sum rescale and O rescale. - # row_max is already in log2(scaled) space (S * scale_log2), - # so the difference old_row_max - row_max_safe is already the - # correct exponent for exp2. No extra scale_log2 factor. - acc_scale = cute.math.exp2(old_row_max - row_max_safe, fastmath=True) - if old_row_max == -cutlass.Float32.inf: - acc_scale = Float32(0.0) - row_sum *= acc_scale - - # Pass 2: P = exp2((S - new_max) * log2), accumulate row_sum, - # store BF16 P through the FP32-backed register bridge. - rP_words = cute.make_rmem_tensor(tTMEM_STOREcP.shape, self.qk_acc_dtype) - rP_bf16 = cute.make_tensor(cute.recast_ptr(rP_words.iterator, dtype=self.q_dtype), tTMEM_LOADrS.layout) - minus_row_max_scale = (Float32(0.0) - row_max_safe) * scale_log2 - - rP_bf16_frg = cute.logical_divide(rP_bf16, cute.make_layout(frg_tile)) - for j in range(frg_cnt): - for k in range(cute.size(tTMEM_LOADrS_frg, mode=[0])): - tTMEM_LOADrS_frg[k, j] = tTMEM_LOADrS_frg[k, j] * scale_log2 + minus_row_max_scale - tTMEM_LOADrS_frg[k, j] = cute.math.exp2(tTMEM_LOADrS_frg[k, j], fastmath=True) - row_sum = row_sum + tTMEM_LOADrS_frg[k, j] - s_vec = tTMEM_LOADrS_frg[None, j].load() - rP_bf16_frg[None, j].store(s_vec.to(self.q_dtype)) - - cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP) - cute.arch.fence_view_async_tmem_store() - - # O rescale TEMPORARILY DISABLED for debugging NaN - # if kt > 0: - # cute.arch.fence_view_async_tmem_load() - # for i in range(o_col_tiles): - # ... - # cute.arch.fence_view_async_tmem_store() - - si_handle.release() - softmax_done_bar.arrive() - - # Wait for MMA's last PV to commit before reading O for normalize. - # Without this barrier softmax can race MMA's PV[N-1]. - final_o_bar.arrive_and_wait() - - # Final O = O / row_sum - inv_row_sum = Float32(1.0) / row_sum - for i in range(o_col_tiles): - tTMEM_LOAD_O_i = cute.make_tensor( - tTMEM_LOAD_OtO.iterator + i * corr_tile_size, - tTMEM_LOAD_OtO.layout, - ) - tTMEM_STORE_O_i = cute.make_tensor( - tTMEM_STORE_OtO.iterator + i * corr_tile_size, - tTMEM_STORE_OtO.layout, - ) - tTMrO = cute.make_rmem_tensor(tTMEM_LOAD_OcO.shape, self.acc_dtype) - cute.copy(tiled_tmem_load_o, tTMEM_LOAD_O_i, tTMrO) - cute.arch.fence_view_async_tmem_load() - for k in cutlass.range(cute.size(tTMrO), vectorize=True): - tTMrO[k] = tTMrO[k] * inv_row_sum - cute.copy(tiled_tmem_store_o, tTMrO, tTMEM_STORE_O_i) - cute.arch.fence_view_async_tmem_store() - - # Epilogue: TMEM -> SMEM -> GMEM via TMA store - tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout) - acc_cons_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage) - c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) - c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp) - acc_cons_st = utils.gemm.sm100.epilogue_tma_store(self, tidx, warp_idx, tma_c, tCtO_base, sC, tCgC, epi_tile, 0, const_expr(lambda x: x), (0,0,0), acc_cons_st, acc_pipe, c_pipe) - c_pipe.producer_tail() - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def test(): - torch.manual_seed(42) - for n in [128, 256, 512, 1024]: - torch.manual_seed(42) - m, hd = 128, HEAD_DIM - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(n, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(n, hd, dtype=torch.bfloat16, device='cuda') - v_kernel = v.unsqueeze(-1) - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - - # PyTorch reference - qf = q[:, :, 0].float() - kf = k[:, :, 0].float() - scale = 1.0 / math.sqrt(hd) - attn = qf @ kf.T * scale - attn = torch.softmax(attn, dim=-1) - ref = attn @ v.float() - - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - # Each n requires its own compiled kernel (s_k is compile-time). - kernel = FmhaV3StageC(s_k=n) - print(f'n={n}: Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - print(f'n={n}: tmem s0={kernel.tmem_s0_offset} p0={kernel.tmem_p0_offset} ' - f'o0={kernel.tmem_o0_offset} alloc={kernel.num_tmem_alloc_cols} ' - f'kv_tx_bytes={kernel.kv_tx_bytes}', flush=True) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - - out = c[:, :, 0].float() - has_nan = out.isnan().any().item() - has_inf = out.isinf().any().item() - cos = torch.nn.functional.cosine_similarity( - out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0) - ).item() - max_abs = (out - ref).abs().max().item() - n_tiles = n // 128 - print(f'FMHA Stage-C n={n} ({n_tiles} kv tiles): ' - f'cos {cos:.6f} max_abs {max_abs:.4f} ' - f'nan={has_nan} inf={has_inf} ' - f'{"PASS" if cos >= 0.99 else "FAIL"}') - if cos < 0.99: - print(f' out[0,:4]={out[0,:4].tolist()}') - print(f' ref[0,:4]={ref[0,:4].tolist()}') - - -if __name__ == '__main__': - test() diff --git a/tests/archive/unit_test_fmha_v3_stage_c_min.py b/tests/archive/unit_test_fmha_v3_stage_c_min.py deleted file mode 100644 index a915bb57..00000000 --- a/tests/archive/unit_test_fmha_v3_stage_c_min.py +++ /dev/null @@ -1,305 +0,0 @@ -""" -FMHA v3 Stage-C minimal: 12-warps, identity softmax, identity correction. -Validates pipeline chain: mma_s, s_corr, mma_corr, corr_epi. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - -HEAD_DIM = 64 - -class FmhaV3StageCMin: - def __init__(self, s_k=128): - self.s_k = s_k - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.softmax_warp_ids = (0,1,2,3) - self.correction_warp_ids = (4,5,6,7) - self.mma_warp_id = 8; self.tma_warp_id = 9 - self.epilogue_warp_id = (10,) - self.epi_warp_id = 10; self.empty_warp_id = 11 - self.threads_per_cta = 32 * 12 - self.mma_softmax_stage = 1; self.softmax_corr_stage = 1 - self.mma_corr_stage = 2; self.epi_stage = 2 - self.kv_stage = 2; self.q_stage = 1; self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (128, 128, qk_ik * 4) - pv_ik = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (128, HEAD_DIM, pv_ik * (128 // pv_ik)) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0]//cute.size(qk_mma.thr_id.shape), HEAD_DIM, self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - self.q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, self.q_dtype, self.q_stage) - self.k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, self.q_dtype, self.kv_stage) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, self.kv_stage) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - qk_thr = qk_mma.get_slice(0); qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - pv_thr = pv_mma.get_slice(0); pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - self.tmem_s0_offset = 0; self.tmem_p0_offset = 32 - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - p_end = self.tmem_p0_offset + p_cols_fp32 - s_cols = self.qk_mma_tiler[1]; o_after = max(s_cols, p_end) - self.tmem_o0_offset = ((o_after + 31) // 32) * 32 - o_cols = find_tmem_tensor_col_offset(tOtO) - total = self.tmem_o0_offset + o_cols - self.num_tmem_alloc_cols = 1 - while self.num_tmem_alloc_cols < total: - self.num_tmem_alloc_cols *= 2 - cta = cute.size(qk_mma.thr_id.shape) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)) - k_s = cute.slice_(self.k_smem_s,(None,None,None,0)) - self.q_tx_bytes = cute.size_in_bytes(self.q_dtype, q_s) * cta - self.kv_tx_bytes = cute.size_in_bytes(self.q_dtype, k_s) * cta - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - v_fmha = cute.make_tensor(v.iterator, cute.make_layout((HEAD_DIM, self.s_k, 1), stride=(1, HEAD_DIM, HEAD_DIM * self.s_k))) - self.v_major = LayoutEnum.from_tensor(v_fmha).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, self.a_major, self.b_major, self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, self.qk_acc_dtype, self.cta_group, (128,HEAD_DIM), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)); v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id),q,q_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id),k,k_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v_fmha,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape) - epi_s = cute.select(self.c_smem_s,mode=[0,1]) - tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile) - self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.v_smem_s,self.p_tmem_s,self.c_smem_s,self.epi_tile).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, q_smem_s, k_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx,_,_ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k); cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - @cute.struct - class SS: - q_bar: cute.struct.MemRange[cutlass.Int64, self.q_stage*2] - kv_bar: cute.struct.MemRange[cutlass.Int64, self.kv_stage*2] - mma_s_bar: cute.struct.MemRange[cutlass.Int64, self.mma_softmax_stage*2] - s_corr_bar: cute.struct.MemRange[cutlass.Int64, self.softmax_corr_stage*2] - mma_corr_bar: cute.struct.MemRange[cutlass.Int64, self.mma_corr_stage*2] - corr_epi_bar: cute.struct.MemRange[cutlass.Int64, self.epi_stage*2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - def cg(n): return pipeline.CooperativeGroup(pipeline.Agent.Thread, n) - qp,qc = pipeline.PipelineTmaUmma.create(barrier_storage=st.q_bar.data_ptr(),num_stages=self.q_stage,producer_group=cg(1),consumer_group=cg(1),tx_count=self.q_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - kvp,kvc = pipeline.PipelineTmaUmma.create(barrier_storage=st.kv_bar.data_ptr(),num_stages=self.kv_stage,producer_group=cg(1),consumer_group=cg(1),tx_count=self.kv_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - mma_s_prod,mma_s_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.mma_s_bar.data_ptr(),num_stages=self.mma_softmax_stage,producer_group=cg(1),consumer_group=cg(32*len(self.softmax_warp_ids)),cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - s_corr_prod,s_corr_cons = pipeline.PipelineAsync.create(barrier_storage=st.s_corr_bar.data_ptr(),num_stages=self.softmax_corr_stage,producer_group=cg(32*len(self.softmax_warp_ids)),consumer_group=cg(32*len(self.correction_warp_ids))).make_participants() - mma_corr_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.mma_corr_bar.data_ptr(),num_stages=self.mma_corr_stage,producer_group=cg(1),consumer_group=cg(32*len(self.correction_warp_ids)),cta_layout_vmnk=cl_vmnk,defer_sync=True) - corr_epi_prod,corr_epi_cons = pipeline.PipelineAsync.create(barrier_storage=st.corr_epi_bar.data_ptr(),num_stages=self.epi_stage,producer_group=cg(32*len(self.correction_warp_ids)),consumer_group=cg(32)).make_participants() - tmem_bar = pipeline.NamedBarrier(barrier_id=2,num_threads=32*len((*self.softmax_warp_ids,*self.correction_warp_ids,self.mma_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr,barrier_for_retrieve=tmem_bar,allocator_warp_id=self.softmax_warp_ids[0],is_two_cta=cute.size(qk_mma.thr_id.shape)==2,two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc) - if warp_idx == self.empty_warp_id: - cute.arch.mbarrier_init(st.tmem_dealloc, 32*len((*self.softmax_warp_ids,*self.correction_warp_ids))) - cute.arch.mbarrier_init_fence() - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk,is_relaxed=True) - sQ = smem.allocate_tensor(element_type=self.q_dtype,layout=q_smem_s.outer,byte_alignment=128,swizzle=q_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype,layout=k_smem_s.outer,byte_alignment=128,swizzle=k_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype,layout=v_smem_s.outer,byte_alignment=128,swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype,layout=c_smem_s.outer,byte_alignment=128,swizzle=c_smem_s.inner) - gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None)) - gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None)) - gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None)) - gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None)) - n_kv_tiles = cute.size(gK, mode=[3]) - qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - tCgV = pv_thr.partition_B(gV); tCgC = pv_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,0,None,0)).shape) - tAsQ,tAgQ = cpasync.tma_partition(tma_q,0,a_lay,cute.group_modes(sQ,0,3),cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape) - tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3)) - tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)]; tVgV = tVgV[(None,0,None,0)] - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None,None,None,0)] - tOrP0 = cute.make_tensor(tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, tOrP.layout) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_as, 1)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ==================== TMA WARP (9) ==================== - if warp_idx == self.tma_warp_id: - qp.reset(); qh = qp.acquire_and_advance() - cute.copy(tma_q,tAgQ[(None,qh.count)],tAsQ[(None,qh.index)],tma_bar_ptr=qh.barrier) - qp.tail() - kvp.reset(); pk = kvp.try_acquire() - for kt in cutlass.range(n_kv_tiles,unroll=1): - kh = kvp.acquire_and_advance(pk) - cute.copy(tma_k,tBgK[(None,kh.count)],tBsK[(None,kh.index)],tma_bar_ptr=kh.barrier) - pk = cutlass.Boolean(1) - vh = kvp.acquire_and_advance(pk) - cute.copy(tma_v,tVgV[(None,vh.count)],tVsV[(None,vh.index)],tma_bar_ptr=vh.barrier) - pk = cutlass.Boolean(1) - kvp.tail() - - # ==================== MMA WARP (8) ==================== - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - qc.reset(); qh = qc.wait_and_advance(); qh.release() - kvc.reset(); pk = kvc.try_wait() - for kt in range(n_kv_tiles): - kh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - sh = mma_s_prod.acquire_and_advance() - qk_mma.set(tcgen05.Field.ACCUMULATE, False) - for kb in cutlass.range(cute.size(tCrQ,mode=[2]), unroll_full=True): - cute.gemm(qk_mma, tStS0, tCrQ[(None,None,kb,0)], tCrK[(None,None,kb,kh.index)], tStS0) - qk_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - sh.commit(); kh.release() - # MMA waits for softmax to produce P (softmax consumes S, releases when P ready) - # In the pipeline model, the S release by softmax IS the P-ready signal - # But with PipelineUmmaAsync, the consumer release releases the producer handle - # So after the softmax releases, the MMA can acquire the next S handle - - vh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - oh = mma_corr_pipe.producer_acquire(pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.mma_corr_stage)) - pv_mma.set(tcgen05.Field.ACCUMULATE, kt != 0) - for kb in cutlass.range(cute.size(tOrP0,mode=[2]), unroll_full=True): - cute.gemm(pv_mma, tOtO0, tOrP0[(None,None,kb)], tCrV[(None,None,kb,vh.index)], tOtO0) - pv_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - oh.commit(); vh.release() - mma_s_prod.tail() - mma_corr_prod.tail() - cute.arch.relinquish_tmem_alloc_permit() - cute.arch.mbarrier_wait(st.tmem_dealloc, 0) - tmem_ptr = cute.arch.retrieve_tmem_ptr(self.qk_acc_dtype, alignment=16, ptr_to_buffer_holding_addr=st.holding) - cute.arch.dealloc_tmem(tmem_ptr, Int32(self.num_tmem_alloc_cols)) - - # ==================== SOFTMAX WARPS (0-3) — identity ==================== - if warp_idx < len(self.softmax_warp_ids): - tmem.allocate(self.num_tmem_alloc_cols) - tmem.wait_for_alloc() - sfw_idx = tidx % (32 * len(self.softmax_warp_ids)) - tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS0) - thr_load = tiled_tmem_load.get_slice(sfw_idx) - tTMEM_LOADtS = thr_load.partition_S(tStS0) - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tTMEM_LOADcS = thr_load.partition_D(tScS) - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - tStP_layout = cute.composition(tStS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tStP0 = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStP_layout) - tmem_store_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStP0) - thr_store = tiled_tmem_store.get_slice(sfw_idx) - tTMEM_STOREtP = thr_store.partition_D(tStP0) - tScP_layout = cute.composition(tScS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tTMEM_STOREcP = thr_store.partition_S(cute.make_tensor(tScS.iterator, tScP_layout)) - vec_handle = s_corr_prod.acquire_and_advance() - for kt in range(n_kv_tiles): - si_handle = mma_s_cons.wait_and_advance() - tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS) - cute.arch.fence_view_async_tmem_load() - rP_words = cute.make_rmem_tensor(tTMEM_STOREcP.shape, self.qk_acc_dtype) - rP_bf16 = cute.make_tensor(cute.recast_ptr(rP_words.iterator, dtype=self.q_dtype), tTMEM_LOADrS.layout) - frg_cnt = 4 - frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt - tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile)) - rP_bf16_frg = cute.logical_divide(rP_bf16, cute.make_layout(frg_tile)) - for j in range(frg_cnt): - s_vec = tTMEM_LOADrS_frg[None, j].load() - rP_bf16_frg[None, j].store(s_vec.to(self.q_dtype)) - cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP) - cute.arch.fence_view_async_tmem_store() - vec_handle.commit() - si_handle.release() - vec_handle = s_corr_prod.acquire_and_advance() - s_corr_prod.tail() - cute.arch.mbarrier_arrive(st.tmem_dealloc) - tmem.relinquish_alloc_permit() - - # ==================== CORRECTION WARPS (4-7) — identity, no epilogue ==================== - if warp_idx >= len(self.softmax_warp_ids) and warp_idx < len(self.softmax_warp_ids) + len(self.correction_warp_ids): - tmem.wait_for_alloc() - corr_idx = tidx % (32 * len(self.correction_warp_ids)) - first_vec = s_corr_cons.wait_and_advance() - first_vec.release() - for kt in range(n_kv_tiles - 1): - vec = s_corr_cons.wait_and_advance() - o = mma_corr_cons.wait_and_advance() - o.release() - vec.release() - final_vec = s_corr_cons.wait_and_advance() - final_vec.release() - final_o = mma_corr_cons.wait_and_advance() - # Write O from TMEM to output using the epilogue pipeline - epi_handle = corr_epi_prod.acquire_and_advance() - tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) - tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout) - # Use epilogue_tma_store with a fresh consumer state - # The acc_pipe is the pipeline we already consumed from, but - # epilogue_tma_store wants a consumer. Since we already have O, - # skip the acc_pipe wait by using a dummy pipeline. - # Actually, just do a simple TMA copy from sC - # For the minimal test, just signal the epilogue and move on - epi_handle.commit() - final_o.release() - cute.arch.mbarrier_arrive(st.tmem_dealloc) - - # ==================== EPILOGUE WARP (10) ==================== - if warp_idx == self.epi_warp_id: - epi_handle = corr_epi_cons.wait_and_advance() - epi_handle.release() - - -def test(): - torch.manual_seed(42) - for n in [128]: - m, hd = 128, HEAD_DIM - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(n, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(n, hd, dtype=torch.bfloat16, device='cuda') - v_kernel = v.unsqueeze(-1) - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - qf = q[:,:,0].float(); kf = k[:,:,0].float() - ref = (qf @ kf.T).bfloat16().float() @ v.float() - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - kernel = FmhaV3StageCMin(s_k=n) - print(f'n={n}: Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - print(f'n={n}: Running...', flush=True) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - out = c[:,:,0].float() - cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - print(f'FMHA stage-C min n={n}: cosine {cos:.6f}') - -if __name__ == '__main__': - test() diff --git a/tests/archive/unit_test_fmha_v3_vec_c9.py b/tests/archive/unit_test_fmha_v3_vec_c9.py deleted file mode 100644 index 061f77fe..00000000 --- a/tests/archive/unit_test_fmha_v3_vec_c9.py +++ /dev/null @@ -1,488 +0,0 @@ -""" -FMHA v3 + Stage C: QK -> online softmax -> PV with KV-tile interleaving. -Stage C: row_max, exp2, O rescale, row_sum, final normalization. -FMHA pattern P store preserved from Stage B. -""" -import math -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - -HEAD_DIM = 64 - -class FmhaV3Softmax: - def __init__(self): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0,1,2,3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.num_c_stage = 2 - self.kv_stage = 2; self.q_stage = 1; self.num_c_stage = 2 - - def _setup(self, qk_mma, pv_mma): - qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (128, 128, qk_ik * 4) - pv_ik = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (128, HEAD_DIM, pv_ik * (128 // pv_ik)) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0]//cute.size(qk_mma.thr_id.shape), HEAD_DIM, self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - self.q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, self.q_dtype, self.q_stage) - self.k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, self.q_dtype, self.kv_stage) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, self.kv_stage) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - qk_thr = qk_mma.get_slice(0); qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - pv_thr = pv_mma.get_slice(0); pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - self.tmem_s0_offset = 0; self.tmem_p0_offset = 32 - # P occupies [tmem_p0_offset, tmem_p0_offset + p_cols_fp32) - # S occupies [0, qk_mma_tiler[1]) = [0, 128) - # O must NOT overlap P. Place O after max(S end, P end), aligned to 32. - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - p_end = self.tmem_p0_offset + p_cols_fp32 # 32 + 64 = 96 - s_cols = self.qk_mma_tiler[1] # 128 - o_after = max(s_cols, p_end) # 128 - self.tmem_o0_offset = ((o_after + 31) // 32) * 32 - self.tmem_vec_offset = 0 # Reuse S region for per-row inv_row_sum vector # align to 32 = 128 - self.tmem_vec_offset = 0 # Reuse S region (free after softmax loop) - o_cols = find_tmem_tensor_col_offset(tOtO) # footprint of O - total = self.tmem_o0_offset + o_cols - # Must be multiple of 32 AND power of 2 - self.num_tmem_alloc_cols = 1 - while self.num_tmem_alloc_cols < total: - self.num_tmem_alloc_cols *= 2 - cta = cute.size(qk_mma.thr_id.shape) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)) - self.q_tx_bytes = cute.size_in_bytes(self.q_dtype, q_s) * cta - self.kv_tx_bytes = cute.size_in_bytes(self.q_dtype, k_s) * cta - self.scale_softmax_log2 = Float32(1.0 / math.sqrt(HEAD_DIM) * math.log2(math.e)) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - # # s_k hardcoded # BROKEN in @cute.jit - # FMHA-style V: reconstruct as (HEAD_DIM, s_k, 1) MN-major - v_fmha = cute.make_tensor( - v.iterator, - cute.make_layout( - (HEAD_DIM, 128, 1), - stride=(1, HEAD_DIM, HEAD_DIM * 128), - ), - ) - self.v_major = LayoutEnum.from_tensor(v_fmha).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, self.a_major, self.b_major, self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, self.qk_acc_dtype, self.cta_group, (128,HEAD_DIM), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)); v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id),q,q_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id),k,k_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v_fmha,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape) - epi_s = cute.select(self.c_smem_s,mode=[0,1]) - tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile) - self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.v_smem_s,self.p_tmem_s,self.c_smem_s,self.epi_tile).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, q_smem_s, k_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx,_,_ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k); cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - q_bar: cute.struct.MemRange[cutlass.Int64, self.q_stage*2] - kv_bar: cute.struct.MemRange[cutlass.Int64, self.kv_stage*2] - s_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage*2] - tmem_dealloc: cutlass.Int64; holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - qp,qc = pipeline.PipelineTmaUmma.create(barrier_storage=st.q_bar.data_ptr(),num_stages=self.q_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.q_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - kvp,kvc = pipeline.PipelineTmaUmma.create(barrier_storage=st.kv_bar.data_ptr(),num_stages=self.kv_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.kv_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - s_prod,s_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.s_bar.data_ptr(),num_stages=1,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.epilogue_warp_id))).make_participants() - softmax_done_bar = pipeline.NamedBarrier(barrier_id=3, num_threads=32 + 32*len(self.epilogue_warp_id)) - pv_done_bar = pipeline.NamedBarrier(barrier_id=4, num_threads=32 + 32*len(self.epilogue_warp_id)) - acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(),num_stages=self.num_acc_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,len(self.epilogue_warp_id)),cta_layout_vmnk=cl_vmnk,defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2,num_threads=32*len((self.mma_warp_id,*self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr,barrier_for_retrieve=tmem_bar,allocator_warp_id=self.epilogue_warp_id[0],is_two_cta=cute.size(qk_mma.thr_id.shape)==2,two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk,is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype,layout=q_smem_s.outer,byte_alignment=128,swizzle=q_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype,layout=k_smem_s.outer,byte_alignment=128,swizzle=k_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype,layout=v_smem_s.outer,byte_alignment=128,swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype,layout=c_smem_s.outer,byte_alignment=128,swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None)) - gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None)) - gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None)) - gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None)) - n_kv_tiles = cute.size(gK, mode=[3]) - - qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - tCgV = pv_thr.partition_B(gV); tCgC = pv_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,0,None,0)).shape) - tAsQ,tAgQ = cpasync.tma_partition(tma_q,0,a_lay,cute.group_modes(sQ,0,3),cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape) - tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3)) - tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)]; tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - # --- PV read view (for MMA only, NOT for softmax store) --- - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None,None,None,0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # TMA LOAD - if warp_idx == self.tma_warp_id: - qp.reset(); qh = qp.acquire_and_advance() - cute.copy(tma_q,tAgQ[(None,qh.count)],tAsQ[(None,qh.index)],tma_bar_ptr=qh.barrier) - qp.tail() - kvp.reset(); pk = kvp.try_acquire() - for kt in cutlass.range(n_kv_tiles,unroll=1): - kh = kvp.acquire_and_advance(pk) - cute.copy(tma_k,tBgK[(None,kh.count)],tBsK[(None,kh.index)],tma_bar_ptr=kh.barrier) - pk = cutlass.Boolean(1) - vh = kvp.acquire_and_advance(pk) - cute.copy(tma_v,tVgV[(None,vh.count)],tVsV[(None,vh.index)],tma_bar_ptr=vh.barrier) - pk = cutlass.Boolean(1) - kvp.tail() - - # MMA - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - qc.reset(); qh = qc.wait_and_advance(); qh.release() - kvc.reset(); pk = kvc.try_wait() - acc_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_stage) - acc_pipe.producer_acquire(acc_st) - for kt in range(n_kv_tiles): - kh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - sh = s_prod.acquire_and_advance() - qk_mma.set(tcgen05.Field.ACCUMULATE, False) - for kb in cutlass.range(cute.size(tCrQ,mode=[2]), unroll_full=True): - cute.gemm(qk_mma, tStS0, tCrQ[(None,None,kb,0)], tCrK[(None,None,kb,kh.index)], tStS0) - qk_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - sh.commit(); kh.release() - softmax_done_bar.arrive_and_wait() - vh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - pv_mma.set(tcgen05.Field.ACCUMULATE, kt != 0) - for kb in cutlass.range(cute.size(tOrP0,mode=[2]), unroll_full=True): - cute.gemm(pv_mma, tOtO0, tOrP0[(None,None,kb)], tCrV[(None,None,kb,vh.index)], tOtO0) - pv_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - vh.release() - pv_done_bar.arrive() - acc_pipe.producer_commit(acc_st); acc_st.advance() - acc_pipe.producer_tail(acc_st) - - # ===================== EPILOGUE WARPS (STAGE C: ONLINE SOFTMAX) ===================== - if warp_idx < self.mma_warp_id: - tmem.allocate(self.num_tmem_alloc_cols) - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) - sfw_idx = tidx % (32 * len(self.epilogue_warp_id)) - - # --- S load (QK C-fragment) --- - tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS0) - thr_load = tiled_tmem_load.get_slice(sfw_idx) - tTMEM_LOADtS = thr_load.partition_S(tStS0) - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tTMEM_LOADcS = thr_load.partition_D(tScS) - - # --- P store (QK C-fragment composition, FMHA pattern) --- - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - tStP_layout = cute.composition(tStS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tStP0 = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStP_layout) - tmem_store_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStP0) - thr_store = tiled_tmem_store.get_slice(sfw_idx) - tTMEM_STOREtP = thr_store.partition_D(tStP0) - tScP_layout = cute.composition(tScS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tScP = cute.make_tensor(tScS.iterator, tScP_layout) - tTMEM_STOREcP = thr_store.partition_S(tScP) - - # --- Vector TMEM (per-row row_sum storage, FMHA pattern) --- - # composition(tStS.layout, (128, 2)) = 2 FP32 columns per logical row - # vec[0] = row_sum (final, after loop), vec[1] = unused - # Reuses S TMEM region (offset 0), free after softmax loop writes - - tStS_vec_layout = cute.composition(tStS.layout, cute.make_layout((128, 2))) - tStS_vec = cute.make_tensor(tStS.iterator + self.tmem_vec_offset, tStS_vec_layout) - tScS_vec_layout = cute.composition(tScS.layout, cute.make_layout((128, 2))) - tScS_vec = cute.make_tensor(tScS.iterator, tScS_vec_layout) - tmem_store_vec_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(2)), self.qk_acc_dtype) - tiled_tmem_store_vec = tcgen05.make_tmem_copy(tmem_store_vec_atom, tStS_vec) - thr_tmem_store_vec = tiled_tmem_store_vec.get_slice(sfw_idx) - tTMEM_STORE_VECtS = thr_tmem_store_vec.partition_D(tStS_vec) - tTMEM_STORE_VECcS = thr_tmem_store_vec.partition_S(tScS_vec) - tmem_load_vec_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(2)), self.qk_acc_dtype) - tiled_tmem_load_vec = tcgen05.make_tmem_copy(tmem_load_vec_atom, tStS_vec) - thr_tmem_load_vec = tiled_tmem_load_vec.get_slice(sfw_idx) - tTMEM_LOAD_VECtS = thr_tmem_load_vec.partition_S(tStS_vec) - tTMEM_LOAD_VECcS = thr_tmem_load_vec.partition_D(tScS_vec) - - # --- C6: O TMEM load/store for rescale (correction_rescale pattern) --- - corr_tile_size = 16 - cO = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1])) - tOcO = pv_thr.partition_C(cO) - o_tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), self.qk_acc_dtype) - o_tmem_store_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), self.qk_acc_dtype) - tOtO_i_layout = cute.composition(tOtO0.layout, cute.make_layout((128, corr_tile_size))) - tOcO_i_layout = cute.composition(tOcO.layout, cute.make_layout((128, corr_tile_size))) - tOtO_i = cute.make_tensor(tOtO0.iterator, tOtO_i_layout) - tOcO_i = cute.make_tensor(tOcO.iterator, tOcO_i_layout) - o_tiled_tmem_load = tcgen05.make_tmem_copy(o_tmem_load_atom, tOtO_i) - o_tiled_tmem_store = tcgen05.make_tmem_copy(o_tmem_store_atom, tOtO_i) - o_thr_load = o_tiled_tmem_load.get_slice(sfw_idx) - o_thr_store = o_tiled_tmem_store.get_slice(sfw_idx) - tTMEM_LOADtO = o_thr_load.partition_S(tOtO_i) - tTMEM_LOADcO = o_thr_load.partition_D(tOcO_i) - tTMEM_STOREtO = o_thr_store.partition_D(tOtO_i) - o_col_tiles = self.pv_mma_tiler[1] // corr_tile_size - - # --- C2: Per-thread row state (persist across KV tiles) --- - row_max = -cutlass.Float32.inf - row_sum = cutlass.Float32(0.0) - - # --- C3: QK scale = 1/sqrt(HEAD_DIM) * log2(e) for exp2 --- - scale = self.scale_softmax_log2 - - # ============================================================= - # Per-KV-tile online softmax loop - # ============================================================= - for kt in range(n_kv_tiles): - si_handle = s_cons.wait_and_advance() - - # Load S from TMEM (FP32, QK C-fragment layout) - tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS) - - # --- C4: Compute tile_max via .reduce(MAX) --- - old_row_max = row_max - row_max = tTMEM_LOADrS.load().reduce(cute.ReductionOp.MAX, row_max, 0) - row_max_safe = row_max - if row_max == -cutlass.Float32.inf: - row_max_safe = cutlass.Float32(0.0) - - # --- C5: Compute rescale factor --- - acc_scale = cute.math.exp2(scale * (old_row_max - row_max_safe), fastmath=True) - - # --- C6: Rescale O in TMEM (load O, multiply by acc_scale, store O) --- - # acc_scale belongs to QK row (N//4), but O rows are in PV partition (N). - # Store acc_scale to vector by QK row, read by PV row. - if kt > 0: - pv_done_bar.arrive_and_wait() - - # Store acc_scale to vector indexed by QK logical row - qk_row_c6 = tTMEM_LOADcS[0][0] - thr_vs_c6 = tiled_tmem_store_vec.get_slice(qk_row_c6) - tVStore_c6 = thr_vs_c6.partition_D(tStS_vec) - tVStoreSrc_c6 = thr_vs_c6.partition_S(tScS_vec) - tVStoreRmem_c6 = cute.make_rmem_tensor(tVStoreSrc_c6.shape, self.qk_acc_dtype) - tVStoreRmem_c6[0] = acc_scale - cute.copy(tiled_tmem_store_vec, tVStoreRmem_c6, tVStore_c6) - cute.arch.fence_view_async_tmem_store() - - # Read acc_scale from vector indexed by PV logical row - pv_row_c6 = tTMEM_LOADcO[0][0] - thr_vl_c6 = tiled_tmem_load_vec.get_slice(pv_row_c6) - tVLoad_c6 = thr_vl_c6.partition_S(tStS_vec) - tVLoadDst_c6 = thr_vl_c6.partition_D(tScS_vec) - tVLoadRmem_c6 = cute.make_rmem_tensor(tVLoadDst_c6.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load_vec, tVLoad_c6, tVLoadRmem_c6) - cute.arch.fence_view_async_tmem_load() - acc_scale_pv = tVLoadRmem_c6[0] - - tTMrO = cute.make_rmem_tensor((tTMEM_LOADcO.shape, o_col_tiles), self.qk_acc_dtype) - for i in range(o_col_tiles): - tTMrO_i_ = tTMrO[None, i] - tTMrO_i_layout = cute.composition(tTMrO_i_.layout, cute.make_layout(tTMrO.shape[0])) - tTMrO_i = cute.make_tensor(tTMrO_i_.iterator, tTMrO_i_layout) - tTMEM_LOADtO_i = cute.make_tensor(tTMEM_LOADtO.iterator + i * corr_tile_size, tTMEM_LOADtO.layout) - tTMEM_STOREtO_i = cute.make_tensor(tTMEM_STOREtO.iterator + i * corr_tile_size, tTMEM_STOREtO.layout) - cute.copy(o_tiled_tmem_load, tTMEM_LOADtO_i, tTMrO_i) - for j in cutlass.range(cute.size(tTMrO_i), vectorize=True): - tTMrO_i[j] = tTMrO_i[j] * acc_scale_pv - cute.copy(o_tiled_tmem_store, tTMrO_i, tTMEM_STOREtO_i) - cute.arch.fence_view_async_tmem_store() - - # Rescale row_sum - row_sum = row_sum * acc_scale - - # --- C7: Compute P = exp2((S - row_max_safe) * scale) --- - minus_row_max_scale = (cutlass.Float32(0.0) - row_max_safe) * scale - - # Register bridge (FMHA pattern: FP32 backing + BF16 view) - rP_words = cute.make_rmem_tensor(tTMEM_STOREcP.shape, self.qk_acc_dtype) - rP_bf16 = cute.make_tensor(cute.recast_ptr(rP_words.iterator, dtype=self.q_dtype), tTMEM_LOADrS.layout) - - frg_cnt = 4 - frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt - tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile)) - rP_bf16_frg = cute.logical_divide(rP_bf16, cute.make_layout(frg_tile)) - - # Scale S, compute exp2, store through register bridge - for j in range(frg_cnt): - for k in cutlass.range(cute.size(tTMEM_LOADrS_frg, mode=[0]), vectorize=True): - tTMEM_LOADrS_frg[k, j] = tTMEM_LOADrS_frg[k, j] * scale + minus_row_max_scale - tTMEM_LOADrS_frg[k, j] = cute.math.exp2(tTMEM_LOADrS_frg[k, j], fastmath=True) - s_vec = tTMEM_LOADrS_frg[None, j].load() - rP_bf16_frg[None, j].store(s_vec.to(self.q_dtype)) - - # Store P to TMEM - cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP) - cute.arch.fence_view_async_tmem_store() - si_handle.release() - softmax_done_bar.arrive() - - # --- C8: Row sum accumulation (CUTLASS FMHA packed f32x2 pattern) --- - # P values still in tTMEM_LOADrS registers. - # 4 accumulators for 4 reduction_unroll columns. - local_row_sum_0 = (cutlass.Float32(0.0), cutlass.Float32(0.0)) - local_row_sum_1 = (cutlass.Float32(0.0), cutlass.Float32(0.0)) - local_row_sum_2 = (cutlass.Float32(0.0), cutlass.Float32(0.0)) - local_row_sum_3 = (cutlass.Float32(0.0), cutlass.Float32(0.0)) - - reduction_unroll = 4 - rfrg_tile = cute.size(tTMEM_LOADrS) // reduction_unroll - tTMEM_LOADrS_rfrg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(rfrg_tile)) - - for j in cutlass.range_constexpr(0, cute.size(tTMEM_LOADrS_rfrg, mode=[0]), 2): - local_row_sum_0 = cute.arch.add_packed_f32x2( - local_row_sum_0, (tTMEM_LOADrS_rfrg[j, 0], tTMEM_LOADrS_rfrg[j + 1, 0])) - local_row_sum_1 = cute.arch.add_packed_f32x2( - local_row_sum_1, (tTMEM_LOADrS_rfrg[j, 1], tTMEM_LOADrS_rfrg[j + 1, 1])) - local_row_sum_2 = cute.arch.add_packed_f32x2( - local_row_sum_2, (tTMEM_LOADrS_rfrg[j, 2], tTMEM_LOADrS_rfrg[j + 1, 2])) - local_row_sum_3 = cute.arch.add_packed_f32x2( - local_row_sum_3, (tTMEM_LOADrS_rfrg[j, 3], tTMEM_LOADrS_rfrg[j + 1, 3])) - - local_row_sum_0 = cute.arch.add_packed_f32x2(local_row_sum_0, local_row_sum_1) - local_row_sum_2 = cute.arch.add_packed_f32x2(local_row_sum_2, local_row_sum_3) - local_row_sum_0 = cute.arch.add_packed_f32x2(local_row_sum_0, local_row_sum_2) - tile_sum = local_row_sum_0[0] + local_row_sum_0[1] - - row_sum = row_sum + tile_sum - - # --- C9: Final normalization via O TMEM rescale --- - pv_done_bar.arrive_and_wait() - - # Write row_sum to TMEM vector using QK partition (correct row mapping) - qk_row_c9 = tTMEM_LOADcS[0][0] - thr_vs_c9 = tiled_tmem_store_vec.get_slice(qk_row_c9) - tVStore_c9 = thr_vs_c9.partition_D(tStS_vec) - tVStoreSrc_c9 = thr_vs_c9.partition_S(tScS_vec) - tVStoreRmem_c9 = cute.make_rmem_tensor(tVStoreSrc_c9.shape, self.qk_acc_dtype) - tVStoreRmem_c9[0] = row_sum - cute.copy(tiled_tmem_store_vec, tVStoreRmem_c9, tVStore_c9) - cute.arch.fence_view_async_tmem_store() - - # Read row_sum from TMEM vector using PV partition (correct for O rows) - pv_row_c9 = tTMEM_LOADcO[0][0] - thr_vl_c9 = tiled_tmem_load_vec.get_slice(pv_row_c9) - tVLoad_c9 = thr_vl_c9.partition_S(tStS_vec) - tVLoadDst_c9 = thr_vl_c9.partition_D(tScS_vec) - tVLoadRmem_c9 = cute.make_rmem_tensor(tVLoadDst_c9.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load_vec, tVLoad_c9, tVLoadRmem_c9) - cute.arch.fence_view_async_tmem_load() - pv_row_sum = tVLoadRmem_c9[0] - - inv_row_sum = cutlass.Float32(1.0) / pv_row_sum - - # Normalize O in TMEM using PV-correct inv_row_sum - tTMrO_final = cute.make_rmem_tensor((tTMEM_LOADcO.shape, o_col_tiles), self.qk_acc_dtype) - for i in range(o_col_tiles): - tTMrO_i_ = tTMrO_final[None, i] - tTMrO_i_layout = cute.composition(tTMrO_i_.layout, cute.make_layout(tTMrO_final.shape[0])) - tTMrO_i = cute.make_tensor(tTMrO_i_.iterator, tTMrO_i_layout) - tTMEM_LOADtO_i = cute.make_tensor( - tTMEM_LOADtO.iterator + i * corr_tile_size, tTMEM_LOADtO.layout) - tTMEM_STOREtO_i = cute.make_tensor( - tTMEM_STOREtO.iterator + i * corr_tile_size, tTMEM_STOREtO.layout) - cute.copy(o_tiled_tmem_load, tTMEM_LOADtO_i, tTMrO_i) - for j in cutlass.range(cute.size(tTMrO_i), vectorize=True): - tTMrO_i[j] = tTMrO_i[j] * inv_row_sum - cute.copy(o_tiled_tmem_store, tTMrO_i, tTMEM_STOREtO_i) - cute.arch.fence_view_async_tmem_store() - - # Now O in TMEM is normalized. Use standard epilogue_tma_store with identity. - tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout) - acc_cons_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage) - c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) - c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp) - acc_cons_st = utils.gemm.sm100.epilogue_tma_store( - self, tidx, warp_idx, tma_c, tCtO_base, sC, tCgC, epi_tile, 0, - const_expr(lambda x: x), - (0,0,0), acc_cons_st, acc_pipe, c_pipe) - c_pipe.producer_tail() - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def test(): - import math - torch.manual_seed(42) - for n in [128, 256, 384]: - m, hd = 128, HEAD_DIM - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device="cuda") - k = torch.randn(n, hd, 1, dtype=torch.bfloat16, device="cuda") - v = torch.randn(n, hd, dtype=torch.bfloat16, device="cuda") - v_kernel = v.unsqueeze(-1) - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device="cuda") - qf = q[:,:,0].float(); kf = k[:,:,0].float() - attn = qf @ kf.T / math.sqrt(hd) - ref = torch.softmax(attn, dim=-1) @ v.float() - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - kernel = FmhaV3Softmax() - print(f"n={n}: Compiling...", flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - print(f"n={n}: tmem: s0={kernel.tmem_s0_offset} p0={kernel.tmem_p0_offset} o0={kernel.tmem_o0_offset} vec={kernel.tmem_vec_offset} alloc={kernel.num_tmem_alloc_cols}", flush=True) - print(f"n={n}: Running...", flush=True) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - out = c[:,:,0].float() - cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - max_err = (out - ref).abs().max().item() - print(f"FMHA softmax n={n}: cosine {cos:.6f} max_err {max_err:.6f} {'PASS' if cos >= 0.999 else 'FAIL'}", flush=True) - -if __name__ == "__main__": - test() - - diff --git a/tests/archive/unit_test_pv64_with_softmax.py b/tests/archive/unit_test_pv64_with_softmax.py deleted file mode 100644 index a92680fb..00000000 --- a/tests/archive/unit_test_pv64_with_softmax.py +++ /dev/null @@ -1,257 +0,0 @@ -""" -Test (128,64) PV WITH identity softmax, single AB pipeline. -This is test_pv64.py but we add a print to see if softmax actually writes. -The key: if P/A alias works (proven above), then 0.67 must be from softmax -writing to wrong columns or V being wrong. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - -HEAD_DIM = 64 - -class Pv64WithSoftmax: - def __init__(self): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0,1,2,3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.num_c_stage = 2 - self.num_ab_stage = 1; self.num_acc_stage = 1 - - def _setup(self, qk_mma, pv_mma): - qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (128, 128, qk_ik * 4) - pv_ik = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (128, HEAD_DIM, pv_ik * (128 // pv_ik)) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0]//cute.size(qk_mma.thr_id.shape), HEAD_DIM, self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - qk_thr = qk_mma.get_slice(0); qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - pv_thr = pv_mma.get_slice(0); pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0; self.tmem_p0_offset = 32 - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - p_end = self.tmem_p0_offset + p_cols_fp32 - s_cols = self.qk_mma_tiler[1] - o_after = max(s_cols, p_end) - self.tmem_o0_offset = ((o_after + 31) // 32) * 32 - tCS = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage)) - tCO = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCS, tCO], arch="sm_100") - a_s = cute.slice_(self.a_smem_s,(None,None,None,0)); b_s = cute.slice_(self.b_smem_s,(None,None,None,0)) - v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - self.num_tma_load_bytes = (cute.size_in_bytes(self.q_dtype,a_s)+cute.size_in_bytes(self.q_dtype,b_s)+cute.size_in_bytes(self.q_dtype,v_s))*cute.size(qk_mma.thr_id.shape) - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - v_fmha = cute.make_tensor( - v.iterator, - cute.make_layout( - (HEAD_DIM, 128, 1), - stride=(1, HEAD_DIM, HEAD_DIM * 128), - ), - ) - self.v_major = LayoutEnum.from_tensor(v_fmha).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, self.a_major, self.b_major, self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, self.qk_acc_dtype, self.cta_group, (128,HEAD_DIM), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - q_s = cute.slice_(self.a_smem_s,(None,None,None,0)); k_s = cute.slice_(self.b_smem_s,(None,None,None,0)) - v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id),q,q_s,self.mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id),k,k_s,self.mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v_fmha,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape) - epi_s = cute.select(self.c_smem_s,mode=[0,1]) - tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile) - self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.a_smem_s,self.b_smem_s,self.v_smem_s,self.p_tmem_s,self.c_smem_s,self.epi_tile).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx,_,_ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - @cute.struct - class SS: - ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage*2] - mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage*2] - tmem_dealloc: cutlass.Int64; holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - ab_p,ab_c = pipeline.PipelineTmaUmma.create(barrier_storage=st.ab_bar.data_ptr(),num_stages=self.num_ab_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.num_tma_load_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - mma_si_prod,mma_si_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.mma_si_bar.data_ptr(),num_stages=1,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.epilogue_warp_id))).make_participants() - softmax_done_bar = pipeline.NamedBarrier(barrier_id=3, num_threads=32 + 32*len(self.epilogue_warp_id)) - acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(),num_stages=self.num_acc_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,len(self.epilogue_warp_id)),cta_layout_vmnk=cl_vmnk,defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2,num_threads=32*len((self.mma_warp_id,*self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr,barrier_for_retrieve=tmem_bar,allocator_warp_id=self.epilogue_warp_id[0],is_two_cta=cute.size(qk_mma.thr_id.shape)==2,two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk,is_relaxed=True) - sQ = smem.allocate_tensor(element_type=self.q_dtype,layout=a_smem_s.outer,byte_alignment=128,swizzle=a_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype,layout=b_smem_s.outer,byte_alignment=128,swizzle=b_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype,layout=v_smem_s.outer,byte_alignment=128,swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype,layout=c_smem_s.outer,byte_alignment=128,swizzle=c_smem_s.inner) - gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None)) - gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None)) - gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None)) - gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None)) - k_cnt = cute.size(gQ, mode=[3]) - qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - tCgV = pv_thr.partition_B(gV); tCgC = pv_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,0,None,0)).shape) - tAsQ,tAgQ = cpasync.tma_partition(tma_q,0,a_lay,cute.group_modes(sQ,0,3),cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape) - tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3)) - tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)]; tVgV = tVgV[(None,0,None,0)] - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - tStS0 = cute.make_tensor(tStS.iterator+self.tmem_s0_offset,tStS.layout) - pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - tOtO0 = cute.make_tensor(tOtO.iterator+self.tmem_o0_offset,tOtO.layout) - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None,None,None,0)] - tOrP0 = cute.make_tensor(tOrP.iterator+self.qk_acc_dtype.width//self.q_dtype.width*self.tmem_p0_offset,tOrP.layout) - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_as,self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_as,self.num_acc_stage)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # TMA LOAD - if warp_idx == self.tma_warp_id: - ab_p.reset(); peek = ab_p.try_acquire() - for kt in cutlass.range(k_cnt,unroll=1): - h = ab_p.acquire_and_advance(peek) - cute.copy(tma_q,tAgQ[(None,h.count)],tAsQ[(None,h.index)],tma_bar_ptr=h.barrier) - cute.copy(tma_k,tBgK[(None,h.count)],tBsK[(None,h.index)],tma_bar_ptr=h.barrier) - cute.copy(tma_v,tVgV[(None,h.count)],tVsV[(None,h.index)],tma_bar_ptr=h.barrier) - peek = cutlass.Boolean(1) - if h.count+1= 0.99 else "FAIL"}") - -if __name__ == "__main__": - test() diff --git a/tests/archive/unit_test_qk_softmax.py b/tests/archive/unit_test_qk_softmax.py deleted file mode 100644 index d465f448..00000000 --- a/tests/archive/unit_test_qk_softmax.py +++ /dev/null @@ -1,252 +0,0 @@ -""" -Debug: QK + identity softmax, output P (BF16) to GMEM. -Tests the full QK -> softmax -> TMEM pipeline without PV. -Uses the QK C-fragment store (like FMHA). -Output = S.bf16() which is (Q@K^T).bfloat16(), shape (128, 128). -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - -HEAD_DIM = 64 - -class QkSoftmaxTest: - def __init__(self): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0,1,2,3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.num_c_stage = 2 - self.kv_stage = 2; self.q_stage = 1 - - def _setup(self, qk_mma): - qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (128, 128, qk_ik * 4) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0]//cute.size(qk_mma.thr_id.shape), 128, self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - self.q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, self.q_dtype, self.q_stage) - self.k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, self.q_dtype, self.kv_stage) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - qk_thr = qk_mma.get_slice(0); qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0; self.tmem_p0_offset = 32; self.tmem_o0_offset = 0 - tCS = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCS], arch="sm_100") - cta = cute.size(qk_mma.thr_id.shape) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)) - self.q_tx_bytes = cute.size_in_bytes(self.q_dtype, q_s) * cta - self.kv_tx_bytes = cute.size_in_bytes(self.q_dtype, k_s) * cta - - @cute.jit - def __call__(self, q, k, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, self.a_major, self.b_major, self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM) - self._setup(qk_mma) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)) - tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id),q,q_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id),k,k_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - epi_s = cute.select(self.c_smem_s,mode=[0,1]) - tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile) - self._kernel(qk_mma,tma_q,mQ,tma_k,mK,tma_c,mC,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.c_smem_s,self.epi_tile).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, tma_q, mQ, tma_k, mK, tma_c, mC, cl_vmnk, q_smem_s, k_smem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx,_,_ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - q_bar: cute.struct.MemRange[cutlass.Int64, self.q_stage*2] - kv_bar: cute.struct.MemRange[cutlass.Int64, self.kv_stage*2] - s_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage*2] - tmem_dealloc: cutlass.Int64; holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - qp,qc = pipeline.PipelineTmaUmma.create(barrier_storage=st.q_bar.data_ptr(),num_stages=self.q_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.q_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - kvp,kvc = pipeline.PipelineTmaUmma.create(barrier_storage=st.kv_bar.data_ptr(),num_stages=self.kv_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.kv_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - - s_prod,s_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.s_bar.data_ptr(),num_stages=1,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.epilogue_warp_id))).make_participants() - - # P-ready: softmax -> MMA signal using NamedBarrier - softmax_done_bar = pipeline.NamedBarrier(barrier_id=3, num_threads=32 + 32*len(self.epilogue_warp_id)) - - acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(),num_stages=self.num_acc_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,len(self.epilogue_warp_id)),cta_layout_vmnk=cl_vmnk,defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2,num_threads=32*len((self.mma_warp_id,*self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr,barrier_for_retrieve=tmem_bar,allocator_warp_id=self.epilogue_warp_id[0],is_two_cta=cute.size(qk_mma.thr_id.shape)==2,two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk,is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype,layout=q_smem_s.outer,byte_alignment=128,swizzle=q_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype,layout=k_smem_s.outer,byte_alignment=128,swizzle=k_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype,layout=c_smem_s.outer,byte_alignment=128,swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None)) - gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None)) - gC = cute.local_tile(mC,cute.slice_(self.qk_mma_tiler,(None,None,0)),(None,None,None)) - n_kv_tiles = cute.size(gK, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,0,None,0)).shape) - tAsQ,tAgQ = cpasync.tma_partition(tma_q,0,a_lay,cute.group_modes(sQ,0,3),cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape) - tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - - qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # TMA LOAD - if warp_idx == self.tma_warp_id: - qp.reset(); qh = qp.acquire_and_advance() - cute.copy(tma_q,tAgQ[(None,qh.count)],tAsQ[(None,qh.index)],tma_bar_ptr=qh.barrier) - qp.tail() - kvp.reset(); pk = kvp.try_acquire() - for kt in cutlass.range(n_kv_tiles,unroll=1): - kh = kvp.acquire_and_advance(pk) - cute.copy(tma_k,tBgK[(None,kh.count)],tBsK[(None,kh.index)],tma_bar_ptr=kh.barrier) - pk = cutlass.Boolean(1) - kvp.tail() - - # MMA: QK, then wait for softmax, then signal done - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - qc.reset(); qh = qc.wait_and_advance(); qh.release() - kvc.reset(); pk = kvc.try_wait() - - acc_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_stage) - acc_pipe.producer_acquire(acc_st) - - for kt in range(n_kv_tiles): - kh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - - sh = s_prod.acquire_and_advance() - qk_mma.set(tcgen05.Field.ACCUMULATE, False) - for kb in cutlass.range(cute.size(tCrQ,mode=[2]), unroll_full=True): - cute.gemm(qk_mma, tStS0, tCrQ[(None,None,kb,0)], tCrK[(None,None,kb,kh.index)], tStS0) - qk_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - sh.commit() - kh.release() - - # Wait for softmax to finish - softmax_done_bar.arrive_and_wait() - - # After all tiles: S contains the identity-softmax'd result - # But we want to output P (written by softmax at p0 offset) - # For this test, output what's at tmem_s0_offset (the final S accumulator) - acc_pipe.producer_commit(acc_st); acc_st.advance() - acc_pipe.producer_tail(acc_st) - - # EPILOGUE: identity softmax + output - if warp_idx < self.mma_warp_id: - tmem.allocate(self.num_tmem_alloc_cols) - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) - sfw_idx = tidx % (32 * len(self.epilogue_warp_id)) - - # S load - tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS0) - thr_load = tiled_tmem_load.get_slice(sfw_idx) - tTMEM_LOADtS = thr_load.partition_S(tStS0) - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tTMEM_LOADcS = thr_load.partition_D(tScS) - - # P store - tStS_P_layout = cute.composition(tStS.layout, cute.make_layout((128, self.tilePlikeFP32))) - tStS_P = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStS_P_layout) - tmem_store_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStS_P) - thr_store = tiled_tmem_store.get_slice(sfw_idx) - tTMEM_STOREtS_x4 = thr_store.partition_D(tStS_P) - tScS_P_layout = cute.composition(tScS.layout, cute.make_layout((128, self.tilePlikeFP32))) - tScS_P = cute.make_tensor(tScS.iterator, tScS_P_layout) - tTMEM_STOREcS = thr_store.partition_S(tScS_P) - - for kt in range(n_kv_tiles): - si_handle = s_cons.wait_and_advance() - - # Load S - tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS) - - # Identity softmax: FP32 S -> BF16 P, write to TMEM at p0 offset - tTMEM_STORErS_x4 = cute.make_rmem_tensor(tTMEM_STOREcS.shape, self.qk_acc_dtype) - tTMEM_STORErS_x4_e = cute.make_tensor(cute.recast_ptr(tTMEM_STORErS_x4.iterator, dtype=self.q_dtype), tTMEM_LOADrS.layout) - - frg_cnt = 4; frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt - tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile)) - tTMEM_STORErS_x4_e_frg = cute.logical_divide(tTMEM_STORErS_x4_e, cute.make_layout(frg_tile)) - for j in range(frg_cnt): - s_vec = tTMEM_LOADrS_frg[None, j].load() - tTMEM_STORErS_x4_e_frg[None, j].store(s_vec.to(self.q_dtype)) - cute.copy(tiled_tmem_store, tTMEM_STORErS_x4, tTMEM_STOREtS_x4) - cute.arch.fence_view_async_tmem_store() - - si_handle.release() - # Signal MMA - softmax_done_bar.arrive() - - # Output: read from p0 offset (BF16 P values, but we read as FP32) - # We need to output the BF16 P values. Read from p0 and convert. - # Actually, output the S accumulator (at s0) for now to verify QK works - tCtS_base = cute.make_tensor(tmem_ptr + self.tmem_s0_offset, tCtS_fake.layout) - acc_cons_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage) - c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) - c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp) - acc_cons_st = utils.gemm.sm100.epilogue_tma_store(self, tidx, warp_idx, tma_c, tCtS_base, sC, tCgC, epi_tile, 0, const_expr(lambda x: x), (0,0,0), acc_cons_st, acc_pipe, c_pipe) - c_pipe.producer_tail() - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def test(): - torch.manual_seed(42) - n = 128; m, hd = 128, HEAD_DIM - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(n, hd, 1, dtype=torch.bfloat16, device='cuda') - c = torch.zeros(m, n, 1, dtype=torch.bfloat16, device='cuda') - qf = q[:,:,0].float(); kf = k[:,:,0].float() - ref = (qf @ kf.T) - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - kernel = QkSoftmaxTest() - print('Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mC, stream) - print('Running...', flush=True) - compiled(mQ, mK, mC, stream) - torch.cuda.synchronize() - out = c[:,:,0].float() - cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - print(f'QK+softmax n={n}: cosine {cos:.6f} {"PASS" if cos >= 0.99 else "FAIL"}') - if cos < 0.99: - print(f' out[0,:4]={out[0,:4].tolist()} ref[0,:4]={ref[0,:4].tolist()}') - -if __name__ == '__main__': - test() diff --git a/tests/archive/unit_test_qkonly.py b/tests/archive/unit_test_qkonly.py deleted file mode 100644 index bf5c7a34..00000000 --- a/tests/archive/unit_test_qkonly.py +++ /dev/null @@ -1,269 +0,0 @@ -""" -Debug: QK only (no PV) with KV-tile interleaving pipeline. -Outputs P to GMEM to verify QK+softmax pipeline works. -n=128, single KV tile, identity softmax. -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct - -HEAD_DIM = 64 - - -class QkOnlyTest: - def __init__(self): - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0,1,2,3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.num_c_stage = 2 - self.kv_stage = 2; self.q_stage = 1 - - def _setup(self, qk_mma): - qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (128, 128, qk_ik * 4) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0]//cute.size(qk_mma.thr_id.shape), 128, self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - self.q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, self.q_dtype, self.q_stage) - self.k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, self.q_dtype, self.kv_stage) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - - qk_thr = qk_mma.get_slice(0); qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - - self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width - self.tmem_s0_offset = 0; self.tmem_p0_offset = 32 - self.tmem_o0_offset = 0 # Output is at S offset for QK-only - - tCS = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage)) - self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCS], arch="sm_100") - - cta = cute.size(qk_mma.thr_id.shape) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)) - k_s = cute.slice_(self.k_smem_s,(None,None,None,0)) - self.q_tx_bytes = cute.size_in_bytes(self.q_dtype, q_s) * cta - self.kv_tx_bytes = cute.size_in_bytes(self.q_dtype, k_s) * cta - - @cute.jit - def __call__(self, q, k, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, self.a_major, self.b_major, - self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM) - self._setup(qk_mma) - - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)) - k_s = cute.slice_(self.k_smem_s,(None,None,None,0)) - - tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id), - q,q_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id), - k,k_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - epi_s = cute.select(self.c_smem_s,mode=[0,1]) - tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile) - - self._kernel(qk_mma, tma_q, mQ, tma_k, mK, tma_c, mC, - self.cluster_layout_vmnk, self.q_smem_s, self.k_smem_s, self.c_smem_s, self.epi_tile - ).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, tma_q, mQ, tma_k, mK, tma_c, mC, - cl_vmnk, q_smem_s, k_smem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx,_,_ = cute.arch.thread_idx() - - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) - cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - q_bar: cute.struct.MemRange[cutlass.Int64, self.q_stage*2] - kv_bar: cute.struct.MemRange[cutlass.Int64, self.kv_stage*2] - s_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage*2] - tmem_dealloc: cutlass.Int64; holding: cutlass.Int32 - - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - qp,qc = pipeline.PipelineTmaUmma.create( - barrier_storage=st.q_bar.data_ptr(), num_stages=self.q_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1), - tx_count=self.q_tx_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - kvp,kvc = pipeline.PipelineTmaUmma.create( - barrier_storage=st.kv_bar.data_ptr(), num_stages=self.kv_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1), - tx_count=self.kv_tx_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True - ).make_participants() - - s_prod,s_cons = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.s_bar.data_ptr(), num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32*len(self.epilogue_warp_id)) - ).make_participants() - - softmax_done_bar = pipeline.NamedBarrier(barrier_id=3, num_threads=32 + 32*len(self.epilogue_warp_id)) - - acc_pipe = pipeline.PipelineUmmaAsync.create( - barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, len(self.epilogue_warp_id)), - cta_layout_vmnk=cl_vmnk, defer_sync=True) - - tmem_bar = pipeline.NamedBarrier(barrier_id=2, - num_threads=32*len((self.mma_warp_id,*self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, - allocator_warp_id=self.epilogue_warp_id[0], - is_two_cta=cute.size(qk_mma.thr_id.shape)==2, - two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=q_smem_s.outer, byte_alignment=128, swizzle=q_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype, layout=k_smem_s.outer, byte_alignment=128, swizzle=k_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler,(None,0,None)), (None,None,None)) - gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler,(0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler,(None,None,0)), (None,None,None)) - n_kv_tiles = cute.size(gK, mode=[3]) - - qk_thr = qk_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - tCgC = qk_thr.partition_C(gC) - - a_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,0,None,0)).shape) - tAsQ,tAgQ = cpasync.tma_partition(tma_q,0,a_lay,cute.group_modes(sQ,0,3),cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape) - tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - - qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage)) - - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # TMA LOAD - if warp_idx == self.tma_warp_id: - qp.reset(); qh = qp.acquire_and_advance() - cute.copy(tma_q, tAgQ[(None,qh.count)], tAsQ[(None,qh.index)], tma_bar_ptr=qh.barrier) - qp.tail() - - kvp.reset(); pk = kvp.try_acquire() - for kt in cutlass.range(n_kv_tiles, unroll=1): - kh = kvp.acquire_and_advance(pk) - cute.copy(tma_k, tBgK[(None,kh.count)], tBsK[(None,kh.index)], tma_bar_ptr=kh.barrier) - pk = cutlass.Boolean(1) - kvp.tail() - - # MMA - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - qc.reset(); qh = qc.wait_and_advance(); qh.release() - kvc.reset(); pk = kvc.try_wait() - - acc_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_stage) - acc_pipe.producer_acquire(acc_st) - - for kt in range(n_kv_tiles): - kh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) - - # QK only, accumulate across KV tiles - sh = s_prod.acquire_and_advance() - qk_mma.set(tcgen05.Field.ACCUMULATE, kt != 0) - for kb in cutlass.range(cute.size(tCrQ,mode=[2]), unroll_full=True): - cute.gemm(qk_mma, tStS0, - tCrQ[(None,None,kb,0)], tCrK[(None,None,kb,kh.index)], tStS0) - qk_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - sh.commit() - kh.release() - - # Wait for softmax (identity: just signal done) - softmax_done_bar.arrive_and_wait() - - acc_pipe.producer_commit(acc_st); acc_st.advance() - acc_pipe.producer_tail(acc_st) - - # EPILOGUE - if warp_idx < self.mma_warp_id: - tmem.allocate(self.num_tmem_alloc_cols) - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) - sfw_idx = tidx % (32 * len(self.epilogue_warp_id)) - - for kt in range(n_kv_tiles): - si_handle = s_cons.wait_and_advance() - # Identity softmax: no-op, just signal MMA - si_handle.release() - softmax_done_bar.arrive() - - # Output S (QK result) to GMEM - tCtS_base = cute.make_tensor(tmem_ptr + self.tmem_s0_offset, tCtS_fake.layout) - acc_cons_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage) - c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) - c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp) - acc_cons_st = utils.gemm.sm100.epilogue_tma_store( - self, tidx, warp_idx, tma_c, tCtS_base, sC, tCgC, - epi_tile, 0, const_expr(lambda x: x), (0,0,0), acc_cons_st, acc_pipe, c_pipe) - c_pipe.producer_tail() - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def test(): - torch.manual_seed(42) - n = 128 - m, hd = 128, HEAD_DIM - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(n, hd, 1, dtype=torch.bfloat16, device='cuda') - c = torch.zeros(m, n, 1, dtype=torch.bfloat16, device='cuda') # (128, 128) output = S - - qf = q[:,:,0].float(); kf = k[:,:,0].float() - ref = (qf @ kf.T) - - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - kernel = QkOnlyTest() - print('Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mC, stream) - print('Running...', flush=True) - compiled(mQ, mK, mC, stream) - torch.cuda.synchronize() - out = c[:,:,0].float() - cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - print(f'QK-only n={n}: cosine {cos:.6f} {"PASS" if cos >= 0.99 else "FAIL"}') - if cos < 0.99: - print(f' out[0,:4]={out[0,:4].tolist()} ref[0,:4]={ref[0,:4].tolist()}') - print(f' out stats: min={out.min().item():.4f} max={out.max().item():.4f}') - print(f' ref stats: min={ref.min().item():.4f} max={ref.max().item():.4f}') - -if __name__ == '__main__': - test() diff --git a/tests/fmha_v3_stage_c_example10.py b/tests/fmha_v3_stage_c_example10.py deleted file mode 100644 index 14b12f0e..00000000 --- a/tests/fmha_v3_stage_c_example10.py +++ /dev/null @@ -1,510 +0,0 @@ -""" -FMHA v3 Stage-C Multi-Tile (8-mode TMA indexing, paired-atom epilogue). - -Three structural rules learned the hard way: - -(A) Pipeline handle's `.count` is NOT a GMEM tile coordinate. Whatever it is - at runtime (phase, wrapped slot index, internal state), it is not a - global tile counter and TMA copies don't consume it as one. Use the - loop induction variable for GMEM, handle.index for SMEM. - -(B) Hand-constructed TMEM load/store atoms (Ld32x32bOp + St32x32bOp built - independently) preserve register tile shape across a round-trip only if - they share the same Repetition count. Pair-matching also via - `utils.sm100.get_tmem_load_op` + `get_smem_store_op` works and is what - the CUTLASS Blackwell FMHA reference uses in `correction_rescale`. - -(C) tma_partition produces 4-mode tensors: (((64,128),1), ?, ?, ?). - Mode 2 is the GMEM-tile iteration axis. - Pre-slicing with `tBgK[(None,None,0,0)]` keeps modes 0,1 free but sets - mode 2 to 0, collapsing the KV-tile axis so TMA always reads tile 0. - Fix: pre-slice with `tBgK[(None,0,None,0)]` to keep modes 0,2 free, - then `cute.copy(tma_k, tBgK[None, kt], ...)` indexes the surviving - KV_tiles mode. Verified shapes on B200: after (None,0,None,0), all - three tensors become 2D: (((64,128),1), Int32(?)) or (((64,128),1), 1). - -Kernel structure: - -1. Combined K+V pipeline (tx_count = K_bytes + V_bytes; one acquire per kt; - K and V share the same barrier slot). SMEM slot via kvh.index, GMEM via - the loop's Python int kt (producer is fully unrolled at trace time via - cutlass.range_constexpr, since self.n_kv_tiles is known from __init__). - -2. Reference-style scaled epilogue: TMEM correction_rescale (O *= 1/row_sum - via paired Ld32x32b + St32x32b atoms), then standard epilogue_tma_store - to send O from TMEM through SMEM to GMEM. No TMEM round-trip with - mismatched atoms. - -3. Per-tile O rescale (O *= exp2(old_max - new_max) before PV[kt]) lives in - the softmax warp BEFORE softmax_done_bar.arrive(). Reuses the same - paired-atom pattern as the final normalize. - -4. final_o_bar (32 MMA + 128 softmax threads). MMA arrives between - acc_pipe.producer_commit and producer_tail; softmax arrives_and_waits - before reading O. Order: producer_commit → final_o_bar.arrive() → - producer_tail (reverse deadlocks). -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct -import math - -HEAD_DIM = 64 - - -class FmhaV3StageCMulti: - def __init__(self, s_k=128, scale_softmax=None): - # s_k MUST equal actual sequence length n. - self.s_k = s_k - self.n_kv_tiles = s_k // 128 - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0,1,2,3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.num_c_stage = 2 - self.kv_stage = 2; self.q_stage = 1; self.num_c_stage = 2 - self.scale_softmax = scale_softmax if scale_softmax is not None else 1.0 / math.sqrt(HEAD_DIM) - self.scale_softmax_log2 = self.scale_softmax * math.log2(math.e) - - def _setup(self, qk_mma, pv_mma): - qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (128, 128, qk_ik * 4) - pv_ik = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (128, HEAD_DIM, pv_ik * (128 // pv_ik)) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0]//cute.size(qk_mma.thr_id.shape), HEAD_DIM, self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - self.q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, self.q_dtype, self.q_stage) - self.k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, self.q_dtype, self.kv_stage) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, self.kv_stage) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - qk_thr = qk_mma.get_slice(0); qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - pv_thr = pv_mma.get_slice(0); pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - self.tmem_s0_offset = 0; self.tmem_p0_offset = 32 - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - p_end = self.tmem_p0_offset + p_cols_fp32 - s_cols = self.qk_mma_tiler[1] - o_after = max(s_cols, p_end) - self.tmem_o0_offset = ((o_after + 31) // 32) * 32 - o_cols = find_tmem_tensor_col_offset(tOtO) - total = self.tmem_o0_offset + o_cols - self.num_tmem_alloc_cols = 1 - while self.num_tmem_alloc_cols < total: - self.num_tmem_alloc_cols *= 2 - cta = cute.size(qk_mma.thr_id.shape) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)) - k_s = cute.slice_(self.k_smem_s,(None,None,None,0)) - v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - self.q_tx_bytes = cute.size_in_bytes(self.q_dtype, q_s) * cta - # Combined barrier: tx_count covers BOTH K and V transfers per acquire. - self.kv_tx_bytes = (cute.size_in_bytes(self.q_dtype, k_s) + - cute.size_in_bytes(self.q_dtype, v_s)) * cta - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - v_fmha = cute.make_tensor( - v.iterator, - cute.make_layout( - (HEAD_DIM, self.s_k, 1), - stride=(1, HEAD_DIM, HEAD_DIM * self.s_k), - ), - ) - self.v_major = LayoutEnum.from_tensor(v_fmha).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, self.a_major, self.b_major, self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, self.qk_acc_dtype, self.cta_group, (128,HEAD_DIM), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)); v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id),q,q_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id),k,k_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v_fmha,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape) - epi_s = cute.select(self.c_smem_s,mode=[0,1]) - tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile) - self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.v_smem_s,self.p_tmem_s,self.c_smem_s,self.epi_tile).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, q_smem_s, k_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx,_,_ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k); cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - q_bar: cute.struct.MemRange[cutlass.Int64, self.q_stage*2] - kv_bar: cute.struct.MemRange[cutlass.Int64, self.kv_stage*2] - s_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage*2] - tmem_dealloc: cutlass.Int64; holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - qp,qc = pipeline.PipelineTmaUmma.create(barrier_storage=st.q_bar.data_ptr(),num_stages=self.q_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.q_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - kvp,kvc = pipeline.PipelineTmaUmma.create(barrier_storage=st.kv_bar.data_ptr(),num_stages=self.kv_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.kv_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - s_prod,s_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.s_bar.data_ptr(),num_stages=1,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.epilogue_warp_id))).make_participants() - softmax_done_bar = pipeline.NamedBarrier(barrier_id=3, num_threads=32 + 32*len(self.epilogue_warp_id)) - final_o_bar = pipeline.NamedBarrier(barrier_id=4, num_threads=32 + 32*len(self.epilogue_warp_id)) - acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(),num_stages=self.num_acc_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,len(self.epilogue_warp_id)),cta_layout_vmnk=cl_vmnk,defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2,num_threads=32*len((self.mma_warp_id,*self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr,barrier_for_retrieve=tmem_bar,allocator_warp_id=self.epilogue_warp_id[0],is_two_cta=cute.size(qk_mma.thr_id.shape)==2,two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk,is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype,layout=q_smem_s.outer,byte_alignment=128,swizzle=q_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype,layout=k_smem_s.outer,byte_alignment=128,swizzle=k_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype,layout=v_smem_s.outer,byte_alignment=128,swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype,layout=c_smem_s.outer,byte_alignment=128,swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None)) - gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None)) - gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None)) - gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None)) - n_kv_tiles = cute.size(gK, mode=[3]) - - qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - tCgV = pv_thr.partition_B(gV); tCgC = pv_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,0,None,0)).shape) - tAsQ,tAgQ = cpasync.tma_partition(tma_q,0,a_lay,cute.group_modes(sQ,0,3),cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape) - tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3)) - tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3)) - # SHAPES (from diag test on B200, n=256): - # tAgQ: (((64,128),1), Int32(?), Int32(?), Int32(?)) — 4 modes - # tBgK: (((64,128),1), Int32(?), Int32(?), Int32(?)) — 4 modes - # tVgV: (((64,128),1), 1, 1, 1) — 4 modes - # Mode 2 is the GMEM tile iteration axis. - # (None,0,None,0) keeps modes 0 and 2 free → 2D tensor. - # Then [None, kt] indexes the surviving KV_tiles dim. - # The old (None,None,0,0) kept modes 0,1 free → mode 2 (KV tiles) set to 0. - tAgQ = tAgQ[(None,0,None,0)] - tBgK = tBgK[(None,0,None,0)] - tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None,None,None,0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ===== TMA LOAD warp — fully unrolled ===== - # The old pre-slice (None,None,0,0) kept modes 0,1 free and set - # mode 2 (KV tiles) to 0. Fix: (None,0,None,0) keeps modes 0,2 free. - # Then [None, kt] indexes the KV tile axis correctly. - # Matches the CUTLASS reference FMHA TMA indexing pattern. - if warp_idx == self.tma_warp_id: - qp.reset(); qh = qp.acquire_and_advance() - cute.copy(tma_q, tAgQ[None, Int32(0)], tAsQ[(None, qh.index)], tma_bar_ptr=qh.barrier) - qp.tail() - kvp.reset() - for kt in cutlass.range_constexpr(self.n_kv_tiles): - kvh = kvp.acquire_and_advance() - # After (None,0,None,0) pre-slice, tBgK is 2D: (V_grouped, KV_tiles). - # [None, kt] indexes mode 1 (KV_tiles) with kt. - cute.copy( - tma_k, - tBgK[None, kt], - tBsK[(None, kvh.index)], - tma_bar_ptr=kvh.barrier, - ) - cute.copy( - tma_v, - tVgV[None, kt], - tVsV[(None, kvh.index)], - tma_bar_ptr=kvh.barrier, - ) - kvp.tail() - - # ===== MMA warp ===== - # Outer kt loop unrolled to match the producer. The earlier hypothesis - # was that CuTeDSL 4.5.1 couldn't propagate dynamic TMA coords; the - # actual root cause turned out to be the producer's GMEM-tensor - # pre-slice eating the mode-4 KV-tile axis. The unroll is kept here - # for symmetry with the producer (single concern: one fewer thing - # that could surprise us if the layout assumption shifts), but the - # GMEM indexing fix is what actually makes multi-tile work. - # Inner GEMM K-block loops stay dynamic. kvh.index correctly tracks - # the SMEM ring buffer at runtime. - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - qc.reset(); qh = qc.wait_and_advance(); qh.release() - kvc.reset() - acc_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_stage) - acc_pipe.producer_acquire(acc_st) - for kt in cutlass.range_constexpr(self.n_kv_tiles): - kvh = kvc.wait_and_advance() - sh = s_prod.acquire_and_advance() - qk_mma.set(tcgen05.Field.ACCUMULATE, False) - for kb in cutlass.range(cute.size(tCrQ, mode=[2]), unroll_full=True): - cute.gemm(qk_mma, tStS0, tCrQ[(None,None,kb,0)], tCrK[(None,None,kb,kvh.index)], tStS0) - qk_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - sh.commit() - softmax_done_bar.arrive_and_wait() - pv_mma.set(tcgen05.Field.ACCUMULATE, kt != 0) - for kb in cutlass.range(cute.size(tOrP0, mode=[2]), unroll_full=True): - cute.gemm(pv_mma, tOtO0, tOrP0[(None,None,kb)], tCrV[(None,None,kb,kvh.index)], tOtO0) - pv_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - kvh.release() - acc_pipe.producer_commit(acc_st); acc_st.advance() - final_o_bar.arrive() - acc_pipe.producer_tail(acc_st) - - # ===== SOFTMAX + EPILOGUE warps ===== - if warp_idx < self.mma_warp_id: - tmem.allocate(self.num_tmem_alloc_cols) - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) - sfw_idx = tidx % (32 * len(self.epilogue_warp_id)) - - # S load - tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS0) - thr_load = tiled_tmem_load.get_slice(sfw_idx) - tTMEM_LOADtS = thr_load.partition_S(tStS0) - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tTMEM_LOADcS = thr_load.partition_D(tScS) - - # P store - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - tStP_layout = cute.composition(tStS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tStP0 = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStP_layout) - tmem_store_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStP0) - thr_store = tiled_tmem_store.get_slice(sfw_idx) - tTMEM_STOREtP = thr_store.partition_D(tStP0) - tScP_layout = cute.composition(tScS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tScP = cute.make_tensor(tScS.iterator, tScP_layout) - tTMEM_STOREcP = thr_store.partition_S(tScP) - - # === O rescale path setup (used per-tile AND for final normalize) === - corr_tile_size = 16 - cO = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1])) - tOcO = pv_thr.partition_C(cO) - tOtO_i_layout = cute.composition(tOtO0.layout, cute.make_layout((128, corr_tile_size))) - tOcO_i_layout = cute.composition(tOcO.layout, cute.make_layout((128, corr_tile_size))) - tOtO_i = cute.make_tensor(tOtO0.iterator, tOtO_i_layout) - tOcO_i = cute.make_tensor(tOcO.iterator, tOcO_i_layout) - tmem_load_o_atom = cute.make_copy_atom( - tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), - self.acc_dtype, - ) - tmem_store_o_atom = cute.make_copy_atom( - tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), - self.acc_dtype, - ) - tiled_tmem_load_o = tcgen05.make_tmem_copy(tmem_load_o_atom, tOtO_i) - tiled_tmem_store_o = tcgen05.make_tmem_copy(tmem_store_o_atom, tOtO_i) - thr_tmem_load_o = tiled_tmem_load_o.get_slice(sfw_idx) - thr_tmem_store_o = tiled_tmem_store_o.get_slice(sfw_idx) - tTMEM_LOADtO = thr_tmem_load_o.partition_S(tOtO_i) - tTMEM_LOADcO = thr_tmem_load_o.partition_D(tOcO_i) - tTMEM_STOREtO = thr_tmem_store_o.partition_D(tOtO_i) - n_corr_tiles = HEAD_DIM // corr_tile_size - - row_max = -Float32.inf - row_sum = Float32(0.0) - scale_log2 = Float32(self.scale_softmax_log2) - - # Per-tile softmax loop with online rescale. - # Unrolled for consistency with producer/MMA warps. The `if kt > 0` - # rescale guard now becomes a Python-level conditional at trace - # time (no rescale block emitted for kt=0; rescale block emitted - # in-line for kt=1..N-1). - for kt in cutlass.range_constexpr(self.n_kv_tiles): - si_handle = s_cons.wait_and_advance() - - # Load S[kt] - tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS) - cute.arch.fence_view_async_tmem_load() - - # Pass 1: update row_max in log2-domain. - old_row_max = row_max - frg_cnt = 4 - frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt - tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile)) - for j in range(frg_cnt): - for k in range(cute.size(tTMEM_LOADrS_frg, mode=[0])): - row_max = cute.arch.fmax(row_max, tTMEM_LOADrS_frg[k, j] * scale_log2) - - row_max_safe = row_max - if row_max == -cutlass.Float32.inf: - row_max_safe = Float32(0.0) - - # acc_scale = exp2(old_max - new_max). On first tile this is 0 - # (old_max = -inf), so row_sum stays 0 and rescale is skipped. - # row_max is already in scaled domain, so no extra scale_log2. - acc_scale_ = old_row_max - row_max_safe - acc_scale = cute.math.exp2(acc_scale_, fastmath=True) - if old_row_max == -cutlass.Float32.inf: - acc_scale = Float32(0.0) - row_sum *= acc_scale - - # Pass 2: P = exp2((S - new_max) * log2), accumulate row_sum, - # cast to BF16 via FP32-backed register bridge. - rP_words = cute.make_rmem_tensor(tTMEM_STOREcP.shape, self.qk_acc_dtype) - rP_bf16 = cute.make_tensor(cute.recast_ptr(rP_words.iterator, dtype=self.q_dtype), tTMEM_LOADrS.layout) - minus_row_max = Float32(0.0) - row_max_safe - - rP_bf16_frg = cute.logical_divide(rP_bf16, cute.make_layout(frg_tile)) - for j in range(frg_cnt): - for k in range(cute.size(tTMEM_LOADrS_frg, mode=[0])): - tTMEM_LOADrS_frg[k, j] = tTMEM_LOADrS_frg[k, j] * scale_log2 + minus_row_max - tTMEM_LOADrS_frg[k, j] = cute.math.exp2(tTMEM_LOADrS_frg[k, j], fastmath=True) - row_sum = row_sum + tTMEM_LOADrS_frg[k, j] - s_vec = tTMEM_LOADrS_frg[None, j].load() - rP_bf16_frg[None, j].store(s_vec.to(self.q_dtype)) - - cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP) - cute.arch.fence_view_async_tmem_store() - - # === Per-tile O rescale: O *= acc_scale for kt > 0 === - # Uses the SAME paired-atom pattern as the final normalize. - # Must run BEFORE softmax_done_bar.arrive() so MMA's PV[kt] - # reads the rescaled O. - # Visibility of MMA's PV[kt-1] writes: provided by - # s_cons.wait_and_advance at the top of this iteration, which - # acquires on MMA's S[kt] commit. S[kt] is sequenced after - # PV[kt-1] in MMA's iteration, so PV[kt-1]'s tmem_store_fence - # has been observed by the time we read O here. - if kt > 0: - for i in range(n_corr_tiles): - tTMEM_LOADtO_i = cute.make_tensor( - tTMEM_LOADtO.iterator + i * corr_tile_size, - tTMEM_LOADtO.layout, - ) - tTMEM_STOREtO_i = cute.make_tensor( - tTMEM_STOREtO.iterator + i * corr_tile_size, - tTMEM_STOREtO.layout, - ) - tTMrO = cute.make_rmem_tensor(tTMEM_LOADcO.shape, self.acc_dtype) - cute.copy(tiled_tmem_load_o, tTMEM_LOADtO_i, tTMrO) - cute.arch.fence_view_async_tmem_load() - for k in cutlass.range(cute.size(tTMrO), vectorize=True): - tTMrO[k] = tTMrO[k] * acc_scale - cute.copy(tiled_tmem_store_o, tTMrO, tTMEM_STOREtO_i) - cute.arch.fence_view_async_tmem_store() - - si_handle.release() - softmax_done_bar.arrive() - - # Wait for MMA's PV[N-1] to commit before reading O for normalize. - final_o_bar.arrive_and_wait() - - # === Final O normalization: O *= 1/row_sum === - inv_row_sum = Float32(1.0) / row_sum - for i in range(n_corr_tiles): - tTMEM_LOADtO_i = cute.make_tensor( - tTMEM_LOADtO.iterator + i * corr_tile_size, - tTMEM_LOADtO.layout, - ) - tTMEM_STOREtO_i = cute.make_tensor( - tTMEM_STOREtO.iterator + i * corr_tile_size, - tTMEM_STOREtO.layout, - ) - tTMrO = cute.make_rmem_tensor(tTMEM_LOADcO.shape, self.acc_dtype) - cute.copy(tiled_tmem_load_o, tTMEM_LOADtO_i, tTMrO) - cute.arch.fence_view_async_tmem_load() - for k in cutlass.range(cute.size(tTMrO), vectorize=True): - tTMrO[k] = tTMrO[k] * inv_row_sum - cute.copy(tiled_tmem_store_o, tTMrO, tTMEM_STOREtO_i) - cute.arch.fence_view_async_tmem_store() - - # Standard epilogue: TMEM → SMEM → GMEM via TMA store. - # O in TMEM is now scaled by 1/row_sum. - tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout) - acc_cons_st = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Consumer, self.num_acc_stage - ) - c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) - c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp) - acc_cons_st = utils.gemm.sm100.epilogue_tma_store( - self, tidx, warp_idx, tma_c, tCtO_base, sC, tCgC, epi_tile, - 0, const_expr(lambda x: x), (0, 0, 0), - acc_cons_st, acc_pipe, c_pipe, - ) - c_pipe.producer_tail() - - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def test(): - torch.manual_seed(42) - for n in [128, 256, 512, 1024]: - torch.manual_seed(42) - m, hd = 128, HEAD_DIM - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(n, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(n, hd, dtype=torch.bfloat16, device='cuda') - v_kernel = v.unsqueeze(-1) - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - - qf = q[:, :, 0].float() - kf = k[:, :, 0].float() - scale = 1.0 / math.sqrt(hd) - attn = qf @ kf.T * scale - attn = torch.softmax(attn, dim=-1) - ref = attn @ v.float() - - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - kernel = FmhaV3StageCMulti(s_k=n) - print(f'n={n}: Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - print(f'n={n}: tmem s0={kernel.tmem_s0_offset} p0={kernel.tmem_p0_offset} ' - f'o0={kernel.tmem_o0_offset} alloc={kernel.num_tmem_alloc_cols} ' - f'kv_tx_bytes={kernel.kv_tx_bytes}', flush=True) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - - out = c[:, :, 0].float() - cos = torch.nn.functional.cosine_similarity( - out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0) - ).item() - max_abs = (out - ref).abs().max().item() - n_tiles = n // 128 - print(f'FMHA Stage-C Multi n={n} ({n_tiles} kv tiles): ' - f'cos {cos:.6f} max_abs {max_abs:.4f} ' - f'{"PASS" if cos >= 0.99 else "FAIL"}') - if cos < 0.99: - print(f' out[0,:4]={out[0,:4].tolist()}') - print(f' ref[0,:4]={ref[0,:4].tolist()}') - - -if __name__ == '__main__': - test() \ No newline at end of file diff --git a/tests/fmha_v3_stage_c_example9.py b/tests/fmha_v3_stage_c_example9.py deleted file mode 100644 index e025adfc..00000000 --- a/tests/fmha_v3_stage_c_example9.py +++ /dev/null @@ -1,495 +0,0 @@ -""" -FMHA v3 Stage-C Multi-Tile (paired TMEM/SMEM atoms, reference-style epilogue). - -Two structural rules we had to learn the hard way: - -(A) Pipeline handle's `.count` is NOT a GMEM tile coordinate. Whatever it is at - runtime (phase, wrapped slot index, internal state), it is not a global - tile counter and TMA copies don't consume it as one. Use the loop - induction variable for GMEM, handle.index for SMEM. - -(B) Hand-constructed TMEM load/store atoms (Ld32x32bOp + St32x32bOp built - independently) DO NOT preserve register tile shape across a round-trip. - Use paired atoms (or, as we discovered: independently constructed atoms - DO work if they're built from the SAME `Repetition(N)` count — the - Ld32x32bOp(Rep(16)) + St32x32bOp(Rep(16)) pair preserves the register - tile shape exactly because the atom width matches). This is what the - CUTLASS Blackwell FMHA reference does in `correction_rescale`. - -(C) Multi-tile GMEM indexing: `kt` from cutlass.range constant-folds at trace - time, so all TMA loads address tile 0. Workaround: track an Int32 - coordinate manually, BUT seed it from an SSA expression - (`n_kv_tiles - n_kv_tiles`) rather than a literal `Int32(0)`, so the JIT - sees it as a runtime register and propagates the `+= 1` as a tracked - loop-carried iter_args update. - -Kernel structure: - -1. Combined K+V pipeline (tx_count = K_bytes + V_bytes; one acquire per kt; - K and V share the same barrier slot). SMEM slot via kvh.index, GMEM via - manually-tracked kv_coord (SSA-seeded). - -2. Reference-style scaled epilogue: TMEM correction_rescale (O *= 1/row_sum - via paired Ld32x32b + St32x32b atoms), then standard epilogue_tma_store - to send O from TMEM through SMEM to GMEM. - -3. Per-tile O rescale (multiplying existing O by exp2(old_max - new_max) - before PV[kt]) lives in the softmax warp BEFORE softmax_done_bar.arrive(). - Reuses the same paired-atom pattern as the final normalize. - -4. final_o_bar (32 MMA + 128 softmax threads). MMA arrives between - acc_pipe.producer_commit and producer_tail; softmax arrives_and_waits - before reading O. Order: producer_commit → final_o_bar.arrive() → - producer_tail (reverse deadlocks). -""" -import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, Boolean, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda -import cutlass.torch as ct -import math - -HEAD_DIM = 64 - - -class FmhaV3StageCMulti: - def __init__(self, s_k=128, scale_softmax=None): - # s_k MUST equal actual sequence length n. - self.s_k = s_k - self.n_kv_tiles = s_k // 128 - self.acc_dtype = Float32; self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 - self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE - self.epilogue_warp_id = (0,1,2,3); self.mma_warp_id = 4; self.tma_warp_id = 5 - self.threads_per_cta = 192; self.num_c_stage = 2 - self.kv_stage = 2; self.q_stage = 1; self.num_c_stage = 2 - self.scale_softmax = scale_softmax if scale_softmax is not None else 1.0 / math.sqrt(HEAD_DIM) - self.scale_softmax_log2 = self.scale_softmax * math.log2(math.e) - - def _setup(self, qk_mma, pv_mma): - qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) - self.qk_mma_tiler = (128, 128, qk_ik * 4) - pv_ik = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (128, HEAD_DIM, pv_ik * (128 // pv_ik)) - self.mma_tiler = self.qk_mma_tiler - self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) - self.cta_tile_shape_mnk = (self.qk_mma_tiler[0]//cute.size(qk_mma.thr_id.shape), HEAD_DIM, self.qk_mma_tiler[2]) - self.c_layout = LayoutEnum.ROW_MAJOR - self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) - self.num_ab_stage = 1; self.num_acc_stage = 1 - self.q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, self.q_dtype, self.q_stage) - self.k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, self.q_dtype, self.kv_stage) - self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, self.kv_stage) - self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) - self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - qk_thr = qk_mma.get_slice(0); qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - pv_thr = pv_mma.get_slice(0); pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - self.tmem_s0_offset = 0; self.tmem_p0_offset = 32 - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - p_end = self.tmem_p0_offset + p_cols_fp32 - s_cols = self.qk_mma_tiler[1] - o_after = max(s_cols, p_end) - self.tmem_o0_offset = ((o_after + 31) // 32) * 32 - o_cols = find_tmem_tensor_col_offset(tOtO) - total = self.tmem_o0_offset + o_cols - self.num_tmem_alloc_cols = 1 - while self.num_tmem_alloc_cols < total: - self.num_tmem_alloc_cols *= 2 - cta = cute.size(qk_mma.thr_id.shape) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)) - k_s = cute.slice_(self.k_smem_s,(None,None,None,0)) - v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - self.q_tx_bytes = cute.size_in_bytes(self.q_dtype, q_s) * cta - # Combined barrier: tx_count covers BOTH K and V transfers per acquire. - self.kv_tx_bytes = (cute.size_in_bytes(self.q_dtype, k_s) + - cute.size_in_bytes(self.q_dtype, v_s)) * cta - - @cute.jit - def __call__(self, q, k, v, c, stream): - self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - v_fmha = cute.make_tensor( - v.iterator, - cute.make_layout( - (HEAD_DIM, self.s_k, 1), - stride=(1, HEAD_DIM, HEAD_DIM * self.s_k), - ), - ) - self.v_major = LayoutEnum.from_tensor(v_fmha).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) - qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, self.a_major, self.b_major, self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM) - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, self.qk_acc_dtype, self.cta_group, (128,HEAD_DIM), tcgen05.OperandSource.TMEM) - self._setup(qk_mma, pv_mma) - q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)); v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id),q,q_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id),k,k_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape) - tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v_fmha,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape) - epi_s = cute.select(self.c_smem_s,mode=[0,1]) - tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile) - self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.v_smem_s,self.p_tmem_s,self.c_smem_s,self.epi_tile).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream) - - @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, q_smem_s, k_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tidx,_,_ = cute.arch.thread_idx() - if warp_idx == self.tma_warp_id: - cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k); cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) - - @cute.struct - class SS: - q_bar: cute.struct.MemRange[cutlass.Int64, self.q_stage*2] - kv_bar: cute.struct.MemRange[cutlass.Int64, self.kv_stage*2] - s_bar: cute.struct.MemRange[cutlass.Int64, 2] - acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage*2] - tmem_dealloc: cutlass.Int64; holding: cutlass.Int32 - smem = utils.SmemAllocator(); st = smem.allocate(SS) - - qp,qc = pipeline.PipelineTmaUmma.create(barrier_storage=st.q_bar.data_ptr(),num_stages=self.q_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.q_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - kvp,kvc = pipeline.PipelineTmaUmma.create(barrier_storage=st.kv_bar.data_ptr(),num_stages=self.kv_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.kv_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants() - s_prod,s_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.s_bar.data_ptr(),num_stages=1,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.epilogue_warp_id))).make_participants() - softmax_done_bar = pipeline.NamedBarrier(barrier_id=3, num_threads=32 + 32*len(self.epilogue_warp_id)) - final_o_bar = pipeline.NamedBarrier(barrier_id=4, num_threads=32 + 32*len(self.epilogue_warp_id)) - acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(),num_stages=self.num_acc_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,len(self.epilogue_warp_id)),cta_layout_vmnk=cl_vmnk,defer_sync=True) - tmem_bar = pipeline.NamedBarrier(barrier_id=2,num_threads=32*len((self.mma_warp_id,*self.epilogue_warp_id))) - tmem = utils.TmemAllocator(st.holding.ptr,barrier_for_retrieve=tmem_bar,allocator_warp_id=self.epilogue_warp_id[0],is_two_cta=cute.size(qk_mma.thr_id.shape)==2,two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) - pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk,is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype,layout=q_smem_s.outer,byte_alignment=128,swizzle=q_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype,layout=k_smem_s.outer,byte_alignment=128,swizzle=k_smem_s.inner) - sV = smem.allocate_tensor(element_type=self.q_dtype,layout=v_smem_s.outer,byte_alignment=128,swizzle=v_smem_s.inner) - sC = smem.allocate_tensor(element_type=self.o_dtype,layout=c_smem_s.outer,byte_alignment=128,swizzle=c_smem_s.inner) - - gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None)) - gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None)) - gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None)) - gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None)) - n_kv_tiles = cute.size(gK, mode=[3]) - - qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - tCgV = pv_thr.partition_B(gV); tCgC = pv_thr.partition_C(gC) - a_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,0,None,0)).shape) - tAsQ,tAgQ = cpasync.tma_partition(tma_q,0,a_lay,cute.group_modes(sQ,0,3),cute.group_modes(tCgQ,0,3)) - b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape) - tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3)) - tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3)) - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)]; tVgV = tVgV[(None,0,None,0)] - - tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) - tCrV = pv_mma.make_fragment_B(sV) - - qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) - pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) - - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP) - tOrP = tOrP_base[(None,None,None,0)] - tOrP0 = cute.make_tensor( - tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, - tOrP.layout) - - tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage)) - tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage)) - pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) - - # ===== TMA LOAD warp ===== - # Multi-tile GMEM indexing — combination that finally makes it work: - # - # 1. SSA-seeded kv_coord (n_kv_tiles - n_kv_tiles, not Int32(0)). The - # JIT folds literal Int32(0) at trace time; subtracting an SSA - # value from itself produces an SSA register zero that's tracked. - # - # 2. Drop the try_acquire/pk pattern. We had `pk = kvp.try_acquire()` - # outside the loop and `pk = cutlass.Boolean(1)` at the end of - # each iter — a second loop-carried variable. The JIT's automatic - # iter_args yielding *can* handle one or both, but mixing two - # mutated variables exposed an edge case where `pk` won the - # iter_args slot and `kv_coord` got constant-folded. - # - # The reference uses bare `acquire_and_advance()` with no `pk`, - # so just match that. One loop-carried variable (kv_coord), - # pipeline state managed internally. - if warp_idx == self.tma_warp_id: - qp.reset(); qh = qp.acquire_and_advance() - cute.copy(tma_q, tAgQ[(None, Int32(0))], tAsQ[(None, qh.index)], tma_bar_ptr=qh.barrier) - qp.tail() - kvp.reset() - kv_coord = n_kv_tiles - n_kv_tiles # SSA runtime zero - for kt in cutlass.range(0, n_kv_tiles, 1, unroll=1): - kvh = kvp.acquire_and_advance() - cute.copy(tma_k, tBgK[(None, kv_coord)], tBsK[(None, kvh.index)], tma_bar_ptr=kvh.barrier) - cute.copy(tma_v, tVgV[(None, kv_coord)], tVsV[(None, kvh.index)], tma_bar_ptr=kvh.barrier) - kv_coord = kv_coord + 1 - kvp.tail() - - # ===== MMA warp ===== - # One wait per kt; same slot index used for both K (QK) and V (PV). - # Release happens AFTER PV — combined slot stays held across QK+PV. - # Note: dropped the try_wait/pk pattern here too, matching the TMA - # warp's simplification. Bare wait_and_advance, no loop-carried pk. - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - qc.reset(); qh = qc.wait_and_advance(); qh.release() - kvc.reset() - acc_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_stage) - acc_pipe.producer_acquire(acc_st) - for kt in range(self.n_kv_tiles): - kvh = kvc.wait_and_advance() - sh = s_prod.acquire_and_advance() - qk_mma.set(tcgen05.Field.ACCUMULATE, False) - for kb in cutlass.range(cute.size(tCrQ, mode=[2]), unroll_full=True): - cute.gemm(qk_mma, tStS0, tCrQ[(None,None,kb,0)], tCrK[(None,None,kb,kvh.index)], tStS0) - qk_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - sh.commit() - softmax_done_bar.arrive_and_wait() - pv_mma.set(tcgen05.Field.ACCUMULATE, kt != 0) - for kb in cutlass.range(cute.size(tOrP0, mode=[2]), unroll_full=True): - cute.gemm(pv_mma, tOtO0, tOrP0[(None,None,kb)], tCrV[(None,None,kb,kvh.index)], tOtO0) - pv_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - kvh.release() - acc_pipe.producer_commit(acc_st); acc_st.advance() - final_o_bar.arrive() - acc_pipe.producer_tail(acc_st) - - # ===== SOFTMAX + EPILOGUE warps ===== - if warp_idx < self.mma_warp_id: - tmem.allocate(self.num_tmem_alloc_cols) - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) - sfw_idx = tidx % (32 * len(self.epilogue_warp_id)) - - # S load - tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS0) - thr_load = tiled_tmem_load.get_slice(sfw_idx) - tTMEM_LOADtS = thr_load.partition_S(tStS0) - cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) - tScS = qk_thr.partition_C(cS) - tTMEM_LOADcS = thr_load.partition_D(tScS) - - # P store - p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width - tStP_layout = cute.composition(tStS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tStP0 = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStP_layout) - tmem_store_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype) - tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStP0) - thr_store = tiled_tmem_store.get_slice(sfw_idx) - tTMEM_STOREtP = thr_store.partition_D(tStP0) - tScP_layout = cute.composition(tScS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32))) - tScP = cute.make_tensor(tScS.iterator, tScP_layout) - tTMEM_STOREcP = thr_store.partition_S(tScP) - - # === O rescale path setup (used per-tile AND for final normalize) === - corr_tile_size = 16 - cO = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1])) - tOcO = pv_thr.partition_C(cO) - tOtO_i_layout = cute.composition(tOtO0.layout, cute.make_layout((128, corr_tile_size))) - tOcO_i_layout = cute.composition(tOcO.layout, cute.make_layout((128, corr_tile_size))) - tOtO_i = cute.make_tensor(tOtO0.iterator, tOtO_i_layout) - tOcO_i = cute.make_tensor(tOcO.iterator, tOcO_i_layout) - tmem_load_o_atom = cute.make_copy_atom( - tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), - self.acc_dtype, - ) - tmem_store_o_atom = cute.make_copy_atom( - tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), - self.acc_dtype, - ) - tiled_tmem_load_o = tcgen05.make_tmem_copy(tmem_load_o_atom, tOtO_i) - tiled_tmem_store_o = tcgen05.make_tmem_copy(tmem_store_o_atom, tOtO_i) - thr_tmem_load_o = tiled_tmem_load_o.get_slice(sfw_idx) - thr_tmem_store_o = tiled_tmem_store_o.get_slice(sfw_idx) - tTMEM_LOADtO = thr_tmem_load_o.partition_S(tOtO_i) - tTMEM_LOADcO = thr_tmem_load_o.partition_D(tOcO_i) - tTMEM_STOREtO = thr_tmem_store_o.partition_D(tOtO_i) - n_corr_tiles = HEAD_DIM // corr_tile_size - - row_max = -Float32.inf - row_sum = Float32(0.0) - scale_log2 = Float32(self.scale_softmax_log2) - - # Per-tile softmax loop with online rescale. - # Use self.n_kv_tiles (Python int) not the CuTeDSL symbolic n_kv_tiles - # — Python range() needs a Python int. - for kt in range(self.n_kv_tiles): - si_handle = s_cons.wait_and_advance() - cute.printf("SOFTMAX kt=%d row_max_before=%f row_sum_before=%f\n", kt, row_max, row_sum) - - # Load S[kt] - tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype) - cute.copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS) - cute.arch.fence_view_async_tmem_load() - - # Pass 1: update row_max in log2-domain. - old_row_max = row_max - frg_cnt = 4 - frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt - tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile)) - for j in range(frg_cnt): - for k in range(cute.size(tTMEM_LOADrS_frg, mode=[0])): - row_max = cute.arch.fmax(row_max, tTMEM_LOADrS_frg[k, j] * scale_log2) - - row_max_safe = row_max - if row_max == -cutlass.Float32.inf: - row_max_safe = Float32(0.0) - - # acc_scale = exp2(old_max - new_max). On first tile this is 0 - # (old_max = -inf), so row_sum stays 0 and rescale is skipped. - # row_max is already in scaled domain, so no extra scale_log2. - acc_scale_ = old_row_max - row_max_safe - acc_scale = cute.math.exp2(acc_scale_, fastmath=True) - if old_row_max == -cutlass.Float32.inf: - acc_scale = Float32(0.0) - row_sum *= acc_scale - - # Pass 2: P = exp2((S - new_max) * log2), accumulate row_sum, - # cast to BF16 via FP32-backed register bridge. - rP_words = cute.make_rmem_tensor(tTMEM_STOREcP.shape, self.qk_acc_dtype) - rP_bf16 = cute.make_tensor(cute.recast_ptr(rP_words.iterator, dtype=self.q_dtype), tTMEM_LOADrS.layout) - minus_row_max = Float32(0.0) - row_max_safe - - rP_bf16_frg = cute.logical_divide(rP_bf16, cute.make_layout(frg_tile)) - for j in range(frg_cnt): - for k in range(cute.size(tTMEM_LOADrS_frg, mode=[0])): - tTMEM_LOADrS_frg[k, j] = tTMEM_LOADrS_frg[k, j] * scale_log2 + minus_row_max - tTMEM_LOADrS_frg[k, j] = cute.math.exp2(tTMEM_LOADrS_frg[k, j], fastmath=True) - row_sum = row_sum + tTMEM_LOADrS_frg[k, j] - s_vec = tTMEM_LOADrS_frg[None, j].load() - rP_bf16_frg[None, j].store(s_vec.to(self.q_dtype)) - - cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP) - cute.arch.fence_view_async_tmem_store() - - # === Per-tile O rescale: O *= acc_scale for kt > 0 === - # Uses the SAME paired-atom pattern as the final normalize. - # Must run BEFORE softmax_done_bar.arrive() so MMA's PV[kt] - # reads the rescaled O. - # Visibility of MMA's PV[kt-1] writes: provided by - # s_cons.wait_and_advance at the top of this iteration, which - # acquires on MMA's S[kt] commit. S[kt] is sequenced after - # PV[kt-1] in MMA's iteration, so PV[kt-1]'s tmem_store_fence - # has been observed by the time we read O here. - if kt > 0: - for i in range(n_corr_tiles): - tTMEM_LOADtO_i = cute.make_tensor( - tTMEM_LOADtO.iterator + i * corr_tile_size, - tTMEM_LOADtO.layout, - ) - tTMEM_STOREtO_i = cute.make_tensor( - tTMEM_STOREtO.iterator + i * corr_tile_size, - tTMEM_STOREtO.layout, - ) - tTMrO = cute.make_rmem_tensor(tTMEM_LOADcO.shape, self.acc_dtype) - cute.copy(tiled_tmem_load_o, tTMEM_LOADtO_i, tTMrO) - cute.arch.fence_view_async_tmem_load() - for k in cutlass.range(cute.size(tTMrO), vectorize=True): - tTMrO[k] = tTMrO[k] * acc_scale - cute.copy(tiled_tmem_store_o, tTMrO, tTMEM_STOREtO_i) - cute.arch.fence_view_async_tmem_store() - - si_handle.release() - cute.printf("SOFTMAX kt=%d row_max_after=%f row_sum_after=%f\n", kt, row_max, row_sum) - softmax_done_bar.arrive() - - # Wait for MMA's PV[N-1] to commit before reading O for normalize. - final_o_bar.arrive_and_wait() - - # === Final O normalization: O *= 1/row_sum === - inv_row_sum = Float32(1.0) / row_sum - for i in range(n_corr_tiles): - tTMEM_LOADtO_i = cute.make_tensor( - tTMEM_LOADtO.iterator + i * corr_tile_size, - tTMEM_LOADtO.layout, - ) - tTMEM_STOREtO_i = cute.make_tensor( - tTMEM_STOREtO.iterator + i * corr_tile_size, - tTMEM_STOREtO.layout, - ) - tTMrO = cute.make_rmem_tensor(tTMEM_LOADcO.shape, self.acc_dtype) - cute.copy(tiled_tmem_load_o, tTMEM_LOADtO_i, tTMrO) - cute.arch.fence_view_async_tmem_load() - for k in cutlass.range(cute.size(tTMrO), vectorize=True): - tTMrO[k] = tTMrO[k] * inv_row_sum - cute.copy(tiled_tmem_store_o, tTMrO, tTMEM_STOREtO_i) - cute.arch.fence_view_async_tmem_store() - - # Standard epilogue: TMEM → SMEM → GMEM via TMA store. - # O in TMEM is now scaled by 1/row_sum. - tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout) - acc_cons_st = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Consumer, self.num_acc_stage - ) - c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) - c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp) - acc_cons_st = utils.gemm.sm100.epilogue_tma_store( - self, tidx, warp_idx, tma_c, tCtO_base, sC, tCgC, epi_tile, - 0, const_expr(lambda x: x), (0, 0, 0), - acc_cons_st, acc_pipe, c_pipe, - ) - c_pipe.producer_tail() - - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def test(): - torch.manual_seed(42) - for n in [128, 256, 512, 1024]: - torch.manual_seed(42) - m, hd = 128, HEAD_DIM - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(n, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(n, hd, dtype=torch.bfloat16, device='cuda') - v_kernel = v.unsqueeze(-1) - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - - qf = q[:, :, 0].float() - kf = k[:, :, 0].float() - scale = 1.0 / math.sqrt(hd) - attn = qf @ kf.T * scale - attn = torch.softmax(attn, dim=-1) - ref = attn @ v.float() - - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - kernel = FmhaV3StageCMulti(s_k=n) - print(f'n={n}: Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) - print(f'n={n}: tmem s0={kernel.tmem_s0_offset} p0={kernel.tmem_p0_offset} ' - f'o0={kernel.tmem_o0_offset} alloc={kernel.num_tmem_alloc_cols} ' - f'kv_tx_bytes={kernel.kv_tx_bytes}', flush=True) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - - out = c[:, :, 0].float() - cos = torch.nn.functional.cosine_similarity( - out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0) - ).item() - max_abs = (out - ref).abs().max().item() - n_tiles = n // 128 - print(f'FMHA Stage-C Multi n={n} ({n_tiles} kv tiles): ' - f'cos {cos:.6f} max_abs {max_abs:.4f} ' - f'{"PASS" if cos >= 0.99 else "FAIL"}') - if cos < 0.99: - print(f' out[0,:4]={out[0,:4].tolist()}') - print(f' ref[0,:4]={ref[0,:4].tolist()}') - - -if __name__ == '__main__': - test() \ No newline at end of file