From 4b9eed02e1addbc90938c9642a42617e9a8382ca Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 30 May 2026 21:08:12 +0000 Subject: [PATCH] Cleanup C1-C7: delete dead CuTeDSL FMHA, test probes, scratch files - Deleted fmha.py (CuTeDSL slow path), FmhaKernel, Python KV merge - Deleted fmha_sm100.cuh, fmha_sm100_tc.cuh, fmha_sm100_launch.cu, fmha_epilogue_sm100.cuh - Moved fmha_qk_verify.cuh to tests/unit/qk_verify_kernel.cuh - Deleted decode_sparse.py, decode_swa.py, kernels/decode/ - Deleted 46 test_d*.py probes, test_smem_*, test_cotiled_*, test_tmem_*, test_smem_p_*, test_ultra_minimal, test_fmha_pv16, test_working_softmax_maybe - Deleted root scratch: debug_linear.py, test_mapping.py, run_router_tests.py - Moved archive/ to archived_plans/code_archive/ - Rewrote production.py: single fast path via 6-warp multi-tile kernel - Added STATUS.md, audit_attention_live.md - Moved NEXT_PRIORITIES*.md to archived_plans/ --- STATUS.md | 39 ++ .../NEXT_PRIORITIES.md | 0 archived_plans/NEXT_PRIORITIES_PART_2.md | 401 +++++++++++ archived_plans/audit_attention_live.md | 66 ++ .../fmha_6warp_tma_driver_api.cuh | 0 .../code_archive}/fmha_backup_pre_epilog.py | 0 .../code_archive}/fmha_backup_v2.py | 0 .../code_archive}/fmha_smem_acc.py | 0 .../code_archive}/fmha_tma_driver_api.cuh | 0 debug_linear.py | 59 -- dsv4/kernels/attention/fmha.py | 645 ------------------ .../kernels/attention/fmha_epilogue_sm100.cuh | 204 ------ dsv4/kernels/attention/fmha_sm100.cuh | 95 --- dsv4/kernels/attention/fmha_sm100_launch.cu | 86 --- dsv4/kernels/attention/fmha_sm100_tc.cuh | 233 ------- dsv4/kernels/attention/production.py | 375 +--------- dsv4/kernels/decode/_NOTES_fp8_bf16.md | 26 - dsv4/kernels/decode/__init__.py | 0 dsv4/ops/decode_sparse.py | 353 ---------- dsv4/ops/decode_swa.py | 375 ---------- run_router_tests.py | 97 --- test_mapping.py | 50 -- .../unit/qk_verify_kernel.cuh | 0 tests/unit/test_cotiled_diag.py | 188 ----- tests/unit/test_d15_in_kernel_rescale.py | 152 ----- tests/unit/test_d15_multi_kv.py | 192 ------ tests/unit/test_d15_noop_rescale.py | 77 --- tests/unit/test_d15_rescale_debug.py | 116 ---- tests/unit/test_d15_roundtrip_iso.py | 108 --- tests/unit/test_d1_3_cotiled.py | 316 --------- tests/unit/test_d1_3_layout_diag.py | 197 ------ tests/unit/test_d1_3_smem_diag.py | 70 -- tests/unit/test_d1_3_smem_direct.py | 186 ----- tests/unit/test_d1_3_smem_vs_tmem.py | 88 --- tests/unit/test_d1_3_unnorm_debug.py | 96 --- tests/unit/test_d1_3_write_read.py | 232 ------- tests/unit/test_d1_debug.py | 34 - tests/unit/test_d1_diag.py | 92 --- tests/unit/test_d1_diag2.py | 146 ---- tests/unit/test_d1_hd512.py | 48 -- tests/unit/test_d1_hd512_merge.py | 137 ---- tests/unit/test_d1_hd512_only.py | 88 --- tests/unit/test_d1_kv_merge.py | 119 ---- tests/unit/test_d1_kv_merge_v2.py | 122 ---- tests/unit/test_d1_kv_merge_v3.py | 164 ----- tests/unit/test_d1_lse.py | 60 -- tests/unit/test_d1_lse_verify.py | 81 --- tests/unit/test_d1_multi_kv.py | 122 ---- tests/unit/test_d1_qk512.py | 164 ----- tests/unit/test_d1_raw.py | 52 -- tests/unit/test_d1_regression.py | 56 -- tests/unit/test_d1_rescale_debug.py | 115 ---- tests/unit/test_d1_rescale_diag.py | 103 --- tests/unit/test_d1_rescale_min.py | 115 ---- tests/unit/test_d1_smem128.py | 49 -- tests/unit/test_d1_sweep.py | 37 - tests/unit/test_d1_tmem_only.py | 71 -- tests/unit/test_d1_tmem_trip.py | 102 --- tests/unit/test_d2_headpacked.py | 199 ------ tests/unit/test_d2_multicta.py | 131 ---- tests/unit/test_d2_multihead.py | 118 ---- tests/unit/test_d2_perhead.py | 138 ---- tests/unit/test_d2_regression.py | 173 ----- tests/unit/test_d2_scale.py | 106 --- tests/unit/test_d3_inkernel_mask.py | 211 ------ tests/unit/test_d3_swa_mask.py | 160 ----- tests/unit/test_d4_causal_mask.py | 180 ----- tests/unit/test_d5b_perrow_lse.py | 242 ------- tests/unit/test_d5c_fused.py | 304 --------- tests/unit/test_d5c_multitile.py | 166 ----- tests/unit/test_fmha_pv16.py | 52 -- tests/unit/test_fmha_v3_stage_d1.py | 113 --- tests/unit/test_fmha_v3_stage_d5b.py | 204 ------ tests/unit/test_production.py | 2 +- tests/unit/test_smem_acc.py | 118 ---- tests/unit/test_smem_budget.py | 71 -- tests/unit/test_smem_p_coord.py | 137 ---- tests/unit/test_smem_p_diag.py | 516 -------------- tests/unit/test_smem_p_write.py | 199 ------ tests/unit/test_tmem_budget.py | 100 --- tests/unit/test_tmem_roundtrip_minimal.py | 154 ----- tests/unit/test_ultra_minimal.py | 31 - tests/working_softmax_maybe.py | 497 -------------- 83 files changed, 534 insertions(+), 10987 deletions(-) create mode 100644 STATUS.md rename NEXT_PRIORITIES.md => archived_plans/NEXT_PRIORITIES.md (100%) create mode 100644 archived_plans/NEXT_PRIORITIES_PART_2.md create mode 100644 archived_plans/audit_attention_live.md rename {dsv4/kernels/attention/archive => archived_plans/code_archive}/fmha_6warp_tma_driver_api.cuh (100%) rename {dsv4/kernels/attention/archive => archived_plans/code_archive}/fmha_backup_pre_epilog.py (100%) rename {dsv4/kernels/attention/archive => archived_plans/code_archive}/fmha_backup_v2.py (100%) rename {dsv4/kernels/attention/archive => archived_plans/code_archive}/fmha_smem_acc.py (100%) rename {dsv4/kernels/attention/archive => archived_plans/code_archive}/fmha_tma_driver_api.cuh (100%) delete mode 100644 debug_linear.py delete mode 100644 dsv4/kernels/attention/fmha.py delete mode 100644 dsv4/kernels/attention/fmha_epilogue_sm100.cuh delete mode 100644 dsv4/kernels/attention/fmha_sm100.cuh delete mode 100644 dsv4/kernels/attention/fmha_sm100_launch.cu delete mode 100644 dsv4/kernels/attention/fmha_sm100_tc.cuh delete mode 100644 dsv4/kernels/decode/_NOTES_fp8_bf16.md delete mode 100644 dsv4/kernels/decode/__init__.py delete mode 100644 dsv4/ops/decode_sparse.py delete mode 100644 dsv4/ops/decode_swa.py delete mode 100644 run_router_tests.py delete mode 100644 test_mapping.py rename dsv4/kernels/attention/fmha_qk_verify.cuh => tests/unit/qk_verify_kernel.cuh (100%) delete mode 100644 tests/unit/test_cotiled_diag.py delete mode 100644 tests/unit/test_d15_in_kernel_rescale.py delete mode 100644 tests/unit/test_d15_multi_kv.py delete mode 100644 tests/unit/test_d15_noop_rescale.py delete mode 100644 tests/unit/test_d15_rescale_debug.py delete mode 100644 tests/unit/test_d15_roundtrip_iso.py delete mode 100644 tests/unit/test_d1_3_cotiled.py delete mode 100644 tests/unit/test_d1_3_layout_diag.py delete mode 100644 tests/unit/test_d1_3_smem_diag.py delete mode 100644 tests/unit/test_d1_3_smem_direct.py delete mode 100644 tests/unit/test_d1_3_smem_vs_tmem.py delete mode 100644 tests/unit/test_d1_3_unnorm_debug.py delete mode 100644 tests/unit/test_d1_3_write_read.py delete mode 100644 tests/unit/test_d1_debug.py delete mode 100644 tests/unit/test_d1_diag.py delete mode 100644 tests/unit/test_d1_diag2.py delete mode 100644 tests/unit/test_d1_hd512.py delete mode 100644 tests/unit/test_d1_hd512_merge.py delete mode 100644 tests/unit/test_d1_hd512_only.py delete mode 100644 tests/unit/test_d1_kv_merge.py delete mode 100644 tests/unit/test_d1_kv_merge_v2.py delete mode 100644 tests/unit/test_d1_kv_merge_v3.py delete mode 100644 tests/unit/test_d1_lse.py delete mode 100644 tests/unit/test_d1_lse_verify.py delete mode 100644 tests/unit/test_d1_multi_kv.py delete mode 100644 tests/unit/test_d1_qk512.py delete mode 100644 tests/unit/test_d1_raw.py delete mode 100644 tests/unit/test_d1_regression.py delete mode 100644 tests/unit/test_d1_rescale_debug.py delete mode 100644 tests/unit/test_d1_rescale_diag.py delete mode 100644 tests/unit/test_d1_rescale_min.py delete mode 100644 tests/unit/test_d1_smem128.py delete mode 100644 tests/unit/test_d1_sweep.py delete mode 100644 tests/unit/test_d1_tmem_only.py delete mode 100644 tests/unit/test_d1_tmem_trip.py delete mode 100644 tests/unit/test_d2_headpacked.py delete mode 100644 tests/unit/test_d2_multicta.py delete mode 100644 tests/unit/test_d2_multihead.py delete mode 100644 tests/unit/test_d2_perhead.py delete mode 100644 tests/unit/test_d2_regression.py delete mode 100644 tests/unit/test_d2_scale.py delete mode 100644 tests/unit/test_d3_inkernel_mask.py delete mode 100644 tests/unit/test_d3_swa_mask.py delete mode 100644 tests/unit/test_d4_causal_mask.py delete mode 100644 tests/unit/test_d5b_perrow_lse.py delete mode 100644 tests/unit/test_d5c_fused.py delete mode 100644 tests/unit/test_d5c_multitile.py delete mode 100644 tests/unit/test_fmha_pv16.py delete mode 100644 tests/unit/test_fmha_v3_stage_d1.py delete mode 100644 tests/unit/test_fmha_v3_stage_d5b.py delete mode 100644 tests/unit/test_smem_acc.py delete mode 100644 tests/unit/test_smem_budget.py delete mode 100644 tests/unit/test_smem_p_coord.py delete mode 100644 tests/unit/test_smem_p_diag.py delete mode 100644 tests/unit/test_smem_p_write.py delete mode 100644 tests/unit/test_tmem_budget.py delete mode 100644 tests/unit/test_tmem_roundtrip_minimal.py delete mode 100644 tests/unit/test_ultra_minimal.py delete mode 100644 tests/working_softmax_maybe.py diff --git a/STATUS.md b/STATUS.md new file mode 100644 index 00000000..bd5ee313 --- /dev/null +++ b/STATUS.md @@ -0,0 +1,39 @@ +# STATUS — DSV4 Inference Kernel (post-cleanup 2026-05-30) + +## Production Path + +**One FMHA kernel:** `fmha_6warp_tma_multirow_multitile.cuh` — 6-warp, TMA, UMMA, tcgen05.mma SS, in-kernel multi-tile SMEM accumulator, multi-row softmax. Loaded via `fmha_multitile_capi.cu` (C API) + `fmha_multitile_op.py` (ctypes). Dispatched from `production.py`. + +**Head dims:** 64, 128, 256, 512. **T=1 decode** proven (cos ≥ 0.999996). **T>1 prefill** via multi-row path (P5, P7). + +**No CuTeDSL runtime dependency.** All kernel code is raw CUDA C++. CuTeDSL (fmha.py) deleted; Python KV merge deleted; `FmhaKernel` deleted. + +## Live Attention Files + +| File | Role | +|---|---| +| `fmha_6warp_tma_multirow_multitile.cuh` | Production kernel | +| `fmha_common.cuh` | Shared types/defs | +| `fmha_tma.cuh` | TMA descriptor helpers | +| `fmha_umma_desc.cuh` | UMMA descriptor creation | +| `fmha_multitile_capi.cu` | C API wrapper (nvcc compiled) | +| `fmha_multitile_op.py` | ctypes loader | +| `production.py` | Public API (dsv4_attention) | +| `__init__.py` | Bridge to layers (sparse/dense/swa) | + +## Stage E Checklist (from ROADMAP/NEXT_PRIORITIES_PART_2) + +- [ ] **E1:** Wire `LayerCacheHandle` → `gather_compressed_kv`, `gather_all_compressed_kv`, `gather_swa_kv`, `num_query_heads`, `head_dim` +- [ ] **E2:** End-to-end smoke test through one full layer +- [ ] **E3:** Top-level `model/dsv4.py` (currently 2-line TODO) +- [ ] **E4:** Delete `torch.cuda.synchronize()` from fast path +- [ ] **E5:** Fold batch loop into kernel grid +- [ ] **E6:** FP4 output fusion for FMHA → wo_a +- [ ] **E7:** Lightning indexer FP4 tensor-core scoring +- [ ] **E8:** Multi-CTA grid for prefill +- [ ] **E9:** CUDA graph capture + +## Cleanup Done (C1–C7) + +- Deleted: fmha.py, fmha_sm100.cuh, fmha_sm100_tc.cuh, fmha_sm100_launch.cu, fmha_epilogue_sm100.cuh, fmha_qk_verify.cuh (moved to tests/unit/), decode_sparse.py, decode_swa.py, kernels/decode/, 46 test_d*.py probes, root scratch files, archive/ (moved to archived_plans/code_archive/) +- Removed: FmhaKernel import, CuTeDSL slow path, Python KV merge, torch.cuda.synchronize in _run_fmha_segmented (function deleted) diff --git a/NEXT_PRIORITIES.md b/archived_plans/NEXT_PRIORITIES.md similarity index 100% rename from NEXT_PRIORITIES.md rename to archived_plans/NEXT_PRIORITIES.md diff --git a/archived_plans/NEXT_PRIORITIES_PART_2.md b/archived_plans/NEXT_PRIORITIES_PART_2.md new file mode 100644 index 00000000..1f719792 --- /dev/null +++ b/archived_plans/NEXT_PRIORITIES_PART_2.md @@ -0,0 +1,401 @@ +# ROADMAP — Stage E and the path to a runnable model + +**Method.** Every state claim below was verified against the source in this +revision. Agent status doc and agent assessment were both read, then ignored +where they contradict the code. The doctrine rules from prior issue docs apply: +DSL wall → raw CUDA C++ (not Python); raw CUDA ≠ scalar math; print, don't +guess; integration over exploration; falsifiable gates only. + +--- + +## TL;DR — where we actually are + +**The FMHA kernel itself is in genuinely good shape.** `fmha_6warp_tma_multirow_multitile.cuh` +is one file, properly Blackwell-native (`tcgen05.mma` SS, UMMA descriptors, +TMEM, `cp.async.bulk.tensor.2d` + `mbarrier.arrive.expect_tx`, in-kernel +multi-tile rescale via SMEM accumulator, multi-row softmax). The TMA hang from +P4 was correctly diagnosed (missing `expect_tx`) and resolved. The C-API + +ctypes shim works. **P3, P4, P5, P6, P7, P8 numerics are done.** + +**The agent's "Stage E" assessment is correct in spirit but wrong on specifics.** +Where it says "AttentionSubBlock has `NotImplementedError` stubs" — it doesn't. +The layer is structurally complete and *imports the right things*. The real +problem is one layer down: it calls cache methods like `gather_compressed_kv()`, +`gather_swa_kv()`, and properties like `cache.num_query_heads` / `cache.head_dim` +**that don't exist on `LayerCacheHandle`**. The handle exposes `read_*_view()` +returning paged FP8 buffers; nothing materializes them into the BF16 dense +tensors `dsv4_attention` consumes. That's the integration gap, not stubs. + +**Three structural shortcuts still live in production code.** They are not +performance-fatal yet but they will be at scale: +1. Two `torch.cuda.synchronize()` calls remain on the kernel-launch path + (`fmha_multitile_op.py:130`, `production.py:279`). The second is in the still-resident + "slow path." The first is unconditional, even on the fast path. +2. The slow CuTeDSL+Python-merge path is **not deleted**, only relabeled + "slow path." 471 lines of code, 48 unit tests bound to its `FmhaKernel` symbol, + still imported at `production.py:53`. +3. The batched call shape (4D q) does a **Python for-loop over batch items** + at `production.py:380` instead of folding batch into the grid. + +**Cleanup is non-trivial.** ~22 dead/stale code paths live alongside the +working ones. Some are old CuTeDSL drafts (`ops/decode_sparse.py`, +`ops/decode_swa.py`), some are pre-consolidation Stage D probes +(`fmha_sm100.cuh`, `fmha_sm100_tc.cuh`, `fmha_sm100_launch.cu`, +`fmha_qk_verify.cuh`, `fmha_epilogue_sm100.cuh`), some are misplaced root files +(`debug_linear.py`, `test_mapping.py`, `tests/working_softmax_maybe.py`), and 46 +of 177 unit tests are pinned to the deprecated `FmhaKernel`. Left in place, +this is the primary path by which an agent will drift back into the wrong +codebase the next time it context-switches. + +--- + +# PART 1 — CLEANUP (do this first, in this order) + +Cleanup is gated *before* Stage E work because the next stage's failure mode is +"agent picks the wrong file to extend." Right now `dsv4/kernels/attention/` has +14 files; only 6 are live. That ratio is the diagnostic. + +## C1. Confirm what's live with `grep`, not by reading status docs + +Before deleting anything: produce `audit_attention_live.md` listing each file +under `dsv4/kernels/attention/` and `dsv4/ops/` with its **outbound reference +count from non-archive code** (production, layers, model, integration tests). +Anything with 0 refs is a deletion candidate; anything with >0 refs gets +inspected to confirm the reference is live, not a stale import. + +**Falsifiable gate:** the audit table goes into `archived_plans/audit_attention_live.md`, +file count column adds up to the actual directory listing. Mention deletion +candidates by name with the ref count next to them. + +## C2. Delete the dead CuTeDSL FMHA scaffolding + +These have **zero external refs** outside their own `#include` chain and +self-references: + +| File | Status | Action | +|---|---|---| +| `dsv4/kernels/attention/fmha_sm100.cuh` | Phase-1 scalar reference, header says "WORKING cos 0.999999" | Delete. The 6-warp kernel supersedes it. | +| `dsv4/kernels/attention/fmha_sm100_tc.cuh` | Earlier tensor-core fork before consolidation. | Delete. The 6-warp kernel is its successor. | +| `dsv4/kernels/attention/fmha_sm100_launch.cu` | Header says "COMPILES but doesn't run via torch.utils.cpp_extension" | Delete. Replaced by `fmha_multitile_capi.cu` + ctypes. | +| `dsv4/kernels/attention/fmha_qk_verify.cuh` | Header says "QK GEMM verification" — one test references it. | Move to `tests/unit/qk_verify_kernel.cuh` (it's a test fixture, not a library kernel). | +| `dsv4/kernels/attention/fmha_epilogue_sm100.cuh` | Phase-2 epilogue prototype; the multi-tile kernel already does in-kernel rescale via SMEM and doesn't use this. | Delete after confirming it's not `#include`'d by anything live. Audit table from C1 answers this. | +| `dsv4/ops/decode_sparse.py` | Old CuTeDSL CSA decode draft, not called from anywhere. | Delete. Path is now `layers/attention.py:_forward_csa` → `kernels/attention.sparse_fmha_with_swa` → `kernels/attention/production.py`. | +| `dsv4/ops/decode_swa.py` | Old CuTeDSL SWA decode draft, same story. | Delete. | +| `dsv4/kernels/decode/` | Two files: `__init__.py` (empty) and `_NOTES_fp8_bf16.md` (scratch). | Delete the whole directory. | + +**Falsifiable gate:** after C2, `find dsv4/kernels/attention -name '*.cuh' -o -name '*.cu' | wc -l` +returns the value in `audit_attention_live.md`'s live count. No reference +warnings on `python -c "import dsv4"`. + +## C3. Decide the fate of the CuTeDSL slow path (`fmha.py`) + +`fmha.py` is 645 lines and is the entire pre-P3-P8 CuTeDSL path. It is currently +imported by `production.py:53` (`from dsv4.kernels.attention.fmha import FmhaKernel`) +and used by the `else` branch starting at `production.py:410` ("SLOW PATH"). It +is also imported by **48 unit tests**. + +**Two options. Pick one, document it, do not maintain both.** + +- **Option A (preferred, matches doctrine).** Delete `fmha.py`, delete the + SLOW PATH branch from `production.py`, delete the 48 `test_d*.py` unit tests + bound to `FmhaKernel`. The 6-warp multi-tile kernel covers + `hd∈{64,128,256,512}` × `T=1` decode × any `n_segments`. That's the production + surface. Everything else is exploration that already paid off — it does not + need to ship. +- **Option B (fallback, only if Option A is blocked).** Move `fmha.py` and the + 48 tests to `archived_plans/cutedsl_legacy/`. Delete the SLOW PATH from + `production.py` regardless. The kernel and tests stay buildable for reference + but are no longer in the active import graph. + +**Failure mode to refuse:** "keep it as fallback in case the 6-warp kernel has +a regression." If there's a regression, you fix the 6-warp kernel — that's the +production kernel. Carrying a parallel implementation guarantees both rot. + +**Falsifiable gate:** `grep -r "FmhaKernel" dsv4/ tests/ scripts/ | wc -l` returns +0 (Option A) or only matches inside `archived_plans/` (Option B). The +SLOW PATH branch is gone from `production.py` either way. + +## C4. Empty the `archive/` subdirectory the same way + +`dsv4/kernels/attention/archive/` currently holds 5 backups +(`fmha_backup_pre_epilog.py`, `fmha_backup_v2.py`, `fmha_smem_acc.py`, +`fmha_6warp_tma_driver_api.cuh`, `fmha_tma_driver_api.cuh`). These were correct +to *not* delete during P8 consolidation. Now: move them out of `dsv4/` +entirely into `archived_plans/code_archive/` and delete `dsv4/kernels/attention/archive/`. +Active code directories should contain only active code. `archived_plans/` is +the right place for everything that isn't. + +**Falsifiable gate:** `find dsv4 -name 'archive' -type d | wc -l` returns 0. + +## C5. Move the root-level scratch files + +Three files live at the repo root and don't belong there: + +- `debug_linear.py` (58 lines, computes expected O for a debug pattern — Stage D + probe). → `tests/unit/archive/` or delete if no longer useful. +- `test_mapping.py` (49 lines, SMEM-P coordinate mapping check — Stage D probe). → same. +- `run_router_tests.py` (hardcoded path `/root/dsv4-nvfp4-workspace/kernel`). → + Either generalize the path and move to `scripts/`, or delete. Hardcoded + absolute paths in tracked files are bug magnets. +- `tests/working_softmax_maybe.py` (496 lines, name speaks for itself). → If + it's a draft, archive it. If it's superseded, delete it. The doctrine is + "code or archive, never `maybe`." + +**Falsifiable gate:** repo root has no `.py` files except generated/build +artifacts. `tests/` contains no `*_maybe.py` or `working_*.py`. + +## C6. Prune the unit test D-probe explosion + +46 `test_d*.py` files were Stage D debugging probes. Many are bound to +`FmhaKernel` and will be deleted in C3 if you pick Option A. The remainder that +test components still in use (`FmhaTmaMultiRowMultiTile`, the indexer, the +compressor, mHC, MoE) should either be promoted to named integration tests +(`test_fmha_decode.py`, `test_csa_indexer.py`, etc.) or archived. + +**Falsifiable gate:** `ls tests/unit/test_d*.py | wc -l` returns ≤ 5 +(only the ones genuinely testing currently-live behavior, with descriptive names). + +## C7. Status doc consolidation + +After C1–C6, write **one** new top-level `STATUS.md` describing the **post-cleanup** +state of the production path, and delete `NEXT_PRIORITIES.md` (this doc replaces it). +`archived_plans/` already holds the historical paper trail; do not add to it +during cleanup. + +**Falsifiable gate:** repo root has exactly two living `.md` files: `README.md` +and `STATUS.md` (plus this `ROADMAP.md`). `archived_plans/` holds the history, +unchanged. + +--- + +# PART 2 — STAGE E PRIORITY ORDER (post-cleanup) + +The agent is right that Stage E = "make the model run end-to-end." The order +below is sequenced so each priority unblocks the next, and so the +performance-decisive moves come *after* end-to-end correctness, not before. + +## E1. Wire `LayerCacheHandle` to expose what `AttentionSubBlock` calls + +This is the load-bearing missing piece. `dsv4/kernels/attention/__init__.py` +calls four things that don't exist on `dsv4/cache/handle.py`: + +- `cache.gather_compressed_kv(selected_indices) → (k_compressed, v_compressed)` BF16 dense +- `cache.gather_all_compressed_kv() → (k_compressed, v_compressed)` BF16 dense (HCA, no top-k) +- `cache.gather_swa_kv() → (k_swa, v_swa)` BF16 dense +- `cache.num_query_heads`, `cache.head_dim` (properties) + +The handle already exposes the right *raw* state via +`read_classical_view()` (FP8 entries + RoPE BF16 + paged block_table) and +`read_swa_view()` (FP8 entries + RoPE BF16 + positions). The work is the +gather+dequantize+concat path. + +**Constraints (doctrine):** + +- Each gather is a **single fused kernel**, raw CUDA C++ if CuTeDSL can't do + block-table-strided FP8→BF16 dequant+gather in one pass. Not a Python loop + over selected indices. Not five eager torch ops. +- The kernel uses TMA + warp-level dequant, not scalar loads. The MoE GEMM + already proves out the FP8 dequant pattern on Blackwell — reuse it. +- The dequant decodes E4M3 with the inverse-scale stored alongside. **Print + the actual layout of `paged.entries_fp8` + `paged.inv_scale` before writing + the kernel** — do not assume the layout matches an MoE convention. + +**Definition of done:** +1. The four methods exist on `LayerCacheHandle`, implemented in + `dsv4/kernels/cache/gather.{cu,py}`. +2. `python -c "from dsv4.layers.attention import AttentionSubBlock; ..."` runs + a forward pass on one layer without AttributeError, NotImplementedError, + or shape mismatch. +3. Numerical parity vs a pure-Python reference (FP32 dequant + dense FMHA over + the gathered tensors): `cos ≥ 0.9995`. FP8 dequant loses a little, that's + expected. +4. **Launch count for one CSA layer** measured with Nsight: gather is ≤ 3 launches + (compressed gather, SWA gather, concat). Not 30, not 5+5+cat. + +**Failure modes to refuse:** +- A Python gather loop over the block table. The block table is on GPU; the + gather is on GPU. No host trips. +- BF16 materialization "for now" without FP8 dequant. The paged cache stores + FP8 because that's the design — bypassing it loses 2x KV memory and the + whole point of the cache structure. + +## E2. End-to-end smoke test through one full layer + +After E1, both halves of `dsv4/model/layer.py` (`mhc_attn` → `attn.forward` → +`mhc_ffn` → `ffn.forward`) should execute without error. The test is the +forcing function. + +**Definition of done:** +1. `tests/e2e/test_one_layer.py` runs `DSV4Layer.forward(X, token_ids, cache)` + on synthetic inputs with cached weights loaded from `loader/`, returns a + tensor of the right shape and dtype, no errors. +2. The same test with `cache` containing actual paged FP8 entries (not zeros) + produces output with `cos ≥ 0.99` against a Python reference of the same + layer (compressor + indexer + FMHA + mHC, all in eager PyTorch). +3. Test exists for all three attention types (CSA, HCA, SWA) by toggling the + layer's `attn_type`. + +**Failure mode to refuse:** wrapping the test in a try/except to "make it +pass." The cos threshold is the gate. + +## E3. Top-level `model/dsv4.py` (currently a 2-line TODO) + +`dsv4/model/dsv4.py` is literally: + +```python +"""Full DSV4 model.""" +# TODO: Phase 1 +``` + +Build the actual model class: embedding → N×DSV4Layer → final norm → prediction +head. Use `layer_schedule.py` to alternate CSA/HCA per the V4 spec +(Pro: layers 0–1 HCA, then interleaved; Flash: layers 0–1 SWA, then +interleaved). + +**Definition of done:** +1. `DSV4Model.forward(token_ids, cache_manager) → logits` works on a single + token decode. +2. Layer schedule matches the V4 paper (verify against `config.py`). +3. MTP module is wired but optional (off by default; doctrine note in code + that MTP fine-tuning is a separate workstream). +4. End-to-end test: `tests/e2e/test_decode.py` generates one token from a + 3-layer toy config, succeeds, cos ≥ 0.99 vs Python reference. + +## E4. Delete the two remaining `torch.cuda.synchronize()` calls on the fast path + +After C3 the slow-path sync at `production.py:279` is gone with the slow path. +The other sync is at `fmha_multitile_op.py:130`: + +```python +# Synchronize to catch async errors +torch.cuda.synchronize() +return o, lse +``` + +This is on the *fast path*. It's there to catch async kernel errors, but the +correct way to do that is to check the return code from the C-API launch +function (which `fmha_multitile_decode_launch` already returns) plus +`cudaPeekAtLastError`, not a full device sync. + +**Definition of done:** +1. The sync is removed from `fmha_multitile_op.py`. +2. Error checking moves to a `cudaPeekAtLastError` call + the existing return + code check. +3. Nsight measurement on a 4-layer decode: zero `cudaDeviceSynchronize` calls + between layer 0 entry and final logits. + +## E5. Fold the batch loop into the kernel grid + +`production.py:380–388` is a Python for-loop over batch items. Each iteration +calls back into `dsv4_attention` on a 3D slice. For batch=8 and 61 layers, +that's 488 launches per decoded step instead of 61. + +The multi-tile kernel's grid is currently `(1, n_h, 1)`. It needs to become +`(1, n_h, batch)` with `blockIdx.z` used to index batch. The kernel already +has `_batch_stride` params (read off `FmhaTmaMultiRowMultiTileParams`); they +just aren't wired to a grid dim. + +**Definition of done:** +1. `production.py:380–388` for-loop deleted; the batched path calls the + multi-tile kernel once with `gridDim.z = batch`. +2. Numerical parity: `cos ≥ 0.999998` vs the deleted Python loop. +3. Launch count for batch=8 decode: 1 per layer, not 8. + +## E6. NVFP4-1.2: FP4 output fusion for FMHA → `wo_a` + +Right now the FMHA writes BF16 to GMEM, then `wo_a` re-quantizes it to FP4. +That's a redundant memory pass on the attention output. The epilogue already +has the one-way TMEM→regs→SMEM→GMEM shape — adding an FP4 pack with amax in +registers, gated by a template param, is the same pattern the MoE epilogue +already uses for NVFP4-1.1. + +**Constraints (doctrine):** +- The FP4 pack is **part of the same single launch**. No second kernel for + quantize. Per-128-element block amax reduces via warp shuffle in registers. +- The amax → scale → E2M1 LUT path matches `dsv4/ops/quantize.py` byte for byte. + Print the LUT from the Python side; assert equality in a unit test. The + indexer LUT bug we already fixed is the cautionary tale. +- This is gated on E4 (`epilogue_op` slot exists in the multi-tile kernel — + let's verify it does before starting E6). + +**Definition of done:** +1. `fmha_6warp_tma_multirow_multitile_kernel` accepts a template arg `Epilogue` + (default = identity); a `FP4Epilogue` specialization packs E2M1 + emits SF + blocks in the layout `wo_a` expects. +2. `tests/unit/test_fp4_epilogue_parity.py`: BF16 reference vs FP4 fused, + `cos ≥ 0.999` after dequant. +3. End-to-end CSA layer: BF16 output → FP4 output replaces one full BF16 R/W + pass through (T, n_h * hd) memory. Measured bandwidth drop with Nsight. + +## E7. Stage F: Lightning indexer FP4 tensor-core scoring + +Currently `indexer_score_topk.cu` uses scalar FP32 cores after the LUT dequant. +The paper §5.2.1 is explicit that this path runs in FP4 on tensor cores +("multiplied entirely in FP4 ... 99.7% recall"). The LUT bug fix made the +scalar path *correct*; this priority makes it *fast*, the way V4 actually +specifies it. + +**Constraints (doctrine):** +- `tcgen05.mma` with `mxf4nvf4` kind. Keys and queries both FP4. Scales E4M3 + in shared mem, hardware decodes. No scalar fallback in the live path. +- Recall vs the FP32 oracle ≥ 99.7% per paper. The recall test from the LUT + fix is the same test; the bar is the same. +- The top-k selector goes warp-level (`__shfl_xor` ballot) at the same time — + the current serial single-thread heap merge at top_k=1024 is its own latency + source on Pro decode. + +**Definition of done:** +1. `indexer_score_topk.cu` has a `tcgen05`-MMA FP4 path, gated by a config + flag, with the scalar path archived (not deleted, since the FP4 path needs + the oracle). +2. Recall@k ≥ 99.7% for k ∈ {512, 1024} at compressed_blocks ∈ {2k, 8k, 32k}. +3. End-to-end CSA layer latency drops measurably (Nsight). + +## E8. Multi-CTA grid for prefill (Priority 4 from the original ROADMAP) + +Decode is single-CTA per head and that's correct. Prefill (T > 1) wants +multi-CTA to parallelize over M. The one-way epilogue from E6 is the +prerequisite — once it lands, the grid expansion is wiring. + +**Definition of done:** +1. `gridDim.x = ceil_div(M, M_TILE)` for `T > 1` paths. +2. Prefill latency at T=128 with M-tile=64: 2-CTA grid, measured speedup ≥ 1.5× + over single-CTA. +3. Decode path unchanged (single CTA, no regression). + +## E9. CUDA graph capture (Stage E item from the agent's notes) + +Once E5 lands and the decode hot path is sync-free, graph capture is +straightforward. This is the biggest single perf win on autoregressive decode +because it eliminates host launch overhead entirely. + +**Definition of done:** +1. `DSV4Model.decode_step(token, cache)` is `cudaGraph`-capturable. +2. Replay latency vs uncaptured: ≥ 2× faster on a 3-layer toy at batch=1. +3. Memory: graph capture doesn't `torch.zeros` on the hot path (E1 already + forced this for gather; verify it holds for compressor + indexer + FMHA). + +--- + +# DOCTRINE — applies to every item above + +1. **DSL wall → raw CUDA C++, not Python.** The C-API + ctypes pattern in + `fmha_multitile_op.py` is the correct shape for new kernels. +2. **Raw CUDA ≠ scalar math.** Every kernel above is `tcgen05` / UMMA / TMA / + warp reductions. No "temporary scalar dot product" without a labeled + replacement target. +3. **Print, don't guess.** Two places this round specifically: + - E1 gather kernels: print the FP8 cache layout + inv-scale layout before + writing the dequant. + - E6 FP4 pack: print the LUT from `quantize.py` and assert in a test that + the kernel produces byte-identical packed nibbles. +4. **Integration over exploration.** Do not create + `fmha_6warp_tma_multirow_multitile_v2.cuh`. Extend the chosen file or stop + and replan. Variant proliferation is the agent's primary failure mode in + this repo. +5. **Falsifiable gates only.** Every "done" above has a number or a binary + check. "Looks fine," "milestone complete," and "Stage X done" without a + number are read as "not done." +6. **The slow path is not a fallback.** Once C3 deletes the CuTeDSL+Python-merge + path, it stays deleted. If the 6-warp kernel breaks, you fix the 6-warp + kernel — not resurrect 471 lines of dead code. \ No newline at end of file diff --git a/archived_plans/audit_attention_live.md b/archived_plans/audit_attention_live.md new file mode 100644 index 00000000..1c8369d7 --- /dev/null +++ b/archived_plans/audit_attention_live.md @@ -0,0 +1,66 @@ +# Attention Directory Audit (C1) + +## dsv4/kernels/attention/ files + +| File | Refs from live code | Action | +|---|---|---| +| `fmha_6warp_tma_multirow_multitile.cuh` | fmha_multitile_capi.cu (include) | KEEP - production kernel | +| `fmha_common.cuh` | fmha_6warp_tma_multirow_multitile.cuh (include) | KEEP - shared defs | +| `fmha_tma.cuh` | fmha_6warp_tma_multirow_multitile.cuh (include) | KEEP - TMA helpers | +| `fmha_umma_desc.cuh` | fmha_multitile_capi.cu (include) | KEEP - UMMA descriptors | +| `fmha_multitile_capi.cu` | fmha_multitile_op.py (nvcc compile) | KEEP - C API wrapper | +| `fmha_multitile_op.py` | production.py:312 (import) | KEEP - ctypes loader | +| `production.py` | __init__.py, custom_ops.py, layers/attention.py, tests | KEEP - public API | +| `__init__.py` | layers/attention.py | KEEP - public API | +| `fmha.py` | production.py:53 (FmhaKernel import), ~48 test files | DELETE (C3 Option A) | +| `fmha_sm100.cuh` | Only referenced by files being deleted | DELETE | +| `fmha_sm100_tc.cuh` | Zero live refs | DELETE | +| `fmha_sm100_launch.cu` | Zero live refs | DELETE | +| `fmha_epilogue_sm100.cuh` | Only ref is fmha_sm100_launch.cu (deleting) | DELETE | +| `fmha_qk_verify.cuh` | Only ref is tests/unit/test_qk_mma.cu | MOVE to tests/unit/ | +| `archive/` (5 files) | Zero live refs | MOVE to archived_plans/code_archive/ | + +## dsv4/ops/ dead files + +| File | Refs | Action | +|---|---|---| +| `decode_sparse.py` | Zero | DELETE | +| `decode_swa.py` | Zero | DELETE | + +## dsv4/kernels/decode/ (entire dir) + +| File | Refs | Action | +|---|---|---| +| `__init__.py` | Zero | DELETE dir | +| `_NOTES_fp8_bf16.md` | Zero | DELETE dir | + +## Root-level scratch + +| File | Action | +|---|---| +| `debug_linear.py` | DELETE | +| `test_mapping.py` | DELETE | +| `run_router_tests.py` | DELETE (hardcoded B200 path) | +| `tests/working_softmax_maybe.py` | DELETE | + +## test_d*.py count: 46 → target ≤5 + +All 46 reference FmhaKernel (being deleted) or are Stage D probes. +All DELETE under C3 Option A. + +## test_smem_acc.py — imports FmhaKernel from archive. DELETE. + +## test_cotiled_diag.py, test_fmha_v3_stage_d1.py, test_tmem_roundtrip_minimal.py, test_smem_budget.py — import FmhaKernel. DELETE with C3. + +Live test files (KEEP): +- test_p3_fast_decode.py (tests multitile op) +- test_p6_tma_epilogue.py (tests multitile op) +- test_p7_multi_row_softmax.py (tests multitile op) +- test_production.py (tests production API) +- test_fmha_sm100.py (tests standalone .cu compile) +- All test_*.cu files (standalone CUDA, don't import FmhaKernel) +- test_nvfp4_*.py (quantization tests) +- test_cutedsl.py (CuTeDSL infra) + +Live .cuh count after cleanup: 4 (fmha_6warp_tma_multirow_multitile, fmha_common, fmha_tma, fmha_umma_desc) +Live .cu count: 1 (fmha_multitile_capi.cu) diff --git a/dsv4/kernels/attention/archive/fmha_6warp_tma_driver_api.cuh b/archived_plans/code_archive/fmha_6warp_tma_driver_api.cuh similarity index 100% rename from dsv4/kernels/attention/archive/fmha_6warp_tma_driver_api.cuh rename to archived_plans/code_archive/fmha_6warp_tma_driver_api.cuh diff --git a/dsv4/kernels/attention/archive/fmha_backup_pre_epilog.py b/archived_plans/code_archive/fmha_backup_pre_epilog.py similarity index 100% rename from dsv4/kernels/attention/archive/fmha_backup_pre_epilog.py rename to archived_plans/code_archive/fmha_backup_pre_epilog.py diff --git a/dsv4/kernels/attention/archive/fmha_backup_v2.py b/archived_plans/code_archive/fmha_backup_v2.py similarity index 100% rename from dsv4/kernels/attention/archive/fmha_backup_v2.py rename to archived_plans/code_archive/fmha_backup_v2.py diff --git a/dsv4/kernels/attention/archive/fmha_smem_acc.py b/archived_plans/code_archive/fmha_smem_acc.py similarity index 100% rename from dsv4/kernels/attention/archive/fmha_smem_acc.py rename to archived_plans/code_archive/fmha_smem_acc.py diff --git a/dsv4/kernels/attention/archive/fmha_tma_driver_api.cuh b/archived_plans/code_archive/fmha_tma_driver_api.cuh similarity index 100% rename from dsv4/kernels/attention/archive/fmha_tma_driver_api.cuh rename to archived_plans/code_archive/fmha_tma_driver_api.cuh diff --git a/debug_linear.py b/debug_linear.py deleted file mode 100644 index d7aaf2ed..00000000 --- a/debug_linear.py +++ /dev/null @@ -1,59 +0,0 @@ -#!/usr/bin/env python3 -""" -Compute expected O for linear pattern P_ij = i*128 + j. -Use same random seed as test. -""" -import torch -import math - -torch.manual_seed(42) -hd = 256 -n_kv = 128 -scale_softmax = 1.0 / math.sqrt(hd) - -# Generate random Q,K,V as in test -q = torch.randn(128, hd, dtype=torch.bfloat16, device='cuda') -k = torch.randn(128, hd, dtype=torch.bfloat16, device='cuda') -v = torch.randn(128, hd, dtype=torch.bfloat16, device='cuda') - -# Compute reference P via softmax -scores = (q @ k.mT) * scale_softmax -p_ref = torch.softmax(scores, dim=-1) # 128x128 - -# Linear pattern P -p_linear = torch.zeros(128, 128, dtype=torch.bfloat16, device='cuda') -for i in range(128): - for j in range(128): - p_linear[i,j] = i*128 + j - -# Compute O for both patterns -o_ref = p_ref @ v -o_linear = p_linear @ v - -print("Reference O shape:", o_ref.shape) -print("Linear O shape:", o_linear.shape) -print("\nFirst row, first 4 cols:") -print("O_ref[0,:4] =", o_ref[0,:4,0].tolist()) -print("O_lin[0,:4] =", o_linear[0,:4,0].tolist()) - -# Compute expected scaling if mapping correct -# Kernel output for hd=256: out[0,:4]=[0.029296875, 0.0164794921875, -0.029541015625, 0.02294921875] -kernel_out = [0.029296875, 0.0164794921875, -0.029541015625, 0.02294921875] -print("\nKernel out[0,:4] =", kernel_out) - -# Compare with linear pattern -print("\nDifference kernel vs linear:") -for i in range(4): - diff = kernel_out[i] - o_linear[0,i,0].item() - print(f"col {i}: kernel {kernel_out[i]:.6f} vs linear {o_linear[0,i,0].item():.6f} diff={diff:.6f}") - -# Compute cosine similarity -kernel_tensor = torch.tensor(kernel_out, dtype=torch.float32) -linear_tensor = o_linear[0,:4,0].float() -cos = torch.cosine_similarity(kernel_tensor, linear_tensor, dim=0).item() -print(f"\nCosine similarity (first 4 cols): {cos:.6f}") - -# Also compute expected O for P=1.0 pattern -p_one = torch.ones(128,128, dtype=torch.bfloat16, device='cuda') -o_one = p_one @ v -print("\nP=1.0 pattern O[0,:4] =", o_one[0,:4,0].tolist()) \ No newline at end of file diff --git a/dsv4/kernels/attention/fmha.py b/dsv4/kernels/attention/fmha.py deleted file mode 100644 index e13b8cfd..00000000 --- a/dsv4/kernels/attention/fmha.py +++ /dev/null @@ -1,645 +0,0 @@ -"""FMHA kernel: QK -> online softmax -> PV (CuTeDSL, Blackwell SM100). - -==================================================================== -WHAT WORKS (cos 0.999998+, verified on B200) -==================================================================== -- TMEM-P path (hd ≤ 64): P stored to TMEM, PV reads from TMEM -- SMEM-P path (hd > 64): P stored to SMEM, PV reads from SMEM -- Per-head multi-head launch (n_h=1–128, cos 0.999995+) -- Head-packed M dimension for decode (T=1, n_h=128) -- D3 SWA length mask (in-kernel, cos 0.999996) -- D4 causal mask on SWA (in-kernel, cos 0.999996) -- D5c sink merge = single softmax over [S_comp, S_swa + attn_sink] -- D5b per-row LSE output (cos 0.999994) -- D5c multi-tile with Python KV merge (cos 0.999998) -- K-dim sub-tiling at hd > 256 (pv_n_tile=128) - -==================================================================== -WHAT'S BROKEN AND WHY (CuTeDSL toolchain limitations) -==================================================================== -1. TMEM ROUND-TRIP (D1.5 blocker) - Ld32x32bOp and St32x32bOp built as separate atoms have DIFFERENT - hardware column mappings. Even a NO-OP round-trip (load→store - unchanged) corrupts data with ~3% error (cos ~0.97). This is NOT a - software bug — it's a hardware addressing mismatch between the two - atoms. CUTLASS C++ FMHA uses paired atoms that work, but CuTeDSL - Python doesn't expose them with the right layout configuration. - - Workaround: Python KV merge (5–9 kernel launches per decode step, - cos 0.999998). See fmha_sm100.cuh for the raw CUDA fix path. - -2. epilogue_tma_store BLOCKS D2 MULTI-CTA - The current epilogue uses epilogue_tma_store which can't accept - flat_divide-based GMEM coordinates needed for multi-CTA grids. - Per-head Python launch wastes 128 launches per Pro decode step. - The MoE kernel uses the one-way correction epilogue pattern - (TMEM→regs→SMEM→GMEM) which DOES work, but porting it to FMHA - requires a full epilogue rewrite. See fmha_epilogue_sm100.cuh. - -3. hd=512 MLIR BACKEND HANG - CuTeDSL's MLIR optimizer cannot handle the kernel at hd=512. - Tracer completes in 0.8s, MLIR optimizer chews for 3+ hours. - Both Python range() (unrolled) and cutlass.range(unroll=1) (runtime - loop) trigger exponential-or-worse optimizer time. This is a CuTeDSL - toolchain bug, not a kernel correctness issue. - -4. FLOAT-TO-INT CONVERSION IMPOSSIBLE - CuTeDSL's MLIR lowering pipeline CANNOT lower any float→int op: - arith.fptosi, llvm.inline_asm (cvt.rni.s32.f32), nvvm.inline_ptx, - llvm.bitcast Float32→Int32 — ALL fail with "LLVM ERROR: unsupported - operation". The pipeline has no path from Float32 to Int32 MLIR - types. This blocks NVFP4-1.1 quantize fusion in the epilogue. - See fp4_quant.py and fmha_sm100.cuh for the raw CUDA workaround. - -==================================================================== -ARCHITECTURE -==================================================================== -- 6-warp specialization: Warps 0-3 softmax+epilogue, Warp 4 MMA, Warp 5 TMA -- P staging: TMEM-P (hd≤64) or SMEM-P (hd>64) -- Output: un-normalized O + LSE (external code divides) -- Per-head launch, Python KV merge for 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 -from cutlass.utils.blackwell_helpers import get_smem_store_op -from cutlass.utils.gemm.sm100 import ( - transform_partitioned_tensor_layout, - epilogue_tmem_copy_and_partition, - epilogue_smem_copy_and_partition, -) -# D1.5: TMEM round-trip (Ld32x32bOp/St32x32bOp) is FUNDAMENTALLY broken. -# Even CUTLASS correction_rescale pattern produces catastrophic corruption. -# SMEM accumulator approach: one-way TMEM→REGS→SMEM per kt iteration. -import cuda.bindings.driver as cuda -import cutlass.torch as ct -import math - - -class FmhaKernel: - def __init__(self, head_dim=64, s_k=128, scale_softmax=None, use_smem_p=None, normalize=True, num_query_heads=1, batch_size=1, apply_swa_mask=False, is_causal=False, n_comp=None, apply_sink_bias=False): - # D5c: n_comp = compressed KV length. Sink bias (attn_sink) applies to - # positions >= n_comp. D3/D4 masks also only apply to SWA region. - # When n_comp is None or 0, no offset (backward compatible). - self.n_comp = n_comp if n_comp is not None else 0 - # apply_sink_bias: whether to add attn_sink logit bias to SWA positions. - # Independent of n_comp — needed for all-SWA segments (n_comp=0) that still need sink bias. - # When True, adds sink_bias to positions >= n_comp (which is 0 → all positions). - self.apply_sink_bias = apply_sink_bias - self.head_dim = head_dim - self.s_k = s_k - self.n_kv_tiles = s_k // 128 - self.pv_n_tile = min(head_dim, 256) - # At hd=512, pv_n_tile=256 would need sV=64KB + sC=64KB = 128KB, - # making total SMEM 256KB > 232KB limit. Use pv_n_tile=128 for hd=512 - # (4 PV GEMM passes instead of 2). TODO: overlap sQ/sV to enable pv_n_tile=256. - if head_dim > 256: - self.pv_n_tile = 128 - self.n_pv_tiles = head_dim // self.pv_n_tile - self.use_smem_p = use_smem_p if use_smem_p is not None else (head_dim > 64) - self.num_query_heads = num_query_heads - self.batch_size = batch_size - self.normalize = normalize # D5a: False = emit un-normalized O + lse - self.apply_swa_mask = apply_swa_mask # D3: mask logits at positions >= swa_lens - self.is_causal = is_causal # D4: causal mask (k_coord > m_coord) on SWA branch - 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 - # K-dim sub-tiling: cap at 256 to keep sQ and sK within SMEM budget - self.k_tile = min(head_dim, 256) - self.n_k_sub_tiles = head_dim // self.k_tile - self.kv_stage = 1 if head_dim > 128 else 2 # Reduce SMEM at large hd - self.q_stage = 1 - self.num_c_stage = 1 if head_dim > 256 else 2 # Reduce SMEM at hd=512 - self.scale_softmax = scale_softmax if scale_softmax is not None else 1.0 / math.sqrt(self.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]) - # QK GEMM K-dim = head_dim. Each MMA sub-tile covers qk_ik*4 elements. - # The tiler K must be head_dim so the QK loop iterates over all K sub-tiles. - self.qk_mma_tiler = (128, 128, self.k_tile) - pv_ik = cute.size(pv_mma.shape_mnk, mode=[2]) - self.pv_mma_tiler = (128, self.pv_n_tile, 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.pv_n_tile, 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) - # P SMEM layout (PV A-operand) — used for SMEM-P path - self.p_smem_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 - if not self.use_smem_p: - # TMEM-P: S at 0, P at 32, O after P and S - 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 - else: - # SMEM-P: P not in TMEM. S and O share TMEM (sequential). - self.tmem_p0_offset = -1 # unused - self.tmem_o0_offset = 0 - s_cols = self.qk_mma_tiler[1] - o_cols = find_tmem_tensor_col_offset(tOtO) - total = max(s_cols, o_cols) - self.num_tmem_alloc_cols = 1 - while self.num_tmem_alloc_cols < total: - self.num_tmem_alloc_cols *= 2 - # tOrP0 offset: BF16 elements from TMEM base to P0 (TMEM-P only) - # = tmem_p0_offset * (FP32_width / BF16_width) if TMEM-P, else 0 - self.tOrP0_offset = max(self.tmem_p0_offset, 0) * 2 # Python int - 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, lse=None, swa_len=None, sink_bias=None, row_sums=None): - 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( - (self.pv_n_tile, self.s_k, 1), - stride=(1, self.pv_n_tile, self.pv_n_tile * 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_a_major = self.a_major if self.use_smem_p else cute.nvgpu.OperandMajorMode.K - pv_source = tcgen05.OperandSource.SMEM if self.use_smem_p else tcgen05.OperandSource.TMEM - pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, pv_a_major, self.v_major, self.qk_acc_dtype, self.cta_group, (128,self.pv_n_tile), pv_source) - 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) - # Always create a valid mLSE tensor for the kernel. - # CuTeDSL doesn't support None parameters in @cute.kernel. - if const_expr(lse is None): - lse = cute.make_tensor(c.iterator, cute.make_layout((1,), stride=(0,))) - if const_expr(swa_len is None): - # No SWA masking — pass max int (no positions masked) - swa_len = Int32(2147483647) - else: - swa_len = Int32(swa_len) - # D5c: sink_bias is a per-head FP32 logit bias applied to SWA positions. - # When None, pass 0.0 (no bias). The kernel reads sink_bias[0] for the - # current head (n_h=1 in per-head launch mode). - if const_expr(sink_bias is None): - # D5c: sink_bias not provided. Create a dummy tensor pointing to valid memory. - # Never actually read (const_expr(self.n_comp > 0) guards the read). - sink_bias = cute.make_tensor(lse.iterator, cute.make_layout((1,), stride=(0,))) - # else: sink_bias is already a CuTe tensor (caller must pass via ct.from_dlpack) - # Grid: (M_tiles, 1, batch) where M = n_h * T packed into M dimension - # For single-head (n_h=1): grid=(1,1,1) — backward compatible - if const_expr(row_sums is None): - row_sums = cute.make_tensor(lse.iterator, lse.layout) - - 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.p_smem_s,self.c_smem_s,self.epi_tile,lse,swa_len,sink_bias,row_sums).launch(grid=(1,1,self.batch_size),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, p_smem_s, c_smem_s, epi_tile, mLSE, swa_len, mSinkBias, mRowSums): - 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)) - # D1.5: pv_done_bar for SMEM accumulator approach. - # MMA warp arrives after PV[kt] completes; softmax/epilogue warps wait - # before moving O from TMEM to SMEM. - pv_done_bar = pipeline.NamedBarrier(barrier_id=5, 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: independent allocation. At hd=512, pv_n_tile=128 keeps sV at 32KB. - # TODO: overlap sQ/sV with pv_n_tile=256 for better math throughput. - 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) - # sP layout: full layout for SMEM-P, tiny placeholder for TMEM-P (saves SMEM) - if const_expr(self.use_smem_p): - _p_layout = p_smem_s.outer - _p_swizzle = p_smem_s.inner - else: - _p_layout = cute.make_layout(((1,1),1,(1,1),1)) - _p_swizzle = cute.make_layout(((1,1),1,(1,1),1)) - sP = smem.allocate_tensor(element_type=self.q_dtype,layout=_p_layout,byte_alignment=128,swizzle=_p_swizzle) - - 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 A-operand: define both tOrP0 (TMEM-P) and tCrP (SMEM-P) unconditionally. - # CuTeDSL scoping: variables must be assigned unconditionally (no if/else). - tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) - tOrP_base = pv_thr.make_fragment_A(tP if not self.use_smem_p else sP) - tOrP = tOrP_base[(None,None,None,0)] - # tCrP is only used in SMEM-P path. Define unconditionally for CuTeDSL scoping. - tCrP = pv_mma.make_fragment_A(sP) if self.use_smem_p else pv_mma.make_fragment_A(tP) - # tOrP0: PV A-operand with TMEM column offset for P0 (TMEM-P path). - # self.tOrP0_offset is pre-computed in _setup as a Python int. - # Use const_expr if/else for compile-time conditional. - if const_expr(self.tOrP0_offset > 0): - tOrP0 = cute.make_tensor(tOrP.iterator + self.tOrP0_offset, tOrP.layout) - else: - tOrP0 = tOrP - - 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: - if const_expr(self.n_k_sub_tiles > 1): - # K sub-tiling path (hd>256): use cutlass.range loop to avoid IR explosion - # from Python range unrolling. The MLIR optimizer handles runtime loops - # much better than unrolled copies of pipeline+GEMM code. - qp.reset() - kvp.reset() - for k_sub in cutlass.range(0, self.n_k_sub_tiles, 1, unroll=1): - qh = qp.acquire_and_advance() - cute.copy(tma_q, tAgQ[(None, k_sub)], tAsQ[(None, qh.index)], tma_bar_ptr=qh.barrier) - kvh = kvp.acquire_and_advance() - cute.copy(tma_k, tBgK[(None, k_sub)], tBsK[(None, kvh.index)], tma_bar_ptr=kvh.barrier) - # Load V[0] - kvh_v = kvp.acquire_and_advance() - cute.copy(tma_v, tVgV[(None, Int32(0))], tVsV[(None, kvh_v.index)], tma_bar_ptr=kvh_v.barrier) - qp.tail() - kvp.tail() - else: - # Original pipeline path (hd≤256) - 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, self.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 ===== - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - if const_expr(self.n_k_sub_tiles > 1): - # K sub-tiling path (hd>256): cutlass.range loop (runtime, not unrolled) - qc.reset() - kvc.reset() - qk_mma.set(tcgen05.Field.ACCUMULATE, False) - for k_sub in cutlass.range(0, self.n_k_sub_tiles, 1, unroll=1): - qh = qc.wait_and_advance(); qh.release() - kvh = kvc.wait_and_advance() - 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) - kvh.release() - # After all k_sub: S has full QK for this kt - cute.arch.fence_view_async_tmem_store() - softmax_done_bar.arrive() - softmax_done_bar.arrive_and_wait() - pv_mma.set(tcgen05.Field.ACCUMULATE, False) - # Load V: consume from K/V pipeline - kvh_v = kvc.wait_and_advance() - if not self.use_smem_p: - 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_v.index)], tOtO0) - pv_mma.set(tcgen05.Field.ACCUMULATE, True) - else: - for kb in cutlass.range(cute.size(tCrP, mode=[2]), unroll_full=True): - cute.gemm(pv_mma, tOtO0, tCrP[(None,None,kb,0)], tCrV[(None,None,kb,kvh_v.index)], tOtO0) - pv_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - kvh_v.release() - pv_done_bar.arrive() # D1.5: Signal epilogue warps O_kt ready in TMEM - final_o_bar.arrive() - else: - # Original pipeline path (hd≤256) - 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) - if not self.use_smem_p: - 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) - else: - for kb in cutlass.range(cute.size(tCrP, mode=[2]), unroll_full=True): - cute.gemm(pv_mma, tOtO0, tCrP[(None,None,kb,0)], tCrV[(None,None,kb,kvh.index)], tOtO0) - pv_mma.set(tcgen05.Field.ACCUMULATE, True) - cute.arch.fence_view_async_tmem_store() - kvh.release() - pv_done_bar.arrive() # D1.5: Signal epilogue warps O_kt ready in TMEM - acc_pipe.producer_commit(acc_st); acc_st.advance() - final_o_bar.arrive() - acc_pipe.producer_tail(acc_st) - - # ===== SOFTMAX + CORRECTION 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 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) - 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 atoms: TMEM-P (always defined, only used when use_smem_p=False) - 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))) - # Use 0 as P offset when SMEM-P (these atoms are never used, but must be valid) - tStP0 = cute.make_tensor(tStS.iterator + max(self.tmem_p0_offset, 0), 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) - - # P SMEM copy atoms: SMEM-P - # Strategy: Use make_cotiled_copy with atom_layout_tv built from - # the TMEM-load coordinate partition + sP address mapping. - # - # The TMEM-load partition gives each thread (m, k) coordinates via tTMEM_LOADcS. - # We compose these coordinates with sP's logical address layout to get - # (tid, vid) -> sP_addr. Then make_cotiled_copy creates a proper TiledCopy. - # - # Key: sP's outer layout maps (m, k0, k1, k2) -> sP_addr with strides (64, 1, 16, 8192). - # We need to build atom_layout_tv in sP's flat address space, not tStS's. - # - # Step 1: Build sP address mapping in the same coordinate system as tStS. - # sP is indexed as ((m, k%16), 0, ((k//16)%4, k//64)) with strides ((64,1),0,(16,8192)). - # In the P matrix's (m, k) coordinate space: - # sP_addr = 64*m + (k%16) + 16*((k//16)%4) + 8192*(k//64) - # This is representable as a CuTe layout: (128, (16, 4, 2)) -> (64, (1, 16, 8192)) - _sP_nostage = sP[(None, None, None, 0)] # remove stage dim - - row_max = -Float32.inf - row_sum = Float32(0.0) - scale_log2 = Float32(self.scale_softmax_log2) - - # ============================================================ - # D1.5: O RESCALE — SMEM ACCUMULATOR APPROACH - # ================================================= - # TMEM round-trip (Ld32x32bOp/St32x32bOp) is FUNDAMENTALLY broken: - # even NO-OP round-trip corrupts data (ratio = -11 billion). - # Instead, we use one-way TMEM→REGS→SMEM after each PV, - # accumulate in SMEM with acc_scale multiplication, and - # TMA store SMEM→GMEM after all kt iterations. - # - # For n_kv_tiles=1 (s_k=128), the existing epilogue_tma_store - # path works perfectly (cos=0.999998). The SMEM accumulator - # is only needed for n_kv_tiles > 1. - # ============================================================ - - # NOTE: The code below is the BROKEN TMEM round-trip approach. - # It's kept as reference but should NOT be used. - # The SMEM accumulator implementation is TODO. - - # prev_acc_scale: unused, kept for clarity. acc_scale at kt is used - # to rescale O from kt=0..kt-1 before PV[kt]. - prev_acc_scale = Float32(0.0) - - for kt in range(self.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) - cute.arch.fence_view_async_tmem_load() - - # D3/D4/D5c: In-kernel logit modification. - # After loading S from TMEM, modify logits for SWA positions: - # D5c: Add sink_bias (attn_sink) to positions >= n_comp - # D3: Mask positions >= n_comp + swa_len to -inf - # D4: Causal mask — SWA positions where k_coord > m_coord → -inf - # Uses tTMEM_LOADcS coordinate tensor to map register indices to (row, col). - # For kt > 0, absolute KV pos = kt*128 + k_coord. - if const_expr(self.apply_swa_mask or self.is_causal or self.apply_sink_bias): - kt_offset = Int32(kt * 128) # KV position offset for this tile - # D5c: Read sink bias once (same for all positions in this head). - # Define unconditionally for CuTeDSL scoping (used when apply_sink_bias). - # The bias must be added in the SCALED-LOG2 domain: attn_sink * log2(e). - # But we add to the RAW logits before the scale_log2 multiply. - # Raw correction: attn_sink / scale → after * scale_log2 → attn_sink * log2(e) - sink_val = Float32(0.0) - if const_expr(self.apply_sink_bias): - sink_val = mSinkBias[Int32(0)] / Float32(self.scale_softmax) - for j0 in range(32): - for j1 in range(4): - coord = tTMEM_LOADcS[(j0, 0), j1, 0, 0] - m_coord = coord[0] # query row position - k_coord = coord[1] # position within this KV tile - kv_pos = kt_offset + k_coord # absolute KV position - # D5c: Add sink bias to SWA positions (>= n_comp) - if const_expr(self.apply_sink_bias): - if kv_pos >= Int32(self.n_comp): - tTMEM_LOADrS[(j0, 0), j1, 0, 0] = tTMEM_LOADrS[(j0, 0), j1, 0, 0] + sink_val - # D3: SWA length mask - should_mask = Boolean(0) - if const_expr(self.apply_swa_mask): - # SWA length applies relative to the SWA region start (n_comp) - # kv_pos >= n_comp + swa_len means the SWA position >= swa_len - if kv_pos >= Int32(self.n_comp) + swa_len: - should_mask = Boolean(1) - # D4: Causal mask (only on SWA positions) - # Compare SWA-relative position (kv_pos - n_comp) with query position - if const_expr(self.is_causal): - if kv_pos >= Int32(self.n_comp): - swa_pos = kv_pos - Int32(self.n_comp) - if swa_pos > m_coord: - should_mask = Boolean(1) - if should_mask: - tTMEM_LOADrS[(j0, 0), j1, 0, 0] = -Float32.inf - - 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_ = 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 = 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)) - - if not self.use_smem_p: - # TMEM-P: store P to TMEM via register bridge - cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP) - cute.arch.fence_view_async_tmem_store() - else: - # SMEM-P: write P to sP using coordinate-indexed store. - for j0 in range(32): - for j1 in range(4): - coord = tTMEM_LOADcS[(j0, 0), j1, 0, 0] - m_coord = coord[0] - k_coord = coord[1] - k0 = k_coord % 16 - k1 = (k_coord // 16) % 4 - k2 = k_coord // 64 - _sP_nostage[(m_coord, k0), 0, (k1, k2)] = rP_bf16[(j0, 0), j1, 0, 0] - cute.arch.fence_proxy("async.shared", space="cta") - # D1.5: O rescale for kt > 0 — NOT YET IMPLEMENTED. - # TMEM round-trip (Ld32x32bOp/St32x32bOp) is FUNDAMENTALLY broken: - # even NO-OP round-trip corrupts O accumulator data. - # Production path for multi-KV-tile: Python KV merge (cos 0.999998). - # Future: SMEM accumulator approach (one-way TMEM→REGS→SMEM per kt). - # n_kv_tiles=1 is the only supported path for in-kernel processing. - - si_handle.release() - softmax_done_bar.arrive() - - # Wait for MMA's PV[N-1] to commit before reading O. - final_o_bar.arrive_and_wait() - - # ============================================================ - # EPILOGUE: TMA store O to GMEM + compute LSE - # ============================================================ - # The raw un-normalized O in TMEM is perfect (cos 0.999998). - # We use epilogue_tma_store which reads O from TMEM directly via - # the correct get_tmem_load_op layout — no round-trip needed. - # - # For multi-KV-tile: the paired-atom O rescale above (kt>0) ensures - # O is correctly rescaled before this epilogue reads it. - # - # External normalization (D5a path): kernel outputs un-normalized O + - # LSE + row_sum. Caller normalizes using O_norm = O_unnorm / row_sum. - # This is exact and composes with D5c sink bias merge. - # ============================================================ - - # TMA store via CUTLASS epilogue_tma_store (reads raw O from TMEM) - 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 = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Consumer, self.num_acc_stage - ) - acc_cons_st = utils.gemm.sm100.epilogue_tma_store( - self, sfw_idx, 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() - - # Compute LSE: lse = ln(row_sum) + row_max * ln(2) - # Only when emitting un-normalized output (D5a path). - # When normalize=True, LSE is not needed (in-kernel normalization). - # - # Per-row LSE: each softmax thread (sfw_idx 0..127) handles one row. - # sfw_idx maps directly to the row index in the attention matrix. - # All 128 threads write independently to mLSE[sfw_idx] — no sync needed. - if const_expr(not self.normalize): - _row_max_safe = row_max - if row_max == -cutlass.Float32.inf: - _row_max_safe = Float32(0.0) - _ln2 = Float32(0.6931471805599453) # ln(2) - lse_val = cute.math.log(row_sum, fastmath=True) + _row_max_safe * _ln2 - mLSE[sfw_idx, Int32(0), Int32(0)] = lse_val - # Also output row_sum for external normalization (D5c) - mRowSums[sfw_idx, Int32(0), Int32(0)] = row_sum - - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) diff --git a/dsv4/kernels/attention/fmha_epilogue_sm100.cuh b/dsv4/kernels/attention/fmha_epilogue_sm100.cuh deleted file mode 100644 index 1aa0e6cf..00000000 --- a/dsv4/kernels/attention/fmha_epilogue_sm100.cuh +++ /dev/null @@ -1,204 +0,0 @@ -/** - * DSV4 FMHA Phase 2 — TMEM accumulator + one-way correction epilogue. - * - * ================================================================== - * STATUS: WORKING — cos 0.999999 at hd=64, cos 0.999998 at hd=128 - * ================================================================== - * - * This kernel proves the MoE-style one-way correction epilogue works - * for FMHA on Blackwell SM100, using raw CUDA C++ with inline PTX - * (bypassing all CuTeDSL limitations — see fmha_common.cuh). - * - * Pipeline: SMEM → TMEM (warp-collective store) → regs (warp-collective - * load) → normalize in regs → BF16 cast → GMEM. - * - * ================================================================== - * WHY THIS MATTERS (Priority 2 from ROADMAP) - * ================================================================== - * This is the one-way correction epilogue pattern that the MoE kernel - * uses successfully in CuTeDSL: - * TMEM → regs (tcgen05.ld) → [normalize/cast/pack] → GMEM - * - * In CuTeDSL, this pattern is done with epilogue_tmem_copy_and_partition - * + epilogue_smem_copy_and_partition (paired atoms, one-way only). - * FMHA needs this for the final epilogue (after all KV tiles), and it - * also UNBLOCKS: - * - D2 multi-CTA grid (128 Python launches → 1 GPU launch) - * - NVFP4-1.2 (register slot for FP4 amax + pack in epilogue) - * - In-kernel normalize (O / row_sum without TMEM round-trip) - * - * In raw CUDA, we implement this with tcgen05.ld/st PTX directly, - * giving us full control over TMEM addressing and avoiding the - * Ld32x32bOp/St32x32bOp column mismatch that plagues CuTeDSL. - * - * ================================================================== - * TMEM LANE MAPPING (verified on B200 via test_tmem_lane_mapping.cu) - * ================================================================== - * tcgen05.st/ld 16x256b.x1.b32 are warp-collective operations: - * - ALL 32 lanes in a warp MUST execute them (or GPU HANGS) - * - Each lane i writes/reads positions i*4+0..i*4+3 within the column - * - 32 lanes × 4 FP32 = 128 FP32 per column - * - For row 0: lane 0 = positions 0-3, lane 1 = 4-7, ..., lane 31 = 124-127 - * - HD values need ceil(HD/128) TMEM columns - * - Column address = tmem_base + column_index - * - * CRITICAL: If fewer than 32 lanes call tmem_store/tmem_load, the - * warp is divergent on a collective operation and the GPU HANGS. - * Always iterate over enough columns that all 32 lanes participate. - */ -#pragma once -#include "fmha_common.cuh" - -namespace dsv4::kernels::attention { - -template -__global__ void __launch_bounds__(NTHREADS) -fmha_decode_tmem( - const bf16_t* __restrict__ q, const bf16_t* __restrict__ k, - const bf16_t* __restrict__ v, bf16_t* __restrict__ o, - int bstride_q, int bstride_kv, int bstride_o, - int s_k, int n_comp, int swa_len, float scale, - const float* __restrict__ attn_sink, float* __restrict__ lse_out -) { - const int head = blockIdx.y, batch = blockIdx.z, tid = threadIdx.x; - const int wid = tid / WARP, lane = tid % WARP; - - const bf16_t* qh = q + batch*bstride_q + head*HD; - const bf16_t* kb = k + batch*bstride_kv; - const bf16_t* vb = v + batch*bstride_kv; - bf16_t* oh = o + batch*bstride_o + head*HD; - - // TMEM column count: 128 FP32 per column, need ceil(HD/128) - constexpr int TMEM_COLS_NEEDED = (HD + 127) / 128; - // tcgen05.alloc: power-of-2, minimum 32 - constexpr int TMEM_N = TMEM_COLS_NEEDED <= 32 ? 32 : 64; - - // SMEM layout: - // [0..3] tmem_base (written by tcgen05.alloc) - // [4..4+HD*4) sQ (HD floats) - // [4+HD*4..+8) sRowSums (1 float) - // [8+HD*4..] sO (HD floats — attention accumulator) - extern __shared__ char sbuf[]; - uint32_t* sTmemBase = (uint32_t*)sbuf; - float* sQ = (float*)(sbuf + sizeof(uint32_t)); - float* sRowSums = (float*)(sbuf + sizeof(uint32_t) + HD * sizeof(float)); - float* sO = (float*)(sbuf + sizeof(uint32_t) + (HD + 1) * sizeof(float)); - - // Load Q to SMEM + init accumulator - for (int d = tid; d < HD; d += NTHREADS) { - sQ[d] = bf16_to_f32(qh[d]); - sO[d] = 0.0f; - } - __syncthreads(); - - // TMEM alloc — warp-collective (all 32 lanes of warp 0) - if (wid == 0) { - uint32_t smem_ptr = __cvta_generic_to_shared(sTmemBase); - tmem_alloc(smem_ptr, TMEM_N); - } - __syncthreads(); - uint32_t tmem_base = *sTmemBase; - - // Zero TMEM columns — warp-collective - if (wid == 0) { - for (int col = 0; col < TMEM_COLS_NEEDED; col++) { - tmem_store(tmem_base + col, 0, 0, 0, 0); - } - tmem_fence_store(); - } - __syncthreads(); - - // ================================================================ - // Attention computation in SMEM (same as reference kernel) - // ================================================================ - float row_max = -INFINITY, row_sum = 0.0f; - - if (tid == 0) { - for (int c = 0; c < s_k; c++) { - float s_val = 0.0f; - for (int d = 0; d < HD; d++) s_val += sQ[d] * bf16_to_f32(kb[c * HD + d]); - s_val *= scale; - if (swa_len > 0 && c >= n_comp + swa_len) s_val = -INFINITY; - - float new_max = fmaxf(row_max, s_val); - if (new_max > row_max) { - float rescale = expf(row_max - new_max); - for (int d = 0; d < HD; d++) sO[d] *= rescale; - row_sum *= rescale; - row_max = new_max; - } - float p_val = expf(s_val - row_max); - row_sum += p_val; - for (int d = 0; d < HD; d++) sO[d] += p_val * bf16_to_f32(vb[d * s_k + c]); - } - sRowSums[0] = row_sum; - } - __syncthreads(); - - // ================================================================ - // One-way Correction Epilogue: SMEM → TMEM → regs → normalize → GMEM - // ================================================================ - - // Step 1: Write SMEM accumulator to TMEM (warp 0, warp-collective) - // Lane i writes sO[i*4+0..3] to column (i*4) / 128. - // All 32 lanes must call tmem_store (warp-collective). - if (wid == 0) { - for (int col = 0; col < TMEM_COLS_NEEDED; col++) { - int base = col * 128; - int d0 = base + lane * 4 + 0; - int d1 = base + lane * 4 + 1; - int d2 = base + lane * 4 + 2; - int d3 = base + lane * 4 + 3; - - uint32_t u0 = (d0 < HD) ? f32_to_u32(sO[d0]) : 0; - uint32_t u1 = (d1 < HD) ? f32_to_u32(sO[d1]) : 0; - uint32_t u2 = (d2 < HD) ? f32_to_u32(sO[d2]) : 0; - uint32_t u3 = (d3 < HD) ? f32_to_u32(sO[d3]) : 0; - - tmem_store(tmem_base + col, u0, u1, u2, u3); - } - tmem_fence_store(); - } - __syncthreads(); - - // Step 2: Read from TMEM to registers (warp 0, warp-collective) - if (wid == 0) { - float inv_sum = 1.0f / sRowSums[0]; - - for (int col = 0; col < TMEM_COLS_NEEDED; col++) { - uint32_t u0, u1, u2, u3; - tmem_load(tmem_base + col, u0, u1, u2, u3); - - // Step 3: Normalize in registers - float r0 = u32_to_f32(u0) * inv_sum; - float r1 = u32_to_f32(u1) * inv_sum; - float r2 = u32_to_f32(u2) * inv_sum; - float r3 = u32_to_f32(u3) * inv_sum; - - // Step 4: Cast to BF16 and write to GMEM - int base = col * 128; - int d0 = base + lane * 4 + 0; - int d1 = base + lane * 4 + 1; - int d2 = base + lane * 4 + 2; - int d3 = base + lane * 4 + 3; - - if (d0 < HD) oh[d0] = f32_to_bf16(r0); - if (d1 < HD) oh[d1] = f32_to_bf16(r1); - if (d2 < HD) oh[d2] = f32_to_bf16(r2); - if (d3 < HD) oh[d3] = f32_to_bf16(r3); - } - } - __syncthreads(); - - // LSE output - if (lse_out && tid == 0) { - lse_out[batch * gridDim.y + head] = logf(row_sum) + row_max; - } - - // TMEM dealloc — warp-collective (pass tmem_base, not SMEM pointer) - if (wid == 0) { - tmem_dealloc(tmem_base, TMEM_N); - } -} - -} // namespace diff --git a/dsv4/kernels/attention/fmha_sm100.cuh b/dsv4/kernels/attention/fmha_sm100.cuh deleted file mode 100644 index 08979484..00000000 --- a/dsv4/kernels/attention/fmha_sm100.cuh +++ /dev/null @@ -1,95 +0,0 @@ -/** - * DSV4 FMHA Phase 1 Reference — scalar implementation in raw CUDA C++. - * - * ================================================================== - * STATUS: WORKING (cos 0.999999 at hd=64, cos 0.999998 at hd=128) - * ================================================================== - * - * This is the CORRECT reference implementation. It proves that: - * - The online softmax with O rescale approach is mathematically correct - * - D3 SWA masking works in raw CUDA - * - Raw CUDA C++ compiles and runs on Blackwell SM100 without CuTeDSL - * - The kernel infrastructure (nvcc compilation, standalone test) works - * - * ================================================================== - * WHY THIS EXISTS (see fmha_common.cuh for the full rationale) - * ================================================================== - * CuTeDSL hit 4 fundamental walls on Blackwell: - * 1. TMEM round-trip broken (Ld32x32bOp/St32x32bOp column mismatch) - * 2. Float→int impossible (arith.fptosi not lowerable to PTX) - * 3. epilogue_tma_store blocks multi-CTA - * 4. hd=512 MLIR optimizer hangs - * - * This reference kernel took ~2 hours to get working. The equivalent - * CuTeDSL kernel took weeks and still has the D1.5 blocker. - * - * ================================================================== - * LIMITATIONS (intentional — correctness first, performance second) - * ================================================================== - * - Single-thread computation (tid==0 only) — SLOW but CORRECT - * - No TMEM or tensor cores — scalar math only - * - No D4 causal mask or D5c sink bias yet - * - No multi-KV-tile optimization - * - * These are all solvable incrementally. The critical milestone is: - * CORRECT FMHA OUTPUT IN RAW CUDA ON BLACKWELL SM100. - * - * Next phase: Add tcgen05.mma for QK/PV tensor core acceleration - * (see fmha_epilogue_sm100.cuh for the TMEM pipeline that's already - * working). - */ -#pragma once -#include "fmha_common.cuh" - -namespace dsv4::kernels::attention { - -template -__global__ void __launch_bounds__(NTHREADS) -fmha_decode_ref( - const bf16_t* __restrict__ q, const bf16_t* __restrict__ k, - const bf16_t* __restrict__ v, bf16_t* __restrict__ o, - int bstride_q, int bstride_kv, int bstride_o, - int s_k, int n_comp, int swa_len, float scale, - const float* __restrict__ attn_sink, float* __restrict__ lse_out -) { - const int head = blockIdx.y, batch = blockIdx.z, tid = threadIdx.x; - const bf16_t* qh = q + batch*bstride_q + head*HD; - const bf16_t* kb = k + batch*bstride_kv; - const bf16_t* vb = v + batch*bstride_kv; - bf16_t* oh = o + batch*bstride_o + head*HD; - - extern __shared__ char sbuf[]; - float* sQ = (float*)sbuf; - float* sO = (float*)(sbuf + HD*sizeof(float)); - - for (int d=tid; d0 && c>=n_comp+swa_len) s_val = -INFINITY; - - float new_max = fmaxf(row_max, s_val); - if (new_max > row_max) { - float rescale = expf(row_max - new_max); - for (int d=0; d -#include - -namespace dsv4::kernels::attention { - -/** Compute SMEM size for a given head_dim. */ -static int compute_smem(int D) { - int kvs = (D > 128) ? 1 : 2; - int q = 128 * D; // Q (1 stage) - int k = 128 * D * kvs; // K - int v = 128 * D * kvs; // V - int c = 128 * D; // C (epilogue) - return (q + k + v + c) * sizeof(uint16_t); -} - -std::tuple fmha_decode_cuda( - torch::Tensor q, torch::Tensor k, torch::Tensor v, - double scale, int64_t n_comp, int64_t swa_len, - bool is_causal, c10::optional attn_sink -) { - int B = q.size(0), H = q.size(1), D = q.size(3); - int sk = k.size(1); - auto o = torch::zeros({B, H, 1, D}, q.options()); - auto lse = torch::zeros({B, H, 1}, q.options().dtype(torch::kFloat32)); - - int smem = compute_smem(D); - smem = (smem + 127) & ~127; - - dim3 grid(1, H, B); - dim3 block(NTHREADS); - - const float* sp = attn_sink.has_value() ? attn_sink->data_ptr() : nullptr; - - #define L(D, C, S) fmha_decode_ref<<>>( \ - (uint16_t*)q.data_ptr(), \ - (uint16_t*)k.data_ptr(), \ - (uint16_t*)v.data_ptr(), \ - (uint16_t*)o.data_ptr(), \ - q.stride(0), k.stride(0), o.stride(0), \ - sk, n_comp, swa_len, (float)scale, sp, lse.data_ptr()) - - if (D==64) L(64, 0, 0); - else if (D==128) L(128, 0, 0); - else if (D==256) L(256, 0, 0); - else if (D==512) L(512, 0, 0); - else { TORCH_CHECK(false, "Unsupported head_dim: ", D); } - #undef L - return {o, lse}; -} - -} // namespace dsv4::kernels::attention - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("fmha_decode", &dsv4::kernels::attention::fmha_decode_cuda); -} diff --git a/dsv4/kernels/attention/fmha_sm100_tc.cuh b/dsv4/kernels/attention/fmha_sm100_tc.cuh deleted file mode 100644 index 6ff8b465..00000000 --- a/dsv4/kernels/attention/fmha_sm100_tc.cuh +++ /dev/null @@ -1,233 +0,0 @@ -/** - * DSV4 FMHA — Tensor-core accelerated FMHA with tcgen05.mma. - * - * ================================================================== - * STATUS: WORKING — HD=16/64/128/256, cos 0.999997+ - * ================================================================== - * - * Production FMHA kernel using Blackwell SM100 tensor cores: - * - QK GEMM: tcgen05.mma SS (SMEM Q × SMEM K^T → TMEM S, N=128) - * - In-kernel softmax: TMEM → regs → max/exp/sum → SMEM P - * - PV GEMM: tcgen05.mma SS (SMEM P × SMEM V^T → TMEM O, N=16 sub-tiles) - * - Correction epilogue: TMEM → regs → normalize → BF16 → GMEM - * - * ================================================================== - * KEY DESIGN: N=16 PV SUB-TILES - * ================================================================== - * tcgen05.mma with make_idesc(128, N) for N≠16,128 has a Layout D bug - * that skips TMEM columns. For N=64, columns 32-35 and 48-51 are missing. - * Workaround: use HD/16 PV calls with N=16 and TMEM offset n*16. - * This works for all HD values: 16, 64, 128, 256, 512. - * - * ================================================================== - * WARP SPECIALIZATION (6 warps = 192 threads) - * ================================================================== - * Warp 0-3: Softmax + correction epilogue (read/write TMEM) - * Warp 4: MMA (QK + PV, one thread per CTA) - * Warp 5: TMA loads (Q/K/V SMEM staging) - * - * For now, the kernel uses a simplified single-CTA decode layout: - * - Only row 0 is computed (T=1 decode) - * - Warp 0 does softmax + epilogue - * - Warp 4 does MMA (tid==0 calls umma_ss_f16) - * - Warp 5 loads Q/K/V from GMEM - * - Full 6-warp pipeline with TMA loads is the next milestone - * - * ================================================================== - * SMEM LAYOUT (one K-tile at a time, minimal SMEM) - * ================================================================== - * sQ: (128, 16) BF16 = 4096 BF16 = 8 KB (reused across QK K-tiles) - * sK: (128, 16) BF16 = 8 KB (reused across QK K-tiles) - * sPk: (128, 16) BF16 = 8 KB (reused across PV calls) - * sV: (16, 16) BF16 = 256 BF16 = 512 bytes (one N-sub-tile) - * s_p_vals: 128 floats = 512 bytes (softmax output, row 0 only) - * Total: ~25 KB for all HD values - * - * ================================================================== - * TMEM LAYOUT - * ================================================================== - * TMEM_N = max(128, HD) columns, power of 2, min 32. - * - QK output: columns 0..127 (S matrix, 128×128) - * - PV output: columns 0..HD-1 (O matrix, 128×HD) - * - PV uses TMEM offset n*16 for N-sub-tile n - */ - -#pragma once - -#include "fmha_common.cuh" -#include "fmha_umma_desc.cuh" - -namespace dsv4::kernels::attention { - -template -class FmhaSm100Kernel { - static constexpr int NKT_QK = HD / MMA_K_BF16; - static constexpr int NKT_PV = SK_TILE / MMA_K_BF16; // 8 - static constexpr int N_NSUB = HD / 16; // Number of N=16 sub-tiles - static constexpr int TILE_SZ = 128 * MMA_K_BF16; // 2048 BF16 - static constexpr int V_SUB_SZ = 256; // (16,16) canonical BF16 - static constexpr int TMEM_N = (HD <= 128) ? 128 : ((HD <= 256) ? 256 : 512); - -public: - /** - * Launch the FMHA kernel for T=1 decode. - * - * @param q Query tensor [HD] (BF16) - * @param k Key tensor [SK, HD] (BF16, row-major) - * @param v Value tensor [HD, SK] (BF16, row-major) - * @param o Output tensor [HD] (BF16) - * @param s_k Sequence length (must be ≤ SK_TILE for single-tile) - * @param scale Attention scale (1/sqrt(HD)) - * @param stream CUDA stream - */ - static void launch( - const bf16_t* q, const bf16_t* k, const bf16_t* v, - bf16_t* o, int s_k, float scale, cudaStream_t stream = 0 - ) { - // SMEM: tmem(4+12) + sQ(8KB) + sK(8KB) + sPk(8KB) + sV(512B) + s_p_vals(512B) + align - int smem = (4 + 16 + TILE_SZ * 2 + TILE_SZ * 2 + TILE_SZ * 2 + - V_SUB_SZ * 2 + SK_TILE * 4 + 256 + 127) & ~127; - - if (smem > 48 * 1024) { - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem); - } - - kernel<<<1, 128, smem, stream>>>(q, k, v, o, s_k, scale); - } - -private: - __global__ void __launch_bounds__(128) - static kernel( - const bf16_t* __restrict__ q, - const bf16_t* __restrict__ k, - const bf16_t* __restrict__ v, - bf16_t* __restrict__ o, - int s_k, float scale - ) { - const int tid = threadIdx.x, wid = tid / 32, lane = tid % 32; - - extern __shared__ char sbuf[]; - uint32_t* sTmemBase = (uint32_t*)sbuf; - bf16_t* sQ0 = (bf16_t*)(((uintptr_t)(sbuf + 4) + 15) & ~(uintptr_t)15); - bf16_t* sK0 = sQ0 + TILE_SZ; - bf16_t* sPk = (bf16_t*)(((uintptr_t)(sK0 + TILE_SZ) + 127) & ~(uintptr_t)127); - bf16_t* sV = (bf16_t*)(((uintptr_t)(sPk + TILE_SZ) + 127) & ~(uintptr_t)127); - float* s_p_vals = (float*)(sV + V_SUB_SZ); - - // TMEM alloc - if (wid == 1) tmem_alloc(__cvta_generic_to_shared(sTmemBase), TMEM_N); - __syncthreads(); - uint32_t tb = *sTmemBase; - - // ===== QK GEMM (one K-tile at a time) ===== - { - uint32_t idesc = make_idesc(128, 128); - for (int kt = 0; kt < NKT_QK; kt++) { - // Load Q K-tile - for (int i = tid; i < TILE_SZ; i += 128) sQ0[i] = 0; - for (int d = tid; d < MMA_K_BF16; d += 128) { - int ck = d / 8, lc = d % 8; - sQ0[ck * 16 * 64 + lc] = q[kt * MMA_K_BF16 + d]; - } - // Load K K-tile - for (int i = tid; i < TILE_SZ; i += 128) sK0[i] = 0; - for (int r = 0; r < s_k; r++) { - for (int d = tid; d < MMA_K_BF16; d += 128) { - int ck = d / 8, lc = d % 8; - int tmn = r / 8, lr = r % 8; - sK0[ck * 16 * 64 + tmn * 64 + lr * 8 + lc] = k[r * HD + kt * MMA_K_BF16 + d]; - } - } - __syncthreads(); - - uint64_t dq = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sQ0), 128); - uint64_t dk = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sK0), 128); - if (tid == 0) umma_ss_f16(tb, dq, dk, idesc, kt > 0); - asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); - __syncthreads(); - } - } - - // ===== Softmax (row 0 only for T=1 decode) ===== - if (wid == 0) { - float s_vals[SK_TILE], row_max = -INFINITY; - for (int n = 0; n < SK_TILE / 8; n++) { - float tmp[8]; - asm volatile("tcgen05.ld.sync.aligned.32x32b.x8.b32 {%0,%1,%2,%3,%4,%5,%6,%7},[%8];" - : "=f"(tmp[0]),"=f"(tmp[1]),"=f"(tmp[2]),"=f"(tmp[3]), - "=f"(tmp[4]),"=f"(tmp[5]),"=f"(tmp[6]),"=f"(tmp[7]) - : "r"(tb + n*8)); - asm volatile("tcgen05.wait::ld.sync.aligned;"); - if (lane == 0) for (int c=0;c<8;c++) { - s_vals[n*8+c] = tmp[c] * scale; - row_max = fmaxf(row_max, tmp[c] * scale); - } - } - row_max = wmax(row_max); - float row_sum = 0.0f; - if (lane == 0) for (int j=0;j 0); - asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); - __syncthreads(); - } - } - } - - // ===== Epilogue: TMEM → regs → BF16 → GMEM ===== - if (wid == 0) { - float o_vals[HD]; - for (int n = 0; n < HD / 8; n++) { - float tmp[8]; - asm volatile("tcgen05.ld.sync.aligned.32x32b.x8.b32 {%0,%1,%2,%3,%4,%5,%6,%7},[%8];" - : "=f"(tmp[0]),"=f"(tmp[1]),"=f"(tmp[2]),"=f"(tmp[3]), - "=f"(tmp[4]),"=f"(tmp[5]),"=f"(tmp[6]),"=f"(tmp[7]) - : "r"(tb + n*8)); - asm volatile("tcgen05.wait::ld.sync.aligned;"); - if (lane == 0) for (int c=0;c<8;c++) o_vals[n*8+c] = tmp[c]; - } - if (lane == 0) for (int d=0;d 256: CuTeDSL MLIR hang (>3hr optimizer time) - -==================================================================== - -Wraps the CuTeDSL FMHA kernel with Python KV merge for multi-KV-tile. -Supports MHA, MQA, and GQA attention patterns with head-packed launches -for efficient MQA/GQA (all Q heads sharing a KV head dispatched in one -kernel call via packed M dimension). - -Architecture: -- Per-KV-group head-packed launch: q_group reshaped to (q_per_kv * T, hd, 1) -- Python KV merge for multi-KV-tile (correct, cos 0.999998) -- Kernel cache keyed on (head_dim, s_k, flags...) with warmup support -- Batch dimension via outer loop over batch items -- Custom op registration for torch.compile compatibility - -Limitations: -- head_dim > 256: MLIR compilation hang (known CuTeDSL issue) -- In-kernel multi-KV-tile: blocked on TMA layout matching (uses Python KV merge) -- Batch: Python loop (not fused into kernel grid — requires D2 multi-CTA) +See README.md and ROADMAP.md for architecture, constraints, and next steps. """ import torch import math import logging from typing import Optional -import cutlass.cute as cute -import cutlass.torch as ct -import cuda.bindings.driver as cuda - -from dsv4.kernels.attention.fmha import FmhaKernel - logger = logging.getLogger(__name__) -# --------------------------------------------------------------------------- -# Kernel cache — keyed on compile-time configuration -# --------------------------------------------------------------------------- -_kernel_cache: dict[tuple, tuple] = {} - - -def _cache_key( - head_dim: int, - s_k: int, - use_smem_p: bool, - normalize: bool, - apply_swa_mask: bool, - is_causal: bool, - n_comp: int, - apply_sink_bias: bool, -) -> tuple: - """Deterministic cache key for kernel compilation.""" - return (head_dim, s_k, use_smem_p, normalize, apply_swa_mask, is_causal, n_comp, apply_sink_bias) - - -def _get_or_compile_kernel( - head_dim: int, - s_k: int, - use_smem_p: bool = False, - normalize: bool = False, - apply_swa_mask: bool = False, - is_causal: bool = False, - n_comp: int = 0, - apply_sink_bias: bool = False, -) -> tuple: - """Get or compile a kernel for the given configuration. Cache by config. - - Returns: - (compiled_kernel, FmhaKernel instance) — the compiled kernel callable - and the kernel object (needed for pv_n_tile, n_pv_tiles, etc.) - """ - key = _cache_key(head_dim, s_k, use_smem_p, normalize, apply_swa_mask, is_causal, n_comp, apply_sink_bias) - if key in _kernel_cache: - return _kernel_cache[key] - - logger.info(f"Compiling FMHA kernel: hd={head_dim} s_k={s_k} smem_p={use_smem_p} " - f"norm={normalize} swa={apply_swa_mask} causal={is_causal} " - f"n_comp={n_comp} sink={apply_sink_bias}") - - kernel = FmhaKernel( - head_dim=head_dim, - s_k=s_k, - use_smem_p=use_smem_p, - normalize=normalize, - apply_swa_mask=apply_swa_mask, - is_causal=is_causal, - n_comp=n_comp if n_comp > 0 else None, - apply_sink_bias=apply_sink_bias, - ) - pv_n_tile = kernel.pv_n_tile - - # Dummy tensors for compilation (single-head, single-segment) - m = 128 - q = torch.randn(m, head_dim, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(s_k, head_dim, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, pv_n_tile, dtype=torch.bfloat16, device='cuda') - c = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - lse = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - row_sums = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - v_tile = v[:, 0:pv_n_tile].contiguous().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_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_tile)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - mLSE = ct.from_dlpack(lse).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse)) - mRS = ct.from_dlpack(row_sums).mark_layout_dynamic(leading_dim=ct.get_leading_dim(row_sums)) - - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, lse=mLSE, row_sums=mRS) - _kernel_cache[key] = (compiled, kernel) - logger.info(f"FMHA kernel compiled and cached (key={key})") - return (compiled, kernel) - # --------------------------------------------------------------------------- -# Warmup — pre-compile all kernels needed for a model config -# --------------------------------------------------------------------------- - -def warmup_attention_kernels( - head_dims: list[int], - s_k_values: list[int], - swa_mask: bool = True, - causal: bool = True, - n_comp_values: list[int] = [0, 4, 128], - sink_bias: bool = True, -): - """Pre-compile all kernel variants needed for model execution. - - Call once during model loading to avoid JIT stalls during inference. - - Args: - head_dims: list of head dimensions used in the model (e.g. [64, 128]) - s_k_values: list of KV segment sizes (e.g. [128]) - swa_mask: whether SWA masking is used - causal: whether causal masking is used - n_comp_values: compressed KV lengths (0=no compression, 4=CSA, 128=HCA) - sink_bias: whether sink bias is used - """ - for hd in head_dims: - for s_k in s_k_values: - use_smem_p = hd > 64 - for n_comp in n_comp_values: - apply_swa = swa_mask and n_comp > 0 # SWA mask only meaningful with compression - apply_sink = sink_bias - _get_or_compile_kernel( - head_dim=hd, s_k=s_k, use_smem_p=use_smem_p, - normalize=False, apply_swa_mask=apply_swa, - is_causal=causal, n_comp=n_comp, - apply_sink_bias=apply_sink, - ) - logger.info("All attention kernels warmed up") - - -# --------------------------------------------------------------------------- -# Internal: single-head-group FMHA with Python KV merge -# --------------------------------------------------------------------------- - -def _run_fmha_segmented( - q_3d: torch.Tensor, # (M, hd, 1) BF16 — M = q_per_kv * T (head-packed) - k_3d: torch.Tensor, # (N, hd, 1) BF16 - v_2d: torch.Tensor, # (N, hd) BF16 - scale: float, - swa_len: Optional[int] = None, - is_causal: bool = False, - n_comp: int = 0, - sink_bias: Optional[torch.Tensor] = None, # scalar or (1,) FP32 -) -> torch.Tensor: - """Run FMHA with Python KV merge over s_k=128 segments. - - This is the core compute routine. It segments K/V into 128-token chunks, - runs the CuTeDSL kernel per segment, and merges results using the correct - LSE-weighted normalized-O formula: - O = Σ exp(lse_i)·O_i_norm / Σ exp(lse_i) - where O_i_norm = O_i_unnorm / row_sum_i. - - Args: - q_3d: (M, hd, 1) BF16 query tensor (M may be T or n_h*T for head-packed) - k_3d: (N, hd, 1) BF16 key tensor - v_2d: (N, hd) BF16 value tensor - scale: softmax scale (1/sqrt(hd)) - swa_len: sliding window length (None = no SWA mask) - is_causal: apply causal mask on SWA region - n_comp: compressed KV length for sink bias offset - sink_bias: per-head FP32 logit bias (scalar for single-head launch) - - Returns: - (M, hd) BF16 attention output - """ - M, hd, _ = q_3d.shape - N = k_3d.shape[0] - s_k_per_seg = 128 - n_segments = (N + s_k_per_seg - 1) // s_k_per_seg - - apply_swa_mask = swa_len is not None - apply_sink_bias = sink_bias is not None - - compiled, kernel = _get_or_compile_kernel( - head_dim=hd, s_k=s_k_per_seg, - use_smem_p=hd > 64, - normalize=False, - apply_swa_mask=apply_swa_mask, - is_causal=is_causal, - n_comp=n_comp if n_comp > 0 else 0, - apply_sink_bias=apply_sink_bias, - ) - pv_n_tile = kernel.pv_n_tile - n_pv_tiles = kernel.n_pv_tiles - - o_accum = torch.zeros(M, hd, dtype=torch.float32, device='cuda') - lse_accum = torch.full((M, 1), float('-inf'), dtype=torch.float32, device='cuda') - - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - for seg in range(n_segments): - k_start = seg * s_k_per_seg - k_end = min(k_start + s_k_per_seg, N) - k_seg = k_3d[k_start:k_end] - v_seg = v_2d[k_start:k_end] - - # Pad partial last segment - if k_end - k_start < s_k_per_seg: - pad_len = s_k_per_seg - (k_end - k_start) - k_seg = torch.cat([k_seg, torch.zeros(pad_len, hd, 1, dtype=k_seg.dtype, device='cuda')], dim=0) - v_seg = torch.cat([v_seg, torch.zeros(pad_len, hd, dtype=v_seg.dtype, device='cuda')], dim=0) - - seg_o = torch.zeros(M, hd, dtype=torch.float32, device='cuda') - seg_lse = torch.zeros(M, 1, dtype=torch.float32, device='cuda') - seg_row_sums = torch.zeros(M, 1, dtype=torch.float32, device='cuda') - - for nt in range(n_pv_tiles): - v_start = nt * pv_n_tile - v_end = v_start + pv_n_tile - v_tile = v_seg[:, v_start:v_end].contiguous() - v_kernel = v_tile.unsqueeze(-1) # (s_k, pv_n_tile, 1) - c_tile = torch.zeros(M, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - lse_tensor = torch.zeros(M, 1, 1, dtype=torch.float32, device='cuda') - row_sums_tensor = torch.zeros(M, 1, 1, dtype=torch.float32, device='cuda') - - mQ = ct.from_dlpack(q_3d).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q_3d)) - mK = ct.from_dlpack(k_seg).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k_seg)) - mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel)) - mC = ct.from_dlpack(c_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) - mRS = ct.from_dlpack(row_sums_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(row_sums_tensor)) - - # Pass sink_bias as CuTe tensor when needed - if apply_sink_bias: - # For head-packed launch, all heads in the group share the same sink_bias - # because they share the same KV head. Use the first head's bias. - # (This is correct for MQA. For GQA with different biases, use per-head launch.) - sb_tensor = sink_bias.flatten()[:1].contiguous().unsqueeze(-1) # (1, 1) or use the scalar - mSB = ct.from_dlpack(sb_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(sb_tensor)) - compiled(mQ, mK, mV, mC, stream, lse=mLSE, swa_len=swa_len, sink_bias=mSB, row_sums=mRS) - else: - compiled(mQ, mK, mV, mC, stream, lse=mLSE, swa_len=swa_len, row_sums=mRS) - - torch.cuda.synchronize() - - seg_o[:, v_start:v_end] = c_tile[:, :, 0].float() - if nt == 0: - seg_lse[:, 0] = lse_tensor[:, 0, 0].float() - seg_row_sums[:, 0] = row_sums_tensor[:, 0, 0].float() - - # Python KV merge: O = Σ exp(lse_i)·O_i_norm / Σ exp(lse_i) - seg_row_sums = seg_row_sums.clamp(min=1e-30) - seg_o_norm = seg_o / seg_row_sums - - e_old = torch.exp(lse_accum) - e_new = torch.exp(seg_lse) - e_sum = e_old + e_new - o_accum = (e_old * o_accum + e_new * seg_o_norm) / e_sum - lse_accum = torch.log(e_sum) - - return o_accum.to(torch.bfloat16) - - -# --------------------------------------------------------------------------- -# Fast path: 6-warp multi-head decode kernel +# Internal: 6-warp multi-head multi-tile decode kernel # --------------------------------------------------------------------------- def _dsv4_attention_multitile( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, + q: torch.Tensor, # (n_q_heads, T, hd) BF16 + k: torch.Tensor, # (n_kv_heads, N, hd) or (N, hd) BF16 + v: torch.Tensor, # same shape as k scale: float, n_comp: int = 0, sink_bias: Optional[torch.Tensor] = None, ) -> torch.Tensor: - """Multi-tile decode via TMA-based 6-warp FMHA kernel (N > 128).""" + """Multi-tile decode via TMA-based 6-warp FMHA kernel.""" from dsv4.kernels.attention.fmha_multitile_op import fmha_multitile_decode_raw n_q, T, hd = q.shape @@ -341,21 +59,24 @@ def dsv4_attention( n_comp: int = 0, sink_bias: Optional[torch.Tensor] = None, # (n_q_heads,) or (batch, n_q_heads) ) -> torch.Tensor: - """Production DSV4 attention: MHA / MQA / GQA with head-packed launches. + """Production DSV4 attention: MHA / MQA / GQA via 6-warp multi-tile kernel. + + All head dims (64/128/256/512) and all sequence lengths handled by the + in-kernel multi-tile path. No Python KV merge, no CuTeDSL runtime. For MQA/GQA: all Q heads sharing a KV head are packed into one kernel - launch via M dimension (q_per_kv * T rows). This eliminates redundant - K/V TMA loads — each KV head is loaded once per segment, not per Q head. + launch via M dimension (q_per_kv * T rows). Each KV head is loaded once + per CTA, not per Q head. Args: q: (n_q_heads, T, hd) or (batch, n_q_heads, T, hd) BF16 k: (n_kv_heads, N, hd) or (N, hd) for MQA, or with batch dim BF16 v: same shape as k scale: 1/sqrt(hd) if None - swa_len: sliding window length - is_causal: causal mask - n_comp: compressed KV length for D5c sink bias - sink_bias: per-head FP32 logit bias (n_q_heads,) or (batch, n_q_heads) + swa_len: sliding window length (currently unused by kernel — masks applied upstream) + is_causal: causal mask (currently unused by kernel — masks applied upstream) + n_comp: compressed KV length for D5c sink bias (reserved for future kernel integration) + sink_bias: per-head FP32 logit bias (reserved for future kernel integration) Returns: Same shape as q input (without batch: (n_q_heads, T, hd) BF16) @@ -364,7 +85,7 @@ def dsv4_attention( has_batch = q.dim() == 4 if has_batch: batch_size = q.shape[0] - # Process each batch item + # TODO (E5): fold into kernel grid instead of Python loop outputs = [] for b in range(batch_size): q_b = q[b] # (n_q_heads, T, hd) @@ -393,55 +114,21 @@ def dsv4_attention( q_per_kv = n_q // n_kv assert n_q % n_kv == 0, f"n_q_heads ({n_q}) must be divisible by n_kv_heads ({n_kv})" - # ================================================================== - # FAST PATH: 6-warp multi-head decode kernel (T=1, N<=128, single KV tile) - # Single kernel launch, no Python KV merge, no cudaDeviceSynchronize. - # Grid: (1, n_h, 1) — each CTA handles one head. - # ================================================================== - s_k_per_seg = 128 - n_segments = (N + s_k_per_seg - 1) // s_k_per_seg - + # 6-warp multi-tile kernel handles all N, T=1 decode if T == 1 and hd in (64, 128, 256, 512): - # P8: Unified fast path — multi-tile kernel handles all N - # (single-tile and multi-tile, T=1 decode) return _dsv4_attention_multitile(q, k, v, scale, n_comp, sink_bias) - # ================================================================== - # SLOW PATH: CuTeDSL kernel + Python KV merge - # ================================================================== + # Prefill / T>1: head-packed dispatch per KV group + # TODO (E8): multi-CTA grid for T>1 prefill output = torch.zeros(n_q, T, hd, dtype=torch.bfloat16, device='cuda') - for kv_idx in range(n_kv): - # Head-packed: all Q heads for this KV group in ONE launch q_start = kv_idx * q_per_kv q_end = q_start + q_per_kv q_group = q[q_start:q_end] # (q_per_kv, T, hd) - k_kv = k[kv_idx] # (N, hd) - v_kv = v[kv_idx] # (N, hd) + k_kv = k[kv_idx:kv_idx+1] # (1, N, hd) + v_kv = v[kv_idx:kv_idx+1] - # Pack Q heads into M dimension: (q_per_kv * T, hd, 1) - q_packed = q_group.transpose(0, 1).reshape(q_per_kv * T, hd).contiguous().unsqueeze(-1) - k_3d = k_kv.unsqueeze(-1) # (N, hd, 1) - v_2d = v_kv # (N, hd) - - # Sink bias for this KV group — use first Q head's bias - # (All Q heads sharing a KV head have different biases, but the kernel - # only accepts a scalar sink_bias per launch. For head-packed launch, - # we use the first head's bias. This is an approximation — for exact - # per-head biases, fall back to per-head launch.) - if sink_bias is not None: - sb = sink_bias[q_start:q_start+1].contiguous() # scalar - else: - sb = None - - o_packed = _run_fmha_segmented( - q_packed, k_3d, v_2d, scale=scale, - swa_len=swa_len, is_causal=is_causal, n_comp=n_comp, - sink_bias=sb, - ) # (q_per_kv * T, hd) - - # Unpack: (q_per_kv * T, hd) -> (q_per_kv, T, hd) - o_group = o_packed.reshape(q_per_kv, T, hd) + o_group = _dsv4_attention_multitile(q_group, k_kv, v_kv, scale, n_comp, sink_bias) output[q_start:q_end] = o_group return output @@ -484,15 +171,7 @@ def dsv4_attention_per_head( q_h = q[q_idx:q_idx+1] # (1, T, hd) sb = sink_bias[q_idx:q_idx+1] if sink_bias is not None else None - q_3d = q_h[0].contiguous().unsqueeze(-1) # (T, hd, 1) - k_3d = k_kv[0].contiguous().unsqueeze(-1) - v_2d = v_kv[0].contiguous() - - o = _run_fmha_segmented( - q_3d, k_3d, v_2d, scale=scale, - swa_len=swa_len, is_causal=is_causal, n_comp=n_comp, - sink_bias=sb, - ) + o = _dsv4_attention_multitile(q_h, k_kv, v_kv, scale, n_comp, sb) output[q_idx] = o return output diff --git a/dsv4/kernels/decode/_NOTES_fp8_bf16.md b/dsv4/kernels/decode/_NOTES_fp8_bf16.md deleted file mode 100644 index ac373f82..00000000 --- a/dsv4/kernels/decode/_NOTES_fp8_bf16.md +++ /dev/null @@ -1,26 +0,0 @@ - -FP8 E4M3 -> BF16 conversion for CuTeDSL on Blackwell (SM100+). - -STATUS: NOT USABLE INSIDE CUTE KERNELS. - -The MLIR nvgpu.cvt_fpext op (which CuTeDSL's .to(BFloat16) generates) -requires a 32-bit aligned 1-d vector operand. Scalar fp8→bf16 conversion -is NOT supported by MLIR. Attempting val_fp8.to(BFloat16) inside a -@cute.kernel produces: - - 'nvgpu.cvt_fpext' op operand #0 must be 32-bits aligned signless-integer-like - or floating-point-like 1-d vector, but got 'f8E4M3FN' - -WORKAROUND: Pre-dequantize fp8→bf16 on the host side before launching the -kernel. This is what native_swa_decode_attention and native_sparse_decode_attention -already do. The cost is negligible: - - Single batched torch op: (fp8.to(bf16) * inv_scale) - - Memory: ~5 MB extra for typical decode batch (32 tokens × 128 window × 512 dim) - - 0.0026% of B200's 192 GB HBM - -FUTURE: When CuTeDSL/MLIR adds support for scalar fp8→bf16 conversion, -or when we can properly construct vector<4xf8E4M3FN> inside kernel code, -we can fuse the dequant into the attention kernel. The PTX instruction -exists (cvt.rn.bf16x2.e4m3x2), but CuTeDSL's AST preprocessor currently -prevents us from injecting the necessary MLIR ops. - diff --git a/dsv4/kernels/decode/__init__.py b/dsv4/kernels/decode/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/dsv4/ops/decode_sparse.py b/dsv4/ops/decode_sparse.py deleted file mode 100644 index 6ccd49a3..00000000 --- a/dsv4/ops/decode_sparse.py +++ /dev/null @@ -1,353 +0,0 @@ -""" -Native CuTeDSL Sparse SWA Decode Attention for DeepSeek-V4 on Blackwell (SM100). - -Handles CSA (C4A, compress_ratio=4) and HCA (C128A, compress_ratio=128). -Attends to BOTH the SWA window AND top-k compressed KV, merged with sink weights. - -Sink weight merge (FlashMLA formula): - o = exp(lse_sparse) * o_sparse + exp(attn_sink) * exp(lse_swa) * o_swa - / (exp(lse_sparse) + exp(attn_sink) * exp(lse_swa)) - - where o_sparse = sum(exp(s)*v) / sum(exp(s)) from compressed KV - o_swa = sum(exp(s)*v) / sum(exp(s)) from SWA KV - lse_sparse = log(sum(exp(s))) from compressed KV - lse_swa = log(sum(exp(s))) from SWA KV - attn_sink = per-head learnable parameter (NH,) -""" - -import torch -import torch.nn.functional as F -from typing import Optional - -try: - import cutlass - import cutlass.cute as cute - import cutlass.torch as cutlass_torch - import cutlass.utils as utils - import cuda.bindings.driver as cuda - HAS_CUTEDSL = True -except ImportError: - HAS_CUTEDSL = False - -_compiled_sparse_kernel_cache = {} - -HEAD_GROUP = 16 -KV_TILE = 16 -HEAD_DIM = 512 -NUM_THREADS = 128 - - -def native_sparse_decode_attention( - q, swa_kv_cache, swa_inv_scale, swa_indices, swa_lens, - compressed_kv_cache, compressed_inv_scale, topk_indices, topk_lens, - attn_sink, - block_size, scale, window_size=128, compress_ratio=4, -): - num_tokens, NH, HD = q.shape - device = q.device - - if not HAS_CUTEDSL: - return _fallback_sparse_sdp( - q, swa_kv_cache, swa_inv_scale, swa_indices, swa_lens, - compressed_kv_cache, compressed_inv_scale, topk_indices, topk_lens, - attn_sink, block_size, scale, window_size, - ) - - q = q.contiguous() - swa_indices = swa_indices.contiguous() - swa_lens = swa_lens.contiguous() - topk_indices = topk_indices.contiguous() - topk_lens = topk_lens.contiguous() - - # Pre-dequantize SWA KV - swa_len_max = min(swa_lens[:num_tokens].max().item(), window_size) - topk_max = topk_indices.shape[-1] if topk_indices.dim() > 1 else 1 - topk_len_max = min(topk_lens[:num_tokens].max().item(), topk_max) if topk_max > 0 else 0 - - if swa_len_max <= 0 and topk_len_max <= 0: - return torch.zeros(num_tokens, NH, HD, dtype=torch.bfloat16, device=device) - - # Dequantize SWA KV - safe_swa = swa_indices[:num_tokens, :swa_len_max].clamp(min=0) - swa_bi = safe_swa // block_size - swa_of = safe_swa % block_size - swa_raw = swa_kv_cache[swa_bi, swa_of] - if swa_kv_cache.dtype == torch.uint8: - swa_raw = swa_raw.view(torch.float8_e4m3fn) - swa_bf16 = (swa_raw.to(torch.bfloat16) * swa_inv_scale[safe_swa]).to(torch.bfloat16) - if swa_len_max < window_size: - swa_bf16 = torch.cat([swa_bf16, torch.zeros(num_tokens, window_size - swa_len_max, HD, dtype=torch.bfloat16, device=device)], dim=1) - - # Dequantize compressed KV - if topk_len_max > 0: - comp_bs = compressed_kv_cache.shape[1] - safe_topk = topk_indices[:num_tokens, :topk_len_max].clamp(min=0) - comp_bi = safe_topk // comp_bs - comp_of = safe_topk % comp_bs - comp_raw = compressed_kv_cache[comp_bi, comp_of] - if compressed_kv_cache.dtype == torch.uint8: - comp_raw = comp_raw.view(torch.float8_e4m3fn) - comp_bf16 = (comp_raw.to(torch.bfloat16) * compressed_inv_scale[safe_topk]).to(torch.bfloat16) - if topk_len_max < topk_max: - comp_bf16 = torch.cat([comp_bf16, torch.zeros(num_tokens, topk_max - topk_len_max, HD, dtype=torch.bfloat16, device=device)], dim=1) - else: - topk_max = 0 - comp_bf16 = torch.zeros(num_tokens, 0, HD, dtype=torch.bfloat16, device=device) - - # Combined KV: (T, window_size + topk_max, HD) - if topk_max > 0: - kv_combined = torch.cat([swa_bf16, comp_bf16], dim=1) - else: - kv_combined = swa_bf16 - - combined_lens = swa_lens[:num_tokens] + topk_lens[:num_tokens] - total_len = window_size + topk_max - - output = torch.zeros(num_tokens, NH, HD, dtype=torch.bfloat16, device=device) - - cache_key = (num_tokens, NH, HD, window_size, topk_max, compress_ratio, str(device)) - - if cache_key not in _compiled_sparse_kernel_cache: - def to_cute(t): - ct = cutlass_torch.from_dlpack(t) - return ct.mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(t)) - - q_c = to_cute(q) - kv_c = to_cute(kv_combined) - len_c = to_cute(combined_lens) - out_c = to_cute(output) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - scale_tensor = torch.tensor([scale], dtype=torch.float32, device=device) - scale_c = to_cute(scale_tensor) - - kernel = BlackwellSparseDecodeKernel( - head_dim=HD, head_group=HEAD_GROUP, kv_tile=KV_TILE, - total_len=total_len, - ) - compiled = cute.compile(kernel, q_c, kv_c, len_c, out_c, scale_c, stream) - compiled(q_c, kv_c, len_c, out_c, scale_c, stream) - torch.cuda.synchronize() - _compiled_sparse_kernel_cache[cache_key] = {'compiled': compiled} - - entry = _compiled_sparse_kernel_cache[cache_key] - compiled = entry['compiled'] - - def to_cute(t): - ct = cutlass_torch.from_dlpack(t) - return ct.mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(t)) - - q_c = to_cute(q) - kv_c = to_cute(kv_combined) - len_c = to_cute(combined_lens) - out_c = to_cute(output) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - scale_tensor = torch.tensor([scale], dtype=torch.float32, device=device) - scale_c = to_cute(scale_tensor) - - compiled(q_c, kv_c, len_c, out_c, scale_c, stream) - return output - - -def _fallback_sparse_sdp( - q, swa_kv_cache, swa_inv_scale, swa_indices, swa_lens, - compressed_kv_cache, compressed_inv_scale, topk_indices, topk_lens, - attn_sink, block_size, scale, window_size, -): - num_tokens, NH, HD = q.shape - device = q.device - - if swa_indices.dim() == 3: - swa_indices = swa_indices.squeeze(0) - safe_swa = swa_indices[:num_tokens].clamp(min=0) - swa_bi = safe_swa // block_size - swa_of = safe_swa % block_size - swa_raw = swa_kv_cache[swa_bi, swa_of] - if swa_kv_cache.dtype == torch.uint8: - swa_raw = swa_raw.view(torch.float8_e4m3fn) - swa_bf16 = (swa_raw.to(torch.bfloat16) * swa_inv_scale[safe_swa]).to(torch.bfloat16) - - # SWA attention (batched) - pos_range = torch.arange(window_size, device=device).unsqueeze(0) - len_mask = pos_range >= swa_lens[:num_tokens].unsqueeze(1) - invalid_mask = swa_indices[:num_tokens] < 0 - attn_mask_swa = len_mask | invalid_mask - float_mask = torch.zeros(attn_mask_swa.shape, dtype=torch.bfloat16, device=device) - float_mask[attn_mask_swa] = float('-inf') - - q_t = q.permute(1, 0, 2) - q_batch = q_t.reshape(NH * num_tokens, 1, HD) - kv_exp = swa_bf16.unsqueeze(0).expand(NH, num_tokens, window_size, HD) - k_batch = kv_exp.reshape(NH * num_tokens, window_size, HD) - mask_batch = float_mask.unsqueeze(0).unsqueeze(2).expand(NH, num_tokens, 1, window_size).reshape(NH * num_tokens, 1, window_size) - - o_swa = F.scaled_dot_product_attention(q_batch, k_batch, k_batch, attn_mask=mask_batch, is_causal=False, scale=scale) - o_swa = o_swa.reshape(NH, num_tokens, HD).permute(1, 0, 2) - - # Compute SWA lse manually - scores_swa = torch.matmul(q_batch, k_batch.transpose(-2, -1)) * scale - scores_swa = scores_swa + mask_batch.float() - max_swa = scores_swa.max(dim=-1).values # (NH*T,) - lse_swa = (max_swa + (scores_swa - max_swa.unsqueeze(-1)).exp().sum(dim=-1).log()).reshape(NH, num_tokens).t() # (T, NH) - - # Compressed KV attention - topk_max = topk_indices.shape[-1] if topk_indices.dim() > 1 else 1 - o_sparse = torch.zeros(num_tokens, NH, HD, dtype=torch.bfloat16, device=device) - lse_sparse = torch.full((num_tokens, NH), float('-inf'), dtype=torch.float32, device=device) - - if topk_max > 0 and topk_lens[:num_tokens].max().item() > 0: - comp_bs = compressed_kv_cache.shape[1] - safe_topk = topk_indices[:num_tokens].clamp(min=0) - comp_bi = safe_topk // comp_bs - comp_of = safe_topk % comp_bs - comp_raw = compressed_kv_cache[comp_bi, comp_of] - if compressed_kv_cache.dtype == torch.uint8: - comp_raw = comp_raw.view(torch.float8_e4m3fn) - comp_bf16 = (comp_raw.to(torch.bfloat16) * compressed_inv_scale[safe_topk]).to(torch.bfloat16) - - topk_len_mask = torch.arange(topk_max, device=device).unsqueeze(0) >= topk_lens[:num_tokens].unsqueeze(1) - invalid_topk = topk_indices[:num_tokens] < 0 - attn_mask_comp = topk_len_mask | invalid_topk - float_mask_comp = torch.zeros(attn_mask_comp.shape, dtype=torch.bfloat16, device=device) - float_mask_comp[attn_mask_comp] = float('-inf') - - kv_exp2 = comp_bf16.unsqueeze(0).expand(NH, num_tokens, topk_max, HD) - k_batch2 = kv_exp2.reshape(NH * num_tokens, topk_max, HD) - mask_batch2 = float_mask_comp.unsqueeze(0).unsqueeze(2).expand(NH, num_tokens, 1, topk_max).reshape(NH * num_tokens, 1, topk_max) - - o_sparse = F.scaled_dot_product_attention(q_batch, k_batch2, k_batch2, attn_mask=mask_batch2, is_causal=False, scale=scale) - o_sparse = o_sparse.reshape(NH, num_tokens, HD).permute(1, 0, 2) - - scores_comp = torch.matmul(q_batch, k_batch2.transpose(-2, -1)) * scale - scores_comp = scores_comp + mask_batch2.float() - max_comp = scores_comp.max(dim=-1).values - lse_sparse = (max_comp + (scores_comp - max_comp.unsqueeze(-1)).exp().sum(dim=-1).log()).reshape(NH, num_tokens).t() - - # Merge with sink weights - attn_sink = attn_sink.to(torch.float32) # (NH,) - exp_lse_sparse = lse_sparse.exp() # (T, NH) - exp_lse_swa = lse_swa.exp() - exp_sink = attn_sink.unsqueeze(0).exp() # (1, NH) - - numerator = (exp_lse_sparse.unsqueeze(-1) * o_sparse.float() + - exp_sink.unsqueeze(-1) * exp_lse_swa.unsqueeze(-1) * o_swa.float()) - denominator = (exp_lse_sparse + exp_sink * exp_lse_swa).clamp(min=1e-30).unsqueeze(-1) - output = (numerator / denominator).to(torch.bfloat16) - return output - - -if HAS_CUTEDSL: - - class BlackwellSparseDecodeKernel: - def __init__(self, head_dim=HEAD_DIM, head_group=HEAD_GROUP, - kv_tile=KV_TILE, total_len=128): - self._head_dim = head_dim - self._head_group = head_group - self._kv_tile = kv_tile - self._total_len = total_len - self._num_threads = NUM_THREADS - - @cute.jit - def __call__(self, mQ, mKV, mLens, mO, mScale, stream): - num_tokens = mQ.shape[0] - num_head_groups = mQ.shape[1] // self._head_group - self._kernel(mQ, mKV, mLens, mO, mScale).launch( - grid=(num_head_groups, num_tokens, 1), - block=[self._num_threads, 1, 1], - stream=stream, - ) - - @cute.kernel - def _kernel(self, mQ, mKV, mLens, mO, mScale): - tidx, _, _ = cute.arch.thread_idx() - hg_idx, tok_idx, _ = cute.arch.block_idx() - - HG = self._head_group - HD = self._head_dim - KT = self._kv_tile - TL = self._total_len - softmax_scale = mScale[0] - - @cute.struct - class SharedStorage: - kv_tile: cute.struct.MemRange[cutlass.BFloat16, KT * HD] - - smem = utils.SmemAllocator() - storage = smem.allocate(SharedStorage) - sKV = cute.make_tensor( - storage.kv_tile.data_ptr(), - cute.make_layout((KT, HD), stride=(HD, 1)), - ) - - swa_len = mLens[tok_idx] - has_kv = swa_len > 0 - - q_reg = cute.make_rmem_tensor((HG, HD), cutlass.BFloat16) - for h in cutlass.range_constexpr(HG): - qh = hg_idx * HG + h - for d in range(HD): - q_reg[h, d] = mQ[tok_idx, qh, d] - - acc_O = cute.make_rmem_tensor((HG, HD), cutlass.Float32) - acc_O.fill(0.0) - row_max = cute.make_rmem_tensor((HG,), cutlass.Float32) - row_sum = cute.make_rmem_tensor((HG,), cutlass.Float32) - row_max.fill(-1e30) - row_sum.fill(0.0) - - max_tiles = (TL + KT - 1) // KT - for tile_idx in range(max_tiles): - tile_start = tile_idx * KT - for kv_pos in range(KT): - global_kv = tile_start + kv_pos - for d in range(HD): - valid = global_kv < swa_len - val = cutlass.BFloat16(0.0) - if valid: - val = mKV[tok_idx, global_kv, d] - sKV[kv_pos, d] = val - - cute.arch.sync_threads() - - scores = cute.make_rmem_tensor((HG, KT), cutlass.Float32) - scores.fill(0.0) - for h in cutlass.range_constexpr(HG): - for kv_pos in range(KT): - dot = cutlass.Float32(0.0) - for d in range(HD): - q_val = q_reg[h, d].to(cutlass.Float32) - k_val = sKV[kv_pos, d].to(cutlass.Float32) - dot = dot + q_val * k_val - scores[h, kv_pos] = dot * softmax_scale - - for h in cutlass.range_constexpr(HG): - tile_max = cutlass.Float32(-1e30) - for kv_pos in range(KT): - s = scores[h, kv_pos] - if s > tile_max: - tile_max = s - new_max = row_max[h] - if tile_max > new_max: - new_max = tile_max - rescale = cutlass.Float32(0.0) - if row_max[h] > cutlass.Float32(-1e29): - rescale = cute.exp(row_max[h] - new_max) - for d in range(HD): - acc_O[h, d] = acc_O[h, d] * rescale - row_sum[h] = row_sum[h] * rescale - for kv_pos in range(KT): - exp_score = cute.exp(scores[h, kv_pos] - new_max) - row_sum[h] = row_sum[h] + exp_score - for d in range(HD): - v_val = sKV[kv_pos, d].to(cutlass.Float32) - acc_O[h, d] = acc_O[h, d] + exp_score * v_val - row_max[h] = new_max - - cute.arch.sync_threads() - - for h in cutlass.range_constexpr(HG): - qh = hg_idx * HG + h - for d in range(HD): - val_f32 = cutlass.Float32(0.0) - if has_kv and row_sum[h] > cutlass.Float32(1e-30): - val_f32 = acc_O[h, d] / row_sum[h] - mO[tok_idx, qh, d] = val_f32.to(cutlass.BFloat16) diff --git a/dsv4/ops/decode_swa.py b/dsv4/ops/decode_swa.py deleted file mode 100644 index 6871bfe2..00000000 --- a/dsv4/ops/decode_swa.py +++ /dev/null @@ -1,375 +0,0 @@ -""" -Blackwell SM100 Tensor-Core SWA Decode Attention Kernel for DeepSeek-V4. - -Architecture: Two GEMMs back-to-back sharing TMEM, softmax in registers between them. -Following dense_blockscaled_gemm_persistent.py for all Blackwell idioms: -- tcgen05.mma with TMEM accumulators -- TmemAllocator with holding buffer and dealloc barrier -- Warp specialization: 1 MMA warp + 2 epilogue warps -- Online softmax in epilogue warps between the two GEMMs -- Final normalize in epilogue (divide by row_sum) -""" - -import torch -from typing import Optional -import math - -try: - import cutlass - import cutlass.cute as cute - import cutlass.torch as cutlass_torch - import cutlass.utils as utils - import cuda.bindings.driver as cuda - from cutlass.cute.nvgpu import tcgen05, warp - import cutlass.pipeline as pipeline - from cutlass.utils import blackwell_helpers as sm100_utils - from cutlass import BFloat16, Float32 - HAS_CUTEDSL = True -except ImportError: - HAS_CUTEDSL = False - -HEAD_DIM = 512 -KV_TILE = 16 -WINDOW_SIZE = 128 -LOG2_E = 1.4426950408889634074 - - -def native_swa_decode_attention( - q, swa_kv_cache, swa_inv_scale, swa_indices, swa_lens, - block_size, scale, window_size=128, -): - num_tokens, NH, HD = q.shape - device = q.device - if not HAS_CUTEDSL: - return _fallback_batched_sdp(q, swa_kv_cache, swa_inv_scale, - swa_indices, swa_lens, block_size, scale, window_size) - q = q.contiguous(); swa_indices = swa_indices.contiguous(); swa_lens = swa_lens.contiguous() - if swa_indices.dim() == 3: swa_indices_2d = swa_indices.squeeze(0)[:num_tokens] - else: swa_indices_2d = swa_indices[:num_tokens] - max_len = swa_lens[:num_tokens].max().item() - if max_len <= 0: return torch.zeros(num_tokens, NH, HD, dtype=torch.bfloat16, device=device) - max_len = min(max_len, window_size) - safe_indices = swa_indices_2d[:, :max_len].clamp(min=0) - block_indices = safe_indices // block_size; offsets = safe_indices % block_size - kv_raw = swa_kv_cache[block_indices, offsets] - if swa_kv_cache.dtype == torch.uint8: kv_raw = kv_raw.view(torch.float8_e4m3fn) - inv_scales = swa_inv_scale[safe_indices] - kv_bf16 = (kv_raw.to(torch.bfloat16) * inv_scales).to(torch.bfloat16) - if max_len < window_size: - kv_bf16 = torch.cat([kv_bf16, torch.zeros(num_tokens, window_size-max_len, HD, dtype=torch.bfloat16, device=device)], dim=1) - output = torch.zeros(num_tokens, NH, HD, dtype=torch.bfloat16, device=device) - def to_cute(t): return cutlass_torch.from_dlpack(t).mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(t)) - q_c, kv_c, len_c, out_c = to_cute(q), to_cute(kv_bf16), to_cute(swa_lens[:num_tokens]), to_cute(output) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - scale_c = to_cute(torch.tensor([scale], dtype=torch.float32, device=device)) - kernel = BlackwellSWADecodeKernel(head_dim=HD, num_heads=NH, kv_tile=KV_TILE, window_size=window_size) - compiled = cute.compile(kernel, q_c, kv_c, len_c, out_c, scale_c, stream) - compiled(q_c, kv_c, len_c, out_c, scale_c, stream) - return output - - -def _fallback_batched_sdp(q, swa_kv_cache, swa_inv_scale, swa_indices, swa_lens, - block_size, scale, window_size): - num_tokens, NH, HD = q.shape; device = q.device - if swa_indices.dim() == 3: swa_indices = swa_indices.squeeze(0) - safe_indices = swa_indices[:num_tokens].clamp(min=0) - block_indices = safe_indices // block_size; offsets = safe_indices % block_size - kv_raw = swa_kv_cache[block_indices, offsets] - if swa_kv_cache.dtype == torch.uint8: kv_raw = kv_raw.view(torch.float8_e4m3fn) - inv_scales = swa_inv_scale[safe_indices] - kv_bf16 = (kv_raw.to(torch.bfloat16) * inv_scales).to(torch.bfloat16) - pos_range = torch.arange(window_size, device=device).unsqueeze(0) - len_mask = pos_range >= swa_lens[:num_tokens].unsqueeze(1) - invalid_mask = swa_indices[:num_tokens, :window_size] < 0 - attn_mask = len_mask | invalid_mask - float_mask = torch.zeros(attn_mask.shape, dtype=torch.bfloat16, device=device) - float_mask[attn_mask] = float('-inf') - q_t = q.permute(1, 0, 2) - kv_batch = kv_bf16.expand(NH, window_size, HD) - mask_batch = float_mask.unsqueeze(0).expand(NH, -1, -1) - out = torch.nn.functional.scaled_dot_product_attention(q_t, kv_batch, kv_batch, attn_mask=mask_batch, is_causal=False, scale=scale) - return out.permute(1, 0, 2) - - -if HAS_CUTEDSL: - - class BlackwellSWADecodeKernel: - def __init__(self, head_dim=HEAD_DIM, num_heads=128, - kv_tile=KV_TILE, window_size=WINDOW_SIZE): - self._head_dim = head_dim - self._num_heads = num_heads - self._kv_tile = kv_tile - self._window_size = window_size - self._mma_m = 128 - self._num_threads = 96 - self._cta_group = tcgen05.CtaGroup.ONE - # Warp IDs: 0,1 = epilogue, 2 = MMA - self._mma_warp_id = 2 - self._epi_warp_ids = [0, 1] - - @cute.jit - def __call__(self, mQ, mKV, mLens, mO, mScale, stream): - num_tokens = mQ.shape[0] - M = self._mma_m; HD = self._head_dim; KT = self._kv_tile - - # TiledMma for Q @ K^T: Q(M,HD) x K^T(HD,KT) → S(M,KT) - tiled_mma_qk = sm100_utils.make_trivial_tiled_mma( - BFloat16, tcgen05.OperandMajorMode.K, tcgen05.OperandMajorMode.K, - Float32, self._cta_group, (M, KT)) - - # TiledMma for P @ V: P(M,KT) x V(KT,HD) → O(M,HD) - tiled_mma_pv = sm100_utils.make_trivial_tiled_mma( - BFloat16, tcgen05.OperandMajorMode.K, tcgen05.OperandMajorMode.MN, - Float32, self._cta_group, (M, KT), - tcgen05.OperandSource.TMEM) - - # SMEM layouts - sA_layout_atom = cute.make_composed_layout( - cute.make_swizzle(3, 3, 3), 0, - cute.make_layout((8, 64), stride=(64, 1))) - sQ_layout = cute.tile_to_shape(sA_layout_atom, (M, HD), (0, 1)) - sK_layout = cute.tile_to_shape(sA_layout_atom, (KT, HD), (0, 1)) - sV_layout = cute.tile_to_shape(sA_layout_atom, (KT, HD), (0, 1)) - sO_layout = cute.tile_to_shape(sA_layout_atom, (M, HD), (0, 1)) - - # Named barriers for TMEM allocation and MMA↔epilogue sync - tmem_alloc_barrier = pipeline.NamedBarrier( - barrier_id=2, num_threads=96) # all 3 warps - acc_full_barrier = pipeline.NamedBarrier( - barrier_id=3, num_threads=96) # all 3 warps - - @cute.struct - class SharedStorage: - sQ: cute.struct.Align[cute.struct.MemRange[BFloat16, cute.cosize(sQ_layout)], 1024] - sK: cute.struct.Align[cute.struct.MemRange[BFloat16, cute.cosize(sK_layout)], 1024] - sV: cute.struct.Align[cute.struct.MemRange[BFloat16, cute.cosize(sV_layout)], 1024] - sO: cute.struct.Align[cute.struct.MemRange[BFloat16, cute.cosize(sO_layout)], 1024] - tmem_dealloc_mbar: cutlass.Int64 - tmem_holding_buf: cutlass.Int32 - - self._kernel( - mQ, mKV, mLens, mO, mScale, - sQ_layout, sK_layout, sV_layout, sO_layout, - tiled_mma_qk, tiled_mma_pv, SharedStorage, - tmem_alloc_barrier, acc_full_barrier, - ).launch(grid=(1, num_tokens, 1), block=[self._num_threads, 1, 1], stream=stream) - - @cute.kernel - def _kernel(self, mQ, mKV, mLens, mO, mScale, - sQ_layout, sK_layout, sV_layout, sO_layout, - tiled_mma_qk, tiled_mma_pv, SharedStorage: cutlass.Constexpr, - tmem_alloc_barrier, acc_full_barrier): - tidx, _, _ = cute.arch.thread_idx() - _, tok_idx, _ = cute.arch.block_idx() - - M = self._mma_m; HD = self._head_dim; KT = self._kv_tile - softmax_scale = mScale[0] - swa_len = mLens[tok_idx] - - warp_idx = tidx // 32 - is_mma_warp = warp_idx == self._mma_warp_id - is_epi_warp = warp_idx in self._epi_warp_ids - - smem = utils.SmemAllocator() - storage = smem.allocate(SharedStorage) - sQ = storage.sQ.get_tensor(sQ_layout) - sK = storage.sK.get_tensor(sK_layout) - sV = storage.sV.get_tensor(sV_layout) - sO = storage.sO.get_tensor(sO_layout) - - # TMEM allocator (all warps participate in alloc/dealloc) - tmem = utils.TmemAllocator( - storage.tmem_holding_buf.ptr, - barrier_for_retrieve=tmem_alloc_barrier, - allocator_warp_id=self._epi_warp_ids[0], - ) - - # Allocate TMEM for score and output accumulators - # Score: (M=128, KT=16) = 2048 FP32 - # Output: (M=128, HD=512) = 65536 FP32 - # Total: 67584 FP32 = 264 KB TMEM (within 1 MB budget) - num_tmem_cols = M * KT + M * HD # TMEM columns - tmem.allocate(num_tmem_cols) - - # TMEM layout for scores and output - # The TMEM layout comes from the TiledMma's C operand layout - tCtScores = tiled_mma_qk.make_fragment_C(tiled_mma_qk.partition_C_shape((M, KT))) - tCtOutput = tiled_mma_pv.make_fragment_C(tiled_mma_pv.partition_C_shape((M, HD))) - - # Zero the output accumulator (MMA with ACCUMULATE=False does this for scores) - # For the output accumulator, we zero it before the KV loop - # TODO: zero tCtOutput via tcgen05.st or first PV GEMM with ACCUMULATE=False - - # ─── MMA WARP ──────────────────────────────────────── - if is_mma_warp: - # Load Q to SMEM (once, reused across all KV tiles) - for h in range(M): - for d in range(HD): - sQ[h, d] = mQ[tok_idx, h, d] - cute.arch.sync_threads() - - # Partition Q and K for QK GEMM - thr_qk = tiled_mma_qk.get_slice(tidx) - tCrQ = thr_qk.make_fragment_A(thr_qk.partition_A(sQ)) - tCrK = thr_qk.make_fragment_B(thr_qk.partition_B(sK)) - - smem_copy_Q = cute.make_tiled_copy_A( - cute.make_copy_atom(warp.LdMatrix8x8x16bOp(False, 4), BFloat16), tiled_mma_qk) - thr_smem_Q = smem_copy_Q.get_slice(tidx) - tCsQ = thr_smem_Q.partition_S(sQ); tCrQ_cv = thr_smem_Q.retile(tCrQ) - - smem_copy_K = cute.make_tiled_copy_B( - cute.make_copy_atom(warp.LdMatrix8x8x16bOp(False, 4), BFloat16), tiled_mma_qk) - thr_smem_K = smem_copy_K.get_slice(tidx) - tCsK = thr_smem_K.partition_S(sK); tCrK_cv = thr_smem_K.retile(tCrK) - - # Load Q to registers (once) - for k in cutlass.range_constexpr(cute.size(tCsQ.shape[2])): - cute.copy(smem_copy_Q, tCsQ[None, None, k], tCrQ_cv[None, None, k]) - - # Partition V for PV GEMM - thr_pv = tiled_mma_pv.get_slice(tidx) - sVt = cute.composition(sV, cute.make_layout((HD, KT), stride=(KT, 1))) - tOrVt = thr_pv.make_fragment_B(thr_pv.partition_B(sVt)) - smem_copy_V = cute.make_tiled_copy_B( - cute.make_copy_atom(warp.LdMatrix8x8x16bOp(True, 4), BFloat16), tiled_mma_pv) - thr_smem_V = smem_copy_V.get_slice(tidx) - tOsVt = thr_smem_V.partition_S(sVt); tOrVt_cv = thr_smem_V.retile(tOrVt) - - # ── KV tile loop ───────────────────────────────── - n_block_max = (self._window_size + KT - 1) // KT - - for n_block in range(n_block_max): - tile_start = n_block * KT - - # Load K and V to SMEM - for kv_pos in range(KT): - global_kv = tile_start + kv_pos - for d in range(HD): - val = cutlass.BFloat16(0.0) - if global_kv < swa_len: - val = mKV[tok_idx, global_kv, d] - sK[kv_pos, d] = val - sV[kv_pos, d] = val - - cute.arch.sync_threads() - - # ── Q @ K^T via tcgen05.mma ────────────────── - for k in cutlass.range_constexpr(cute.size(tCsK.shape[2])): - cute.copy(smem_copy_K, tCsK[None, None, k], tCrK_cv[None, None, k]) - tiled_mma_qk.set(tcgen05.Field.ACCUMULATE, k > 0) - cute.gemm(tiled_mma_qk, tCtScores, tCrQ[None, None, k], tCrK[None, None, k], tCtScores) - - # Signal epilogue: scores in TMEM are ready - cute.arch.fence_view_async_tmem_store() - acc_full_barrier.arrive() - - # Wait for epilogue to finish softmax - acc_full_barrier.wait() - - # ── P @ V via tcgen05.mma ───────────────────── - # P from TMEM (softmax output), V from SMEM - for k in cutlass.range_constexpr(cute.size(tOsVt.shape[2])): - cute.copy(smem_copy_V, tOsVt[None, None, k], tOrVt_cv[None, None, k]) - # For the first KV tile, first k tile: ACCUMULATE=False (zero the output) - # Otherwise, ACCUMULATE=True (accumulate into output) - is_first_output_tile = (n_block == 0) and (k == 0) - tiled_mma_pv.set(tcgen05.Field.ACCUMULATE, not is_first_output_tile) - cute.gemm(tiled_mma_pv, tCtOutput, None, tOrVt[None, None, k], tCtOutput) - - cute.arch.fence_view_async_tmem_store() - acc_full_barrier.arrive() - - # Free TMEM - tmem.relinquish_alloc_permit() - acc_full_barrier.arrive() # Sync before dealloc - tmem.free(num_tmem_cols) - - # ─── EPILOGUE WARPS ────────────────────────────────── - if is_epi_warp: - my_row_start = warp_idx * 64 - num_my_rows = 64 - - # Online softmax state - row_max = cute.make_rmem_tensor((num_my_rows,), Float32) - row_sum = cute.make_rmem_tensor((num_my_rows,), Float32) - row_max.fill(-1e30) - row_sum.fill(0.0) - - # TMEM→register copy for scores (tcgen05.ld pattern) - tiled_copy_t2r_scores = tcgen05.make_tmem_copy( - sm100_utils.get_tmem_load_op(self._cta_group), tCtScores) - - # TMEM→register for output - tiled_copy_t2r_output = tcgen05.make_tmem_copy( - sm100_utils.get_tmem_load_op(self._cta_group), tCtOutput) - - # Register fragments - tRrScores = cute.make_fragment_like( - tiled_copy_t2r_scores.partition_D( - cute.make_tensor(tCtScores.iterator, tCtScores.layout)), Float32) - tRrOutput = cute.make_fragment_like( - tiled_copy_t2r_output.partition_D( - cute.make_tensor(tCtOutput.iterator, tCtOutput.layout)), Float32) - - n_block_max = (self._window_size + KT - 1) // KT - - for n_block in range(n_block_max): - # Wait for MMA to finish Q@K^T - acc_full_barrier.wait() - - # ── tcgen05.ld scores from TMEM to registers ── - cute.copy(tiled_copy_t2r_scores, tCtScores, tRrScores) - cute.arch.fence_view_async_tmem_load() - - # ── Softmax in registers ────────────────────── - # For each row this warp owns (64 rows per warp): - # 1. tile_max = max(scores * scale) — reduce across KT positions - # 2. new_max = max(row_max_prev, tile_max) - # 3. prev_exp = exp((row_max_prev - new_max) * scale * log2e) - # 4. Rescale output accumulator in TMEM: O *= prev_exp - # (via tcgen05.st with scaled values, or defer to PV GEMM) - # 5. row_sum *= prev_exp - # 6. exp_scores = exp((scores * scale - new_max * scale) * log2e) - # 7. row_sum += sum(exp_scores) - # 8. Write exp_scores back to TMEM (as P operand for PV GEMM) - # 9. row_max = new_max - - # The register fragment layout after tcgen05.ld is - # (EPI_TILE_M=128, EPI_TILE_N=16) partitioned per epilogue warp. - # Each epilogue warp's fragment covers 64 rows and 16 columns. - # We can iterate over rows and compute softmax per row. - - # TODO: Implement per-row softmax on the register fragment. - # This requires understanding the exact tRrScores layout - # from tcgen05.make_tmem_copy + partition_D. - # The dense GEMM's epilogue shows the pattern for iterating - # over the register fragment. - - # For now, write the P values (softmax output) back to TMEM - # via tcgen05.st (register → TMEM copy) - # tcgen05.st: tRrScores → tCtScores - cute.copy(tiled_copy_t2r_scores, tRrScores, tCtScores) - cute.arch.fence_view_async_tmem_store() - - # Signal MMA: softmax done, P in TMEM ready for PV GEMM - acc_full_barrier.arrive() - - # Wait for MMA to finish P@V - acc_full_barrier.wait() - - # ── Final normalize ────────────────────────────── - # tcgen05.ld output from TMEM, divide by row_sum, store to GMEM - cute.copy(tiled_copy_t2r_output, tCtOutput, tRrOutput) - cute.arch.fence_view_async_tmem_load() - - # Divide each element by row_sum - # TODO: implement per-row normalization on register fragment - - # Cast to BF16 and store to SMEM, then GMEM - # (following dense GEMM's epilogue pattern) - - # Relinquish TMEM - tmem.relinquish_alloc_permit() - acc_full_barrier.arrive() - tmem.free(num_tmem_cols) diff --git a/run_router_tests.py b/run_router_tests.py deleted file mode 100644 index 5ee89d34..00000000 --- a/run_router_tests.py +++ /dev/null @@ -1,97 +0,0 @@ -#!/usr/bin/env python3 -import sys, os -os.chdir('/root/dsv4-nvfp4-workspace/kernel') -sys.path.insert(0, '.') -import torch -from torch.utils.cpp_extension import load - -def pf(ok): - return 'PASS' if ok else 'FAIL' - -print(f'CUDA: {torch.cuda.is_available()}, Device: {torch.cuda.get_device_name(0)}', flush=True) -print(f'GPU mem free: {torch.cuda.mem_get_info()[0]//1024//1024} MB', flush=True) - -# 1. Hash Router -print('\n=== Hash Router ===', flush=True) -try: - hr = load(name='hr2', sources=['dsv4/kernels/cuda/hash_router.cu'], - extra_cuda_cflags=['-O3', '--generate-code=arch=compute_100a,code=[sm_100a]'], verbose=False) - for N in [1, 4, 64, 128, 512]: - vocab, k, E = 128000, 6, 256 - lut = torch.randint(0, E, (vocab, k), dtype=torch.int32, device='cuda') - tids = torch.randint(0, vocab, (N,), dtype=torch.int32, device='cuda') - ow = torch.empty(N, k, dtype=torch.float32, device='cuda') - oi = torch.empty(N, k, dtype=torch.int32, device='cuda') - hr.hash_router(tids, lut, k, ow, oi) - torch.cuda.synchronize() - exp_ids = lut[tids] - exp_w = torch.full((N, k), 1.0/k, dtype=torch.float32, device='cuda') - ids_ok = (oi == exp_ids).all().item() - w_ok = torch.allclose(ow, exp_w, atol=1e-7, rtol=1e-7) - ok = ids_ok and w_ok - print(f' N={N:4d}: IDs={ids_ok} W={w_ok} {pf(ok)}', flush=True) - del lut, tids, ow, oi, exp_ids, exp_w - print('Hash Router: ALL PASS', flush=True) -except Exception as e: - import traceback; traceback.print_exc() - -torch.cuda.empty_cache() - -# 2. Top-k Select -print('\n=== Top-k Select ===', flush=True) -try: - tk = load(name='tk2', sources=['dsv4/kernels/cuda/topk_select.cu'], - extra_cuda_cflags=['-O3', '--generate-code=arch=compute_100a,code=[sm_100a]'], verbose=False) - for N, E in [(1,256), (4,256), (64,256), (64,384), (128,256), (512,256)]: - k = 6 - scores = torch.randn(N, E, dtype=torch.float32, device='cuda') - ov, oidx = tk.topk_select(scores, k) - torch.cuda.synchronize() - exp = scores.topk(k, dim=-1) - ids_ok = (oidx == exp.indices).all().item() - vals_ok = torch.allclose(ov, exp.values, atol=1e-6, rtol=1e-6) - ok = ids_ok and vals_ok - print(f' N={N:4d} E={E}: IDs={ids_ok} V={vals_ok} {pf(ok)}', flush=True) - del scores, ov, oidx, exp - print('Top-k Select: ALL PASS', flush=True) -except Exception as e: - import traceback; traceback.print_exc() - -torch.cuda.empty_cache() - -# 3. Activation + Top-k -print('\n=== Activation + Top-k ===', flush=True) -try: - atk = load(name='atk2', sources=['dsv4/kernels/cuda/activation_topk.cu'], - extra_cuda_cflags=['-O3', '--generate-code=arch=compute_100a,code=[sm_100a]'], verbose=False) - for N, E in [(1,256), (4,256), (64,256), (64,384)]: - k = 6 - logits = torch.randn(N, E, dtype=torch.float32, device='cuda') - bias = torch.randn(E, dtype=torch.float32, device='cuda') - scaling = 2.5 - out_w = torch.empty(N, k, dtype=torch.float32, device='cuda') - out_ids = torch.empty(N, k, dtype=torch.int32, device='cuda') - atk.activation_topk(logits, bias, k, scaling, out_w, out_ids) - torch.cuda.synchronize() - # Oracle - sp = torch.log1p(logits.exp()) + torch.clamp(logits, min=0) - act = sp.sqrt() - score = act + bias - exp_topk = score.topk(k, dim=-1) - exp_ids = exp_topk.indices - exp_w = torch.gather(act, 1, exp_ids) - exp_w = exp_w / exp_w.sum(dim=-1, keepdim=True) * scaling - ids_ok = (out_ids == exp_ids).all().item() - vals_ok = torch.allclose(out_w, exp_w, atol=1e-5, rtol=1e-5) - ok = ids_ok and vals_ok - print(f' N={N:4d} E={E}: IDs={ids_ok} V={vals_ok} {pf(ok)}', flush=True) - if not ids_ok: - print(f' ID mismatches: {(out_ids != exp_ids).sum().item()}/{out_ids.numel()}', flush=True) - if not vals_ok: - print(f' Max diff: {(out_w - exp_w).abs().max().item():.2e}', flush=True) - del logits, bias, out_w, out_ids, sp, act, score, exp_topk, exp_ids, exp_w - print('Activation+Topk: ALL PASS', flush=True) -except Exception as e: - import traceback; traceback.print_exc() - -print('\n=== DONE ===', flush=True) diff --git a/test_mapping.py b/test_mapping.py deleted file mode 100644 index 11eb7cf1..00000000 --- a/test_mapping.py +++ /dev/null @@ -1,50 +0,0 @@ -#!/usr/bin/env python3 -""" -Test SMEM-P coordinate mapping for collisions. -Mapping: (m,n) -> ((m, n%16), 0, ((n//16)%4, n//64), 0) -Offset: m*64 + (n%16) + ((n//16)%4)*16 + (n//64)*8192 -""" - -def compute_offset(m, n): - n0 = n % 16 - n1 = (n // 16) % 4 - n2 = n // 64 - offset = m * 64 + n0 + n1 * 16 + n2 * 8192 - return offset, (m, n0, n1, n2) - -# Test bijection -offsets = set() -collisions = [] - -for m in range(128): - for n in range(128): - offset, _ = compute_offset(m, n) - if offset in offsets: - collisions.append(((m,n), offset)) - offsets.add(offset) - -print(f"Total positions: {128*128} = {128*128}") -print(f"Unique offsets: {len(offsets)}") -print(f"Collisions: {len(collisions)}") -if collisions: - print("First few collisions:", collisions[:5]) - -# Check offset range -min_offset = min(offsets) -max_offset = max(offsets) -print(f"Offset range: {min_offset} .. {max_offset}") -print(f"Expected range: 0 .. 16383") - -# Test specific coordinates -test_coords = [(0,0), (0,1), (0,15), (0,16), (0,63), (0,64), (0,127), - (1,0), (1,1), (1,127), - (127,0), (127,63), (127,64), (127,127)] -print("\nSample offsets:") -for m,n in test_coords: - offset, decomp = compute_offset(m,n) - print(f"({m},{n}) -> {decomp} = offset {offset}") - -# Verify formula matches layout stride -# Layout: ((128,16),1,(4,2),1):((64,1),0,(16,8192),0) -# Stride for (m,n0,n1,n2): m*64 + n0*1 + n1*16 + n2*8192 -print("\nLayout stride check: OK" if max_offset == 16383 and len(offsets) == 16384 else "Layout mismatch") \ No newline at end of file diff --git a/dsv4/kernels/attention/fmha_qk_verify.cuh b/tests/unit/qk_verify_kernel.cuh similarity index 100% rename from dsv4/kernels/attention/fmha_qk_verify.cuh rename to tests/unit/qk_verify_kernel.cuh diff --git a/tests/unit/test_cotiled_diag.py b/tests/unit/test_cotiled_diag.py deleted file mode 100644 index 1a6e8f92..00000000 --- a/tests/unit/test_cotiled_diag.py +++ /dev/null @@ -1,188 +0,0 @@ -"""Minimal diagnostic: test layout composition for SMEM-P make_cotiled_copy. - -Uses FmhaKernel's __call__ path to set up all layouts, then extracts -the TV layout and sP layout at the right point. -""" -import torch, math -import cutlass, cutlass.cute as cute -import cutlass.utils as utils -from cutlass.cute.nvgpu import tcgen05 -from cutlass import Float32, BFloat16, Int32 -import cutlass.torch as ct -import cuda.bindings.driver as cuda -from dsv4.kernels.attention.fmha import FmhaKernel - - -def main(): - head_dim = 256 - s_k = 128 - m = 128 - pv_n_tile = min(head_dim, 256) - - q = torch.randn(m, head_dim, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(s_k, head_dim, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, head_dim, dtype=torch.bfloat16, device='cuda') - c = torch.zeros(m, head_dim, 1, dtype=torch.bfloat16, device='cuda') - - kernel = FmhaKernel(head_dim=head_dim, s_k=s_k) - - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - # Reproduce the EXACT __call__ setup to get the MMA objects - v_tile = v[:, 0:pv_n_tile].contiguous() - v_kernel = v_tile.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[:, 0:pv_n_tile, :]).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c[:, 0:pv_n_tile, :])) - - # Derive major modes exactly as FmhaKernel does - from cutlass.utils import LayoutEnum - from cutlass.cute.nvgpu import OperandMajorMode - a_major = LayoutEnum.from_tensor(mQ).mma_major_mode() - b_major = LayoutEnum.from_tensor(mK).mma_major_mode() - v_major = LayoutEnum.from_tensor(mV).mma_major_mode() - print(f"a_major: {a_major}, b_major: {b_major}, v_major: {v_major}") # layout (256, 128, 1) stride (1, 256, 32768) = col-major - - c_layout = LayoutEnum.from_tensor(mC) - - qk_mma = utils.sm100.make_trivial_tiled_mma( - BFloat16, BFloat16, a_major, b_major, Float32, - tcgen05.CtaGroup.ONE, (128, 128), tcgen05.OperandSource.SMEM - ) - pv_a_major = a_major # SMEM-P path - pv_source = tcgen05.OperandSource.SMEM - pv_mma = utils.sm100.make_trivial_tiled_mma( - BFloat16, BFloat16, pv_a_major, v_major, Float32, - tcgen05.CtaGroup.ONE, (128, pv_n_tile), pv_source - ) - - 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, pv_n_tile, pv_ik * (128 // pv_ik)) - - # sP layout (PV A-operand SMEM) - p_smem_s = utils.sm100.make_smem_layout_a(pv_mma, pv_mma_tiler, BFloat16, 1) - - # QK C-fragment - qk_thr = qk_mma.get_slice(0) - qk_as = qk_thr.partition_shape_C(qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - tStS0 = cute.make_tensor(tStS.iterator, tStS.layout) - - # TMEM-load copy - 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, tStS0) - - # =========================== - # Print layouts - # =========================== - dst_tv = tiled_tmem_load.layout_dst_tv_tiled - print(f"1. dst_tv shape: {dst_tv.shape}") - print(f" dst_tv stride: {dst_tv.stride}") - print(f" dst_tv: {dst_tv}") - - sP_outer = p_smem_s.outer - sP_coalesced = cute.coalesce(sP_outer) - print(f"\n2. sP outer shape: {cute.shape(sP_outer)}") - print(f" sP outer: {sP_outer}") - print(f" sP coalesced: {sP_coalesced}") - - tStS_coalesced = cute.coalesce(tStS0.layout) - print(f"\n3. tStS layout: {tStS0.layout}") - print(f" tStS coalesced: {tStS_coalesced}") - print(f" tStS coalesced shape: {cute.shape(tStS_coalesced)}") - - # =========================== - # Build sP layout in (128, 128) coordinate space - # =========================== - # sP outer shape is ((128,16),1,(4,2),1) with strides ((64,1),0,(16,8192),0) - # Equivalent to (128, (16, 4, 2)) with strides (64, (1, 16, 8192)) - sP_2d = cute.make_layout( - (128, (16, 4, 2)), - stride=(64, (1, 16, 8192)) - ) - print(f"\n4. sP_2d: {sP_2d}") - print(f" sP_2d size: {cute.size(sP_2d)}") - - # =========================== - # Try left_inverse(tStS_coalesced) - # =========================== - print(f"\n5. Attempting left_inverse(tStS_coalesced)...") - try: - tStS_inv = cute.left_inverse(tStS_coalesced) - print(f" tStS_inv: {tStS_inv}") - print(f" tStS_inv shape: {cute.shape(tStS_inv)}") - except Exception as e: - print(f" FAILED: {e}") - import traceback; traceback.print_exc() - return - - # =========================== - # Try composition: sP_2d ∘ tStS_inv → reindex - # =========================== - print(f"\n6. Attempting composition(sP_2d, tStS_inv)...") - reindex = None - try: - reindex = cute.composition(sP_2d, tStS_inv) - print(f" reindex: {reindex}") - print(f" reindex shape: {cute.shape(reindex)}") - except Exception as e: - print(f" FAILED: {e}") - # Try with sP_coalesced instead - try: - reindex = cute.composition(sP_coalesced, tStS_inv) - print(f" reindex (coalesced): {reindex}") - except Exception as e2: - print(f" ALSO FAILED: {e2}") - import traceback; traceback.print_exc() - return - - # =========================== - # Try composition: reindex ∘ dst_tv → atom_layout_tv - # =========================== - print(f"\n7. Attempting composition(reindex, dst_tv)...") - atom_layout_tv = None - try: - atom_layout_tv = cute.composition(reindex, dst_tv) - print(f" atom_layout_tv: {atom_layout_tv}") - print(f" atom_layout_tv shape: {cute.shape(atom_layout_tv)}") - print(f" atom_layout_tv stride: {atom_layout_tv.stride}") - except Exception as e: - print(f" FAILED: {e}") - import traceback; traceback.print_exc() - return - - # =========================== - # Try make_cotiled_copy - # =========================== - print(f"\n8. Attempting make_cotiled_copy...") - try: - r2s_atom = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), - BFloat16, - num_bits_per_copy=16, - ) - tiled_r2s = cute.make_cotiled_copy(r2s_atom, atom_layout_tv, sP_coalesced) - print(f" make_cotiled_copy SUCCEEDED!") - print(f" layout_tv_tiled: {tiled_r2s.layout_tv_tiled}") - print(f" layout_dst_tv_tiled: {tiled_r2s.layout_dst_tv_tiled}") - - # Try get_slice - try: - thr_r2s = tiled_r2s.get_slice(0) - print(f" get_slice(0) SUCCEEDED!") - except Exception as e: - print(f" get_slice(0) FAILED: {e}") - - except Exception as e: - print(f" FAILED: {e}") - import traceback; traceback.print_exc() - - -if __name__ == '__main__': - main() diff --git a/tests/unit/test_d15_in_kernel_rescale.py b/tests/unit/test_d15_in_kernel_rescale.py deleted file mode 100644 index 13b7c4c3..00000000 --- a/tests/unit/test_d15_in_kernel_rescale.py +++ /dev/null @@ -1,152 +0,0 @@ -""" -D1.5 Phase 4: Test in-kernel O rescale for multi-KV-tile FMHA. - -Tests the CUTLASS correction_rescale pattern: - - Both load and store atoms built from the SAME tOtO_i (composition-tiled) - - Same Repetition(corr_tile_size=16) for both - - Rescale O in TMEM between PV iterations - -Compares against: - 1. FP32 reference (ground truth) - 2. Python KV merge (proven correct, cos 0.999998) - 3. s_k=128 baseline (no rescale, regression check) -""" -import torch, math -import cutlass.cute as cute -import cutlass.torch as ct -import cuda.bindings.driver as cuda -from dsv4.kernels.attention.fmha import FmhaKernel - - -def reference_attention(q, k, v, scale): - """FP32 reference: returns un-normalized O.""" - qf = q.float() - kf = k.float() - attn = qf @ kf.T * scale - attn_max = attn.max(dim=-1, keepdim=True)[0] - attn_exp = torch.exp(attn - attn_max) - ref_unnorm = attn_exp @ v.float() - return ref_unnorm - - -def run_fmha(q, k, v, head_dim, s_k, pv_n_tile, use_smem_p, stream, lse_tensor, row_sums_tensor): - """Run FMHA kernel and return output tensor.""" - m = q.shape[0] - v_tile = v[:, 0:pv_n_tile].contiguous() - v_kernel = v_tile.unsqueeze(-1) - c_tile = torch.zeros(m, pv_n_tile, 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_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) - mRS = ct.from_dlpack(row_sums_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(row_sums_tensor)) - - kernel = FmhaKernel(head_dim=head_dim, s_k=s_k, use_smem_p=use_smem_p, normalize=False) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE, row_sums=mRS) - compiled(mQ, mK, mV, mC, stream, mLSE, row_sums=mRS) - return c_tile, lse_tensor, row_sums_tensor, kernel - - -def test(): - hd = 64 - m = 128 - scale = 1.0 / math.sqrt(hd) - torch.manual_seed(42) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - - # ===== Test 1: s_k=128 baseline (no rescale) ===== - s_k1 = 128 - k1 = torch.randn(s_k1, hd, 1, dtype=torch.bfloat16, device='cuda') - v1 = torch.randn(s_k1, hd, dtype=torch.bfloat16, device='cuda') - lse1 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - rs1 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - - # Need a dummy run to get pv_n_tile - kernel0 = FmhaKernel(head_dim=hd, s_k=s_k1, use_smem_p=False, normalize=False) - pv_n_tile = kernel0.pv_n_tile - - c1, lse1, rs1, _ = run_fmha(q, k1, v1, hd, s_k1, pv_n_tile, False, stream, lse1, rs1) - torch.cuda.synchronize() - - ref1 = reference_attention(q[:, :, 0], k1[:, :, 0], v1, scale) - cos1 = torch.nn.functional.cosine_similarity( - c1[:, :, 0].float().flatten().unsqueeze(0), ref1.flatten().unsqueeze(0) - ).item() - status1 = "PASS" if cos1 >= 0.999 else "FAIL" - print(f'Test 1: s_k=128 baseline: cos={cos1:.6f} {status1}', flush=True) - - # ===== Test 2: s_k=256 with in-kernel rescale (CUTLASS correction_rescale) ===== - s_k2 = 256 - k2 = torch.randn(s_k2, hd, 1, dtype=torch.bfloat16, device='cuda') - v2 = torch.randn(s_k2, hd, dtype=torch.bfloat16, device='cuda') - lse2 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - rs2 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - - c2, lse2, rs2, _ = run_fmha(q, k2, v2, hd, s_k2, pv_n_tile, False, stream, lse2, rs2) - torch.cuda.synchronize() - - ref2 = reference_attention(q[:, :, 0], k2[:, :, 0], v2, scale) - cos2 = torch.nn.functional.cosine_similarity( - c2[:, :, 0].float().flatten().unsqueeze(0), ref2.flatten().unsqueeze(0) - ).item() - status2 = "PASS" if cos2 >= 0.999 else "FAIL" - print(f'Test 2: s_k=256 in-kernel rescale: cos={cos2:.6f} {status2}', flush=True) - - # ===== Test 3: Python KV merge (oracle) ===== - c_s0 = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - lse_s0 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - rs_s0 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - c_s0, lse_s0, rs_s0, _ = run_fmha(q, k2[:128], v2[:128], hd, 128, pv_n_tile, False, stream, lse_s0, rs_s0) - - c_s1 = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - lse_s1 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - rs_s1 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - c_s1, lse_s1, rs_s1, _ = run_fmha(q, k2[128:], v2[128:], hd, 128, pv_n_tile, False, stream, lse_s1, rs_s1) - torch.cuda.synchronize() - - # D5 merge: O = sum(exp(lse_i) * O_i_norm) / sum(exp(lse_i)) - o0 = c_s0[:, :, 0].float() - o1 = c_s1[:, :, 0].float() - r0 = rs_s0[:, 0, 0].float() - r1 = rs_s1[:, 0, 0].float() - l0 = lse_s0[:, 0, 0].float() - l1 = lse_s1[:, 0, 0].float() - o0_norm = o0 / r0.unsqueeze(1).clamp(min=1e-30) - o1_norm = o1 / r1.unsqueeze(1).clamp(min=1e-30) - w0 = torch.exp(l0).unsqueeze(1) - w1 = torch.exp(l1).unsqueeze(1) - oracle = (w0 * o0_norm + w1 * o1_norm) / (w0 + w1) - - cos_oracle = torch.nn.functional.cosine_similarity( - oracle.flatten().unsqueeze(0), ref2.flatten().unsqueeze(0) - ).item() - print(f'Oracle: Python KV merge: cos={cos_oracle:.6f}', flush=True) - - # ===== Test 4: s_k=384 (3 KV tiles) ===== - s_k3 = 384 - k3 = torch.randn(s_k3, hd, 1, dtype=torch.bfloat16, device='cuda') - v3 = torch.randn(s_k3, hd, dtype=torch.bfloat16, device='cuda') - lse3 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - rs3 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - - c3, lse3, rs3, _ = run_fmha(q, k3, v3, hd, s_k3, pv_n_tile, False, stream, lse3, rs3) - torch.cuda.synchronize() - - ref3 = reference_attention(q[:, :, 0], k3[:, :, 0], v3, scale) - cos3 = torch.nn.functional.cosine_similarity( - c3[:, :, 0].float().flatten().unsqueeze(0), ref3.flatten().unsqueeze(0) - ).item() - status3 = "PASS" if cos3 >= 0.999 else "FAIL" - print(f'Test 4: s_k=384 in-kernel rescale: cos={cos3:.6f} {status3}', flush=True) - - # ===== Summary ===== - all_pass = cos1 >= 0.999 and cos2 >= 0.999 and cos3 >= 0.999 - print(f'\nSummary: {"ALL PASS ✅" if all_pass else "SOME FAIL ❌"}', flush=True) - - -if __name__ == '__main__': - test() diff --git a/tests/unit/test_d15_multi_kv.py b/tests/unit/test_d15_multi_kv.py deleted file mode 100644 index cdb41380..00000000 --- a/tests/unit/test_d15_multi_kv.py +++ /dev/null @@ -1,192 +0,0 @@ -""" -FMHA D1.5: Multi-KV-tile attention with Python KV merge. - -The kernel processes one KV tile at a time (s_k=128 per tile). -For s_k>128, we run the kernel multiple times and merge the results -using per-tile LSE values. - -Merge formula (D5 merge for same Q, different KV segments): - O = sum_i [exp(lse_i) * O_i_norm] / sum_i [exp(lse_i)] - -Where O_i_norm is the normalized output for segment i, and lse_i is the -log-sum-exp for that segment. - -Run: ~/.openclaw/workspace/fire_b200_test tests/unit/test_d15_multi_kv.py -""" -import torch -import math -import cutlass.cute as cute -import cutlass.torch as ct -import cuda.bindings.driver as cuda -from dsv4.kernels.attention.fmha import FmhaKernel - - -def reference_attention(q, k, v, scale): - """FP32 reference: q (M, hd), k (s_k, hd), v (s_k, hd) → o (M, hd), lse (M,)""" - scores = torch.matmul(q.float(), k.float().T) * scale - max_s = scores.max(dim=-1, keepdim=True).values - exp_s = (scores - max_s).exp() - sum_s = exp_s.sum(dim=-1, keepdim=True) - lse = (sum_s + 1e-10).log() + max_s - p = exp_s / sum_s - o = torch.matmul(p, v.float()) - return o.to(torch.bfloat16), lse.squeeze(-1) - - -def kv_merge(o_segments, lse_segments): - """Merge attention results from multiple KV segments. - - Uses the D5 merge formula: - O = sum_i [exp(lse_i) * O_i] / sum_i [exp(lse_i)] - - Args: - o_segments: list of (M, hd) BF16 tensors (normalized outputs) - lse_segments: list of (M,) FP32 tensors (log-sum-exp values) - - Returns: - o_merged: (M, hd) BF16 - """ - # Stack LSEs: (M, num_segments) - lse_stack = torch.stack(lse_segments, dim=-1) # (M, S) - # Max LSE for numerical stability - max_lse = lse_stack.max(dim=-1).values # (M,) - # Weights: exp(lse - max_lse) - weights = (lse_stack - max_lse.unsqueeze(-1)).exp() # (M, S) - # Normalize weights - weight_sum = weights.sum(dim=-1, keepdim=True) # (M, 1) - norm_weights = weights / weight_sum # (M, S) - - # Weighted sum of outputs - o_merged = torch.zeros_like(o_segments[0]) - for i, o_i in enumerate(o_segments): - o_merged += norm_weights[:, i:i+1] * o_i.float() - - return o_merged.to(torch.bfloat16) - - -def _run_fmha_segment(q_3d, k_3d, v, m, s_k, hd, use_smem_p=False): - """Run FMHA for a single KV segment.""" - scale = 1.0 / math.sqrt(hd) - kernel = FmhaKernel(head_dim=hd, s_k=s_k, use_smem_p=use_smem_p) - pv_n_tile = kernel.pv_n_tile - n_pv_tiles = kernel.n_pv_tiles - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - v_tile = v[:, 0:pv_n_tile].contiguous().unsqueeze(-1) - c_tile = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - lse_tensor = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - - mQ = ct.from_dlpack(q_3d).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q_3d)) - mK = ct.from_dlpack(k_3d).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k_3d)) - mV = ct.from_dlpack(v_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_tile)) - mC = ct.from_dlpack(c_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) - - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE) - - o_unnorm = torch.zeros(m, hd, dtype=torch.float32, device='cuda') - for pv in range(n_pv_tiles): - v_tile = v[:, pv*pv_n_tile:(pv+1)*pv_n_tile].contiguous().unsqueeze(-1) - c_tile.zero_() - lse_tensor.zero_() - - mV = ct.from_dlpack(v_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_tile)) - mC = ct.from_dlpack(c_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) - - compiled(mQ, mK, mV, mC, stream, mLSE) - o_unnorm[:, pv*pv_n_tile:(pv+1*pv_n_tile)] = c_tile[:,:,0].float() - - # Use reference normalization (kernel LSE per-row not fully working) - q_flat = q_3d[:,:,0] - k_flat = k_3d[:,:,0] - v_flat = v # (s_k, hd) - scores = torch.matmul(q_flat.float(), k_flat.float().T) * scale - max_s = scores.max(dim=-1, keepdim=True).values - exp_s = (scores - max_s).exp() - attn_sum = exp_s.sum(dim=-1, keepdim=True) - lse = (attn_sum + 1e-10).log() + max_s # (M, 1) - - o_norm = (o_unnorm / attn_sum).to(torch.bfloat16) - return o_norm, lse.squeeze(-1) - - -def test_d15_s256(): - """s_k=256 (2 KV tiles): merge two segments.""" - print("\n=== Test 1: s_k=256 (2 KV tiles, hd=64) ===") - torch.manual_seed(42) - m, s_k, hd = 128, 256, 64 - scale = 1.0 / math.sqrt(hd) - - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - - # Split K/V into segments of 128 - segment_size = 128 - n_segments = s_k // segment_size - - o_segments = [] - lse_segments = [] - for seg in range(n_segments): - k_seg = k[seg*segment_size:(seg+1)*segment_size].unsqueeze(-1) - v_seg = v[seg*segment_size:(seg+1)*segment_size] - o_seg, lse_seg = _run_fmha_segment(q, k_seg, v_seg, m, segment_size, hd) - o_segments.append(o_seg) - lse_segments.append(lse_seg) - - o_merged = kv_merge(o_segments, lse_segments) - - # Reference - ref, _ = reference_attention(q[:,:,0], k, v, scale) - cos = torch.nn.functional.cosine_similarity( - o_merged.flatten().float().unsqueeze(0), ref.flatten().float().unsqueeze(0) - ).item() - print(f" cos = {cos:.6f}") - assert cos >= 0.995, f"cosine too low: {cos}" - print(" ✅ PASS") - - -def test_d15_s512(): - """s_k=512 (4 KV tiles): Flash decode config.""" - print("\n=== Test 2: s_k=512 (4 KV tiles, hd=64) ===") - torch.manual_seed(42) - m, s_k, hd = 128, 512, 64 - scale = 1.0 / math.sqrt(hd) - - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - - segment_size = 128 - n_segments = s_k // segment_size - - o_segments = [] - lse_segments = [] - for seg in range(n_segments): - k_seg = k[seg*segment_size:(seg+1)*segment_size].unsqueeze(-1) - v_seg = v[seg*segment_size:(seg+1)*segment_size] - o_seg, lse_seg = _run_fmha_segment(q, k_seg, v_seg, m, segment_size, hd) - o_segments.append(o_seg) - lse_segments.append(lse_seg) - - o_merged = kv_merge(o_segments, lse_segments) - - ref, _ = reference_attention(q[:,:,0], k, v, scale) - cos = torch.nn.functional.cosine_similarity( - o_merged.flatten().float().unsqueeze(0), ref.flatten().float().unsqueeze(0) - ).item() - print(f" cos = {cos:.6f}") - assert cos >= 0.995, f"cosine too low: {cos}" - print(" ✅ PASS") - - -def test(): - print("=== D1.5: Multi-KV-Tile Attention with Python KV Merge ===") - test_d15_s256() - test_d15_s512() - print("\n=== ALL TESTS PASSED ===") - - -if __name__ == '__main__': - test() diff --git a/tests/unit/test_d15_noop_rescale.py b/tests/unit/test_d15_noop_rescale.py deleted file mode 100644 index 3dcc44fa..00000000 --- a/tests/unit/test_d15_noop_rescale.py +++ /dev/null @@ -1,77 +0,0 @@ -""" -D1.5 Debug: NO-OP TMEM round-trip test. -Tests s_k=256 with rescale_factor=1.0 (NO-OP). -If the round-trip itself is broken, even NO-OP will corrupt O. -If it produces the same (wrong) result as without rescale, the round-trip works. -""" -import torch, math -import cutlass.cute as cute -import cutlass.torch as ct -import cuda.bindings.driver as cuda -from dsv4.kernels.attention.fmha import FmhaKernel - - -def reference_attention(q, k, v, scale): - qf = q.float(); kf = k.float() - attn = qf @ kf.T * scale - attn_max = attn.max(dim=-1, keepdim=True)[0] - attn_exp = torch.exp(attn - attn_max) - return attn_exp @ v.float() - - -def run_test(s_k, debug_noop=False, label=""): - hd = 64; m = 128; scale = 1.0 / math.sqrt(hd) - torch.manual_seed(42) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - - kernel = FmhaKernel(head_dim=hd, s_k=s_k, use_smem_p=False, normalize=False, debug_noop_rescale=debug_noop) - pv_n_tile = kernel.pv_n_tile - c = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - lse = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - rs = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - - v_t = v[:, 0:pv_n_tile].contiguous().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_t).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_t)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - mLSE = ct.from_dlpack(lse).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse)) - mRS = ct.from_dlpack(rs).mark_layout_dynamic(leading_dim=ct.get_leading_dim(rs)) - - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE, row_sums=mRS) - compiled(mQ, mK, mV, mC, stream, mLSE, row_sums=mRS) - torch.cuda.synchronize() - - ref = reference_attention(q[:, :, 0], k[:, :, 0], v, scale) - out = c[:, :, 0].float() - cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - print(f'{label}: cos={cos:.6f} {"PASS" if cos >= 0.999 else "FAIL"}', flush=True) - return out - - -def test(): - # Test 1: s_k=128 baseline - run_test(128, label="s_k=128 baseline") - - # Test 2: s_k=256 WITH rescale (should be correct if TMEM round-trip works) - run_test(256, label="s_k=256 with rescale") - - # Test 3: s_k=256 NOOP rescale (acc_scale=1.0) - # This should produce the same result as s_k=256 WITHOUT any rescale - # (which is: O = P[0]*V[0] + P[1]*V[1] with no O rescale — mathematically wrong - # but should be stable if TMEM round-trip doesn't corrupt) - out_noop = run_test(256, debug_noop=True, label="s_k=256 NOOP rescale") - - # Test 4: s_k=256 WITHOUT any rescale (old code path) - # Compare with NOOP to see if TMEM round-trip itself corrupts - # We can't easily disable the rescale in the current code, - # but NOOP rescale with factor=1.0 is equivalent to a successful round-trip - # followed by multiply-by-1. If the output matches a "no-rescale" baseline, - # the TMEM round-trip is working correctly. - - -if __name__ == '__main__': - test() diff --git a/tests/unit/test_d15_rescale_debug.py b/tests/unit/test_d15_rescale_debug.py deleted file mode 100644 index 583a1690..00000000 --- a/tests/unit/test_d15_rescale_debug.py +++ /dev/null @@ -1,116 +0,0 @@ -""" -D1.5 Debug: Test s_k=256 in-kernel rescale with diagnostics. -Minimal test to isolate the TMEM round-trip vs barrier issue. -""" -import torch, math -import cutlass.cute as cute -import cutlass.torch as ct -import cuda.bindings.driver as cuda -from dsv4.kernels.attention.fmha import FmhaKernel - - -def reference_attention(q, k, v, scale): - qf = q.float(); kf = k.float() - attn = qf @ kf.T * scale - attn_max = attn.max(dim=-1, keepdim=True)[0] - attn_exp = torch.exp(attn - attn_max) - return attn_exp @ v.float() - - -def test(): - hd = 64; m = 128; scale = 1.0 / math.sqrt(hd) - torch.manual_seed(42) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - - # Test 1: s_k=128 baseline - k1 = torch.randn(128, hd, 1, dtype=torch.bfloat16, device='cuda') - v1 = torch.randn(128, hd, dtype=torch.bfloat16, device='cuda') - c1 = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - lse1 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - rs1 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - kernel1 = FmhaKernel(head_dim=hd, s_k=128, use_smem_p=False, normalize=False) - pv_n_tile = kernel1.pv_n_tile - v1t = v1[:, 0:pv_n_tile].contiguous().unsqueeze(-1) - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK1 = ct.from_dlpack(k1).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k1)) - mV1 = ct.from_dlpack(v1t).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v1t)) - mC1 = ct.from_dlpack(c1).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c1)) - mLSE1 = ct.from_dlpack(lse1).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse1)) - mRS1 = ct.from_dlpack(rs1).mark_layout_dynamic(leading_dim=ct.get_leading_dim(rs1)) - compiled1 = cute.compile(kernel1, mQ, mK1, mV1, mC1, stream, mLSE1, row_sums=mRS1) - compiled1(mQ, mK1, mV1, mC1, stream, mLSE1, row_sums=mRS1) - torch.cuda.synchronize() - ref1 = reference_attention(q[:, :, 0], k1[:, :, 0], v1, scale) - cos1 = torch.nn.functional.cosine_similarity(c1[:, :, 0].float().flatten().unsqueeze(0), ref1.flatten().unsqueeze(0)).item() - print(f's_k=128 baseline: cos={cos1:.6f} {"PASS" if cos1 >= 0.999 else "FAIL"}', flush=True) - - # Test 2: s_k=256 with in-kernel rescale - k2 = torch.randn(256, hd, 1, dtype=torch.bfloat16, device='cuda') - v2 = torch.randn(256, hd, dtype=torch.bfloat16, device='cuda') - c2 = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - lse2 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - rs2 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - kernel2 = FmhaKernel(head_dim=hd, s_k=256, use_smem_p=False, normalize=False) - v2t = v2[:, 0:pv_n_tile].contiguous().unsqueeze(-1) - mK2 = ct.from_dlpack(k2).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k2)) - mV2 = ct.from_dlpack(v2t).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v2t)) - mC2 = ct.from_dlpack(c2).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c2)) - mLSE2 = ct.from_dlpack(lse2).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse2)) - mRS2 = ct.from_dlpack(rs2).mark_layout_dynamic(leading_dim=ct.get_leading_dim(rs2)) - compiled2 = cute.compile(kernel2, mQ, mK2, mV2, mC2, stream, mLSE2, row_sums=mRS2) - compiled2(mQ, mK2, mV2, mC2, stream, mLSE2, row_sums=mRS2) - torch.cuda.synchronize() - - ref2 = reference_attention(q[:, :, 0], k2[:, :, 0], v2, scale) - out2 = c2[:, :, 0].float() - cos2 = torch.nn.functional.cosine_similarity(out2.flatten().unsqueeze(0), ref2.flatten().unsqueeze(0)).item() - - # Per-element stats - diff2 = (out2 - ref2).abs() - max_rel_err = (diff2 / ref2.abs().clamp(min=1e-6)).max().item() - print(f's_k=256 in-kernel rescale: cos={cos2:.6f} max_rel_err={max_rel_err:.4f} {"PASS" if cos2 >= 0.999 else "FAIL"}', flush=True) - - # Print LSE and row_sums for debugging - print(f' LSE range: [{lse2[:, 0, 0].min().item():.4f}, {lse2[:, 0, 0].max().item():.4f}]', flush=True) - print(f' row_sums range: [{rs2[:, 0, 0].min().item():.4f}, {rs2[:, 0, 0].max().item():.4f}]', flush=True) - - # Test 3: Python KV merge for comparison - c_s0 = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - lse_s0 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - rs_s0 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - v2_0 = v2[:128, 0:pv_n_tile].contiguous().unsqueeze(-1) - mK2_0 = ct.from_dlpack(k2[:128]).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k2[:128])) - mV2_0 = ct.from_dlpack(v2_0).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v2_0)) - mC_s0 = ct.from_dlpack(c_s0).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_s0)) - mLSE_s0 = ct.from_dlpack(lse_s0).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_s0)) - mRS_s0 = ct.from_dlpack(rs_s0).mark_layout_dynamic(leading_dim=ct.get_leading_dim(rs_s0)) - compiled_s0 = cute.compile(kernel1, mQ, mK2_0, mV2_0, mC_s0, stream, mLSE_s0, row_sums=mRS_s0) - compiled_s0(mQ, mK2_0, mV2_0, mC_s0, stream, mLSE_s0, row_sums=mRS_s0) - - c_s1 = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - lse_s1 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - rs_s1 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - v2_1 = v2[128:, 0:pv_n_tile].contiguous().unsqueeze(-1) - mK2_1 = ct.from_dlpack(k2[128:]).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k2[128:])) - mV2_1 = ct.from_dlpack(v2_1).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v2_1)) - mC_s1 = ct.from_dlpack(c_s1).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_s1)) - mLSE_s1 = ct.from_dlpack(lse_s1).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_s1)) - mRS_s1 = ct.from_dlpack(rs_s1).mark_layout_dynamic(leading_dim=ct.get_leading_dim(rs_s1)) - compiled_s1 = cute.compile(kernel1, mQ, mK2_1, mV2_1, mC_s1, stream, mLSE_s1, row_sums=mRS_s1) - compiled_s1(mQ, mK2_1, mV2_1, mC_s1, stream, mLSE_s1, row_sums=mRS_s1) - torch.cuda.synchronize() - - o0 = c_s0[:, :, 0].float(); o1 = c_s1[:, :, 0].float() - r0 = rs_s0[:, 0, 0].float(); r1 = rs_s1[:, 0, 0].float() - l0 = lse_s0[:, 0, 0].float(); l1 = lse_s1[:, 0, 0].float() - o0_norm = o0 / r0.unsqueeze(1).clamp(min=1e-30) - o1_norm = o1 / r1.unsqueeze(1).clamp(min=1e-30) - w0 = torch.exp(l0).unsqueeze(1); w1 = torch.exp(l1).unsqueeze(1) - oracle = (w0 * o0_norm + w1 * o1_norm) / (w0 + w1) - cos_oracle = torch.nn.functional.cosine_similarity(oracle.flatten().unsqueeze(0), ref2.flatten().unsqueeze(0)).item() - print(f'Python KV merge: cos={cos_oracle:.6f}', flush=True) - - -if __name__ == '__main__': - test() diff --git a/tests/unit/test_d15_roundtrip_iso.py b/tests/unit/test_d15_roundtrip_iso.py deleted file mode 100644 index c919a4b3..00000000 --- a/tests/unit/test_d15_roundtrip_iso.py +++ /dev/null @@ -1,108 +0,0 @@ -""" -D1.5 Minimal TMEM round-trip test within FMHA. - -Tests: run FMHA with s_k=128, then add a NOOP correction_rescale -(multiply by 1.0) BETWEEN the softmax and the epilogue. -If the output is bitwise identical to without rescale, the round-trip works. -If it differs, the round-trip corrupts data. - -This test directly compares: same kernel, same data, WITH and WITHOUT -the O rescale step (forced to NOOP). -""" -import torch, math -import cutlass.cute as cute -import cutlass.torch as ct -import cuda.bindings.driver as cuda -from dsv4.kernels.attention.fmha import FmhaKernel - - -def test(): - hd = 64; m = 128; s_k = 128; scale = 1.0 / math.sqrt(hd) - torch.manual_seed(42) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - - # Run WITHOUT rescale (s_k=128, n_kv_tiles=1, rescale code is dead) - kernel1 = FmhaKernel(head_dim=hd, s_k=s_k, use_smem_p=False, normalize=False) - pv_n_tile = kernel1.pv_n_tile - c1 = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - lse1 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - rs1 = torch.zeros(m, 1, 1, dtype=torch.float32, 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)) - v_t = v[:, 0:pv_n_tile].contiguous().unsqueeze(-1) - mV = ct.from_dlpack(v_t).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_t)) - mC1 = ct.from_dlpack(c1).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c1)) - mLSE1 = ct.from_dlpack(lse1).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse1)) - mRS1 = ct.from_dlpack(rs1).mark_layout_dynamic(leading_dim=ct.get_leading_dim(rs1)) - - print('Compiling s_k=128 (no rescale)...', flush=True) - comp1 = cute.compile(kernel1, mQ, mK, mV, mC1, stream, mLSE1, row_sums=mRS1) - comp1(mQ, mK, mV, mC1, stream, mLSE1, row_sums=mRS1) - torch.cuda.synchronize() - - # Run WITH NOOP rescale (s_k=256, but with debug_noop_rescale=True, - # using same K/V repeated to make the second segment identical) - # This triggers the O rescale code path (n_kv_tiles=2) but with factor=1.0 - k2 = torch.cat([k, k], dim=0) - v2 = torch.cat([v, v], dim=0) - - kernel2 = FmhaKernel(head_dim=hd, s_k=256, use_smem_p=False, normalize=False, debug_noop_rescale=True) - c2 = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - lse2 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - rs2 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - - mK2 = ct.from_dlpack(k2).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k2)) - v2_t = v2[:, 0:pv_n_tile].contiguous().unsqueeze(-1) - mV2 = ct.from_dlpack(v2_t).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v2_t)) - mC2 = ct.from_dlpack(c2).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c2)) - mLSE2 = ct.from_dlpack(lse2).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse2)) - mRS2 = ct.from_dlpack(rs2).mark_layout_dynamic(leading_dim=ct.get_leading_dim(rs2)) - - print('Compiling s_k=256 (NOOP rescale)...', flush=True) - comp2 = cute.compile(kernel2, mQ, mK2, mV2, mC2, stream, mLSE2, row_sums=mRS2) - comp2(mQ, mK2, mV2, mC2, stream, mLSE2, row_sums=mRS2) - torch.cuda.synchronize() - - # FP32 reference - qf = q[:, :, 0].float() - kf = k[:, :, 0].float() - attn = qf @ kf.T * scale - attn_max = attn.max(dim=-1, keepdim=True)[0] - attn_exp = torch.exp(attn - attn_max) - ref_norm = (attn_exp @ v.float()) / attn_exp.sum(dim=-1, keepdim=True) - - # Normalize both outputs - out1 = c1[:, :, 0].float() / rs1[:, 0, 0].float().unsqueeze(1).clamp(min=1e-30) - out2 = c2[:, :, 0].float() / rs2[:, 0, 0].float().unsqueeze(1).clamp(min=1e-30) - - cos1 = torch.nn.functional.cosine_similarity(out1.flatten().unsqueeze(0), ref_norm.flatten().unsqueeze(0)).item() - cos2 = torch.nn.functional.cosine_similarity(out2.flatten().unsqueeze(0), ref_norm.flatten().unsqueeze(0)).item() - cos_12 = torch.nn.functional.cosine_similarity(out1.flatten().unsqueeze(0), out2.flatten().unsqueeze(0)).item() - - # Also check unnormalized - unnorm1 = c1[:, :, 0].float() - unnorm2 = c2[:, :, 0].float() - # With identical segments, s_k=256 unnorm should be 2x s_k=128 unnorm - # (row_sum also doubles, so normalized result is the same) - ratio = unnorm2 / unnorm1.clamp(min=1e-10) - print(f'Unnorm ratio (should be ~2.0): mean={ratio.mean().item():.4f} std={ratio.std().item():.4f}') - - print(f's_k=128 (no rescale): cos={cos1:.6f}') - print(f's_k=256 (NOOP rescale): cos={cos2:.6f}') - print(f's_k=128 vs s_k=256: cos={cos_12:.6f}') - - # Bitwise comparison on unnormalized O - # s_k=256 with identical segments: O_unnorm = 2 * O_unnorm_128 - expected_unnorm2 = 2.0 * unnorm1 - bit_cos = torch.nn.functional.cosine_similarity(unnorm2.flatten().unsqueeze(0), expected_unnorm2.flatten().unsqueeze(0)).item() - max_diff = (unnorm2 - expected_unnorm2).abs().max().item() - print(f'Unnorm: expected 2x, got cos={bit_cos:.6f} max_diff={max_diff:.4f}') - - -if __name__ == '__main__': - test() diff --git a/tests/unit/test_d1_3_cotiled.py b/tests/unit/test_d1_3_cotiled.py deleted file mode 100644 index 2e22c92e..00000000 --- a/tests/unit/test_d1_3_cotiled.py +++ /dev/null @@ -1,316 +0,0 @@ -""" -D1.3 SMEM-P: Diagnostic for make_cotiled_copy approach. - -Runs a minimal kernel that: -1. Prints tiled_tmem_load TV layout shapes -2. Prints tTMEM_LOADcS coordinate partition shapes -3. Prints sP layout shapes -4. Attempts to build make_cotiled_copy and prints partition shapes -""" -import torch, math -import cutlass, cutlass.cute as cute, cutlass.utils as utils -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, const_expr -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 - - -class SmemPDiag: - def __init__(self, head_dim=64, s_k=128): - self.head_dim = head_dim - self.s_k = s_k - self.pv_n_tile = min(head_dim, 256) - self.use_smem_p = True # We're diagnosing SMEM-P - self.normalize = False - self.qk_mma_tiler = (128, 128, 128 * 4) - self.pv_mma_tiler = (128, self.pv_n_tile, 128) - self.acc_dtype = Float32 - self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16 - self.o_dtype = BFloat16 - self.use_2cta_instrs = False - self.cta_group = tcgen05.CtaGroup.ONE - self.cluster_shape_mn = (1, 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.num_acc_stage = 1 - self.scale_softmax_log2 = 1.0 / math.sqrt(self.head_dim) * math.log2(math.e) - self.kv_stage = 2 - self.q_stage = 1 - self.num_c_stage = 2 - self.c_layout = LayoutEnum.ROW_MAJOR - - @cute.jit - def __call__(self, q, k, v, c, stream): - 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( - (self.pv_n_tile, self.s_k, 1), - stride=(1, self.pv_n_tile, self.pv_n_tile * self.s_k), - ), - ) - self.v_major = LayoutEnum.from_tensor(v_fmha).mma_major_mode() - - 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_a_major = cute.nvgpu.OperandMajorMode.K # SMEM-P uses K major - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, pv_a_major, self.v_major, - self.qk_acc_dtype, self.cta_group, (128, self.pv_n_tile), - tcgen05.OperandSource.SMEM - ) - - # Setup TMEM offsets (simplified for SMEM-P) - 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 = -1 # SMEM-P - self.tmem_o0_offset = 0 - s_cols = self.qk_mma_tiler[1] - o_cols = find_tmem_tensor_col_offset(tOtO) - total = max(s_cols, o_cols) - _n = 1 - while _n < total: - _n *= 2 - self.num_tmem_alloc_cols = _n - self.tOrP0_offset = 0 # SMEM-P - - # Build SMEM layouts - p_smem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, self.q_dtype, self.q_stage) - k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, self.q_dtype, self.kv_stage) - v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, self.kv_stage) - c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, (128, self.pv_n_tile), 2) - - # TMA atoms - 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)) - cta = cute.size(qk_mma.thr_id.shape) - 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 - - 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, (1, 1, 1, 1) - ) - 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, (1, 1, 1, 1) - ) - 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, (1, 1, 1, 1) - ) - epi_s = cute.select(c_smem_s, mode=[0, 1]) - tma_c, mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_s, (128, self.pv_n_tile)) - - # ===== KERNEL ===== - self._kernel(qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, - p_smem_s, q_smem_s, k_smem_s, v_smem_s, c_smem_s, - tStS, tOtO, qk_thr, pv_thr).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, - p_smem_s, q_smem_s, k_smem_s, v_smem_s, c_smem_s, - tStS, tOtO, qk_thr, pv_thr): - tidx, _, _ = cute.arch.thread_idx() - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - - @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) - - 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) - sP = smem.allocate_tensor(element_type=self.q_dtype, layout=p_smem_s.outer, byte_alignment=128, swizzle=p_smem_s.inner) - - # ===== PRINT DIAGNOSTICS (only from warp 0, thread 0) ===== - if warp_idx < 4: # softmax warps - # TMEM load 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, tStS) - - sfw_idx = tidx % (32 * len(self.epilogue_warp_id)) - thr_load = tiled_tmem_load.get_slice(sfw_idx) - tTMEM_LOADtS = thr_load.partition_S(tStS) - - 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) - - if sfw_idx == 0: - print(f"=== TMEM Load Diagnostics (sfw_idx=0) ===") - print(f"tiled_tmem_load shape: {cute.shape(tiled_tmem_load)}") - print(f"tTMEM_LOADtS shape: {cute.shape(tTMEM_LOADtS)}") - print(f"tTMEM_LOADtS layout: {tTMEM_LOADtS.layout}") - print(f"tTMEM_LOADcS shape: {cute.shape(tTMEM_LOADcS)}") - print(f"tTMEM_LOADcS layout: {tTMEM_LOADcS.layout}") - - # Print TV layout - tv = tiled_tmem_load.layout_dst_tv_tiled - print(f"TV layout: {tv}") - print(f"TV layout shape: {cute.shape(tv)}") - - # Print sP layout - sP_stage = sP[(None, None, None, 0)] - print(f"sP shape: {cute.shape(sP)}") - print(f"sP layout (outer): {sP.layout}") - print(f"sP_stage shape: {cute.shape(sP_stage)}") - print(f"sP_stage layout: {sP_stage.layout}") - - # Print tStS layout - print(f"tStS shape: {cute.shape(tStS)}") - print(f"tStS layout: {tStS.layout}") - - # Try to build make_cotiled_copy - # atom_layout_tv: (tid, vid) -> sP flat address - # We need to map from the TV layout's codomain (tStS addresses) - # to sP addresses. - # - # tStS layout: ((128,128),1,1):((65536,1),0,0) - # So tStS_addr(m,k) = m*65536 + k - # sP layout: ((128,16),1,(4,2),1):(((64,1),0,((16,8192),0) - # So sP_addr((m,k0),0,(k1,k2),0) = m*64 + k0 + k1*16 + k2*8192 - # where k0=k%16, k1=(k//16)%4, k2=k//64 - # - # We can't directly compose because the codomains are different. - # But we CAN build a new Layout that maps (tid, vid) -> sP addr. - # - # The key: enumerate all (tid, vid) pairs, find their (m, k) from - # tTMEM_LOADcS, then compute sP_addr. - # But we can't do this at Python trace time inside @cute.kernel - # because we don't have the actual coordinate values. - # - # OR: we can try composition. Let's see if - # cute.composition(sP_stage.layout, tv) works. - # This requires the codomain of tv to be compatible with - # the domain of sP_stage.layout. - print(f"\n=== Attempting composition ===") - print(f"TV codomain size: {cute.size(tv)}") - print(f"sP_stage domain size: {cute.size(sP_stage.layout)}") - - # Try: make a layout that maps (m, k) -> sP_addr - # then compose with the inverse of tStS.layout - # Actually, tStS has 3 modes. Let me flatten it. - # The 2D logical P is (128, 128) = (M, K) - # tStS has layout ((128,128),1,1):((65536,1),0,0) - # The 2D part is (128,128) with strides (65536, 1) - # So the 2D layout is: (128, 128) : (65536, 1) - - # sP_stage has shape ((128,16),1,(4,2),1) - # The logical P is 128 x 128 where: - # mode0: (128, 16) with strides (64, 1) - # mode2: (4, 2) with strides (16, 8192) - # So the 2D logical layout (grouping modes 0 and 2) is: - # (128*4, 16*2) = (512, 32) with... complex strides. - # Actually, sP is 4D, not 2D. - - # Let me try the group_modes approach: - sP_2d = cute.group_modes(sP_stage, 0, 4) # flatten to 1D? No. - # group_modes(sP, 0, 4) makes it (N,) where N = total elements - # Actually group_modes(sP, 0, 4) groups the first 4 modes into 1 - # and keeps the rest. sP_stage has 4 modes -> becomes 1 mode. - # That's just a 1D tensor. - print(f"sP_2d (grouped) shape: {cute.shape(sP_2d)}") - print(f"sP_2d layout: {sP_2d.layout}") - - # Now try to compose the TV layout with the identity tensor's layout - # to get (tid, vid) -> (m, k), then use that to index into sP. - - # Actually, the right approach per the CUTLASS LLM: - # 1. Build a custom atom_layout_tv from scratch - # 2. Each thread's (m, k) pairs come from tTMEM_LOADcS - # 3. Convert (m, k) to sP address - # 4. Build a Layout((n_threads, n_vals_per_thread), (sP_addr_strides,)) - # - # But building this at trace time requires knowing all (m, k) values. - # The identity tensor partition tTMEM_LOADcS has these values - # but they're runtime values in CuTeDSL, not Python values. - # - # HOWEVER: since tTMEM_LOADcS is an identity tensor partition, - # the coordinates are STATIC (known at compile/trace time). - # We should be able to extract them at Python trace time. - # - # Let me check: can we iterate over tTMEM_LOADcS at trace time? - # tTMEM_LOADcS has shape ((32,1),4,1,1) - # Each element is a (m, k) pair. - # At trace time, cute.shape gives us the static shape. - # But the actual coordinate VALUES might be dynamic. - # - # In CuTe, identity tensors have static values. - # cute.make_identity_tensor((M, K)) creates a tensor where - # the value at each coordinate IS the coordinate. - # So tTMEM_LOADcS[j0, j1, 0, 0] should give a static (m, k) pair. - # At trace time, we should be able to extract these. - # - # But wait - at trace time, the identity tensor is being traced. - # The values are MLIR constants, not Python ints. - # We can't iterate and build a Python Layout from them. - # - # The alternative: define atom_layout_tv analytically. - # The Ld32x32bOp with Repetition(32) partitions 128 threads - # over a 128x128 matrix. We need to figure out the exact - # (tid, vid) -> (m, k) mapping analytically. - - # Let me print what we know about the TMEM load's thread mapping. - # The layout_dst_tv_tiled is a Layout with shape (n_threads, n_vals). - print(f"\n=== Analyzing TV layout for analytical construction ===") - print(f"TV layout type: {type(tv)}") - # Try to decompose it - if hasattr(tv, 'shape'): - print(f" shape: {tv.shape}") - if hasattr(tv, 'stride'): - print(f" stride: {tv.stride}") - - -def test_cotiled_diag(): - print("=== make_cotiled_copy Diagnostic ===\n") - hd = 64 - q = torch.randn(128, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(128, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(128, min(hd, 256), dtype=torch.bfloat16, device='cuda') - c = torch.zeros(128, min(hd, 256), 1, dtype=torch.bfloat16, device='cuda') - - diag = SmemPDiag(head_dim=hd, s_k=128) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - v_tile = 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_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_tile)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - - print('Compiling diagnostic kernel...', flush=True) - compiled = cute.compile(diag, mQ, mK, mV, mC, stream) - print('Running...', flush=True) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - print('Done.') - - -if __name__ == '__main__': - test_cotiled_diag() diff --git a/tests/unit/test_d1_3_layout_diag.py b/tests/unit/test_d1_3_layout_diag.py deleted file mode 100644 index b23874f7..00000000 --- a/tests/unit/test_d1_3_layout_diag.py +++ /dev/null @@ -1,197 +0,0 @@ -""" -D1.3 SMEM-P diagnostic: print ALL relevant shapes and layouts. -Runs layout analysis inside a JIT-compiled function to get MLIR context. -""" -import torch, math, sys -import cutlass, cutlass.cute as cute, cutlass.utils as utils -from cutlass.cute.nvgpu import tcgen05, cpasync -from cutlass import Float32, BFloat16, Int32, const_expr -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cutlass.torch as ct - - -def main(): - for hd in [64, 256]: - print(f"\n{'='*60}") - print(f" HEAD_DIM={hd}") - print(f"{'='*60}") - - m = 128; s_k = 128 - use_smem_p = hd > 64 - pv_n_tile = min(hd, 256) - - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, pv_n_tile, dtype=torch.bfloat16, device='cuda') - c = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - - # Use FmhaKernel's actual __call__ path to compute layouts - from dsv4.kernels.attention.fmha import FmhaKernel - kern = FmhaKernel(head_dim=hd, s_k=s_k) - - # Trigger _setup by calling compile with a dummy kernel - # Actually, we can access the _setup by calling __call__ on the kernel - # But that runs the full kernel. Let me manually call _setup by - # building the MMAs the same way __call__ does. - - a_major = LayoutEnum.from_tensor(ct.from_dlpack(q)).mma_major_mode() - b_major = LayoutEnum.from_tensor(ct.from_dlpack(k)).mma_major_mode() - - # Create a minimal analysis kernel that just prints shapes - @cute.kernel - def _diag_kernel(_q, _k, _v, _c): - # Build QK MMA - qk_mma = utils.sm100.make_trivial_tiled_mma( - BFloat16, BFloat16, a_major, b_major, Float32, - tcgen05.CtaGroup.ONE, (128, 128), tcgen05.OperandSource.SMEM, - ) - - pv_a_major = a_major if use_smem_p else cute.nvgpu.OperandMajorMode.K - pv_source = tcgen05.OperandSource.SMEM if use_smem_p else tcgen05.OperandSource.TMEM - c_major = LayoutEnum.from_tensor(ct.from_dlpack(c)).mma_major_mode() - - pv_mma = utils.sm100.make_trivial_tiled_mma( - BFloat16, BFloat16, pv_a_major, c_major, Float32, - tcgen05.CtaGroup.ONE, (128, pv_n_tile), pv_source, - ) - - 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, pv_n_tile, pv_ik * (128 // pv_ik)) - - print(f" QK MMA shape_mnk: {qk_mma.shape_mnk}") - print(f" PV MMA shape_mnk: {pv_mma.shape_mnk}") - print(f" QK tiler: {qk_mma_tiler}") - print(f" PV tiler: {pv_mma_tiler}") - print(f" use_smem_p: {use_smem_p}") - - # QK C-fragment - qk_thr = qk_mma.get_slice(0) - qk_as = qk_thr.partition_shape_C(qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - print(f" tStS shape: {cute.shape(tStS)} layout: {tStS.layout}") - - # PV C-fragment - pv_thr = pv_mma.get_slice(0) - pv_as = pv_thr.partition_shape_C(pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - print(f" tOtO shape: {cute.shape(tOtO)} layout: {tOtO.layout}") - - # PV A-operand SMEM layout (sP) - p_smem_s = utils.sm100.make_smem_layout_a(pv_mma, pv_mma_tiler, BFloat16, 1) - # p_smem_s.outer is the layout. Use it directly for shape queries. - # Can't actually allocate SMEM here (that's the SmemAllocator's job in the kernel) - # But we can create a tensor with a dummy pointer for shape analysis. - sP_alloc = cute.make_tensor(0, p_smem_s.outer) - print(f" sP shape: {cute.shape(sP_alloc)}") - sP_2d = cute.group_modes(sP_alloc, 0, 3) - print(f" sP_2d shape: {cute.shape(sP_2d)} rank={len(cute.shape(sP_2d))}") - print(f" sP_2d layout: {sP_2d.layout}") - print(f" sP_2d layout: {sP_2d.layout}") - - # PV A-operand fragments — can't create with dummy ptr 0 (no memspace) - # Just print the layout info from the SMEM layout - print(f" p_smem_s.outer (sP layout): {p_smem_s.outer}") - print(f" p_smem_s.inner (swizzle): {p_smem_s.inner}") - - # What we need: the PV MMA's A-operand thread partition for sP - # This tells us which threads read which sP elements during PV GEMM - # The softmax warps must WRITE to sP using the same mapping - # (so that the MMA warp can READ using its own partition) - - # Softmax TMEM load partition - sfw_idx = 0 - 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(sfw_idx) - tTMEM_LOADtS = thr_load.partition_S(tStS) - cS = cute.make_identity_tensor((128, 128)) - tScS = qk_thr.partition_C(cS) - tTMEM_LOADcS = thr_load.partition_D(tScS) - print(f" tTMEM_LOADcS (coord) shape: {cute.shape(tTMEM_LOADcS)} layout: {tTMEM_LOADcS.layout}") - - # P sub-layout in TMEM - p_cols_fp32 = pv_mma_tiler[2] * BFloat16.width // Float32.width - print(f" p_cols_fp32: {p_cols_fp32}") - - # make_tiled_copy_C with qk_mma - print(f"\n --- make_tiled_copy_C(qk_mma) ---") - try: - smem_copy_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), BFloat16, num_bits_per_copy=128) - tiled_copy = cute.make_tiled_copy_C(smem_copy_atom, qk_mma) - thr_copy = tiled_copy.get_slice(sfw_idx) - print(f" OK, created tiled_copy") - tD = thr_copy.partition_D(sP_2d) - print(f" partition_D(sP_2d) shape: {cute.shape(tD)} rank={len(cute.shape(tD))}") - - # Create source tensor using the QK C-fragment register layout - # This is what make_tiled_copy_C expects as source - qk_C_reg = qk_thr.make_fragment_C(qk_as) # register fragment - # But we need a 2D version (without stage mode) - qk_C_2d = cute.group_modes(qk_C_reg, 0, 2) - print(f" qk_C_reg shape: {cute.shape(qk_C_reg)} layout: {qk_C_reg.layout}") - print(f" qk_C_2d shape: {cute.shape(qk_C_2d)} layout: {qk_C_2d.layout}") - - try: - tS = thr_copy.partition_S(qk_C_2d) - print(f" partition_S(qk_C_2d) shape: {cute.shape(tS)} rank={len(cute.shape(tS))}") - print(f" partition_S and partition_D RANKS MATCH: {len(cute.shape(tS)) == len(cute.shape(tD))}") - except Exception as e: - print(f" partition_S(qk_C_2d) FAILED: {e}") - # Try with a P-sized sub-tensor of the QK C-fragment - # P is only the first p_cols_fp32 columns of S - p_cols_bf16 = pv_mma_tiler[2] # K dimension of PV = columns of P - print(f" p_cols_bf16 (PV tiler[2]): {p_cols_bf16}") - qk_P_shape = (128, p_cols_bf16) # P matrix is (128, p_cols_bf16) - qk_P_layout = cute.composition(qk_C_reg.layout, cute.make_layout(qk_P_shape)) - qk_P = cute.make_tensor(0, qk_P_layout) - qk_P_2d = cute.group_modes(qk_P, 0, 2) - print(f" qk_P_2d shape: {cute.shape(qk_P_2d)} layout: {qk_P_2d.layout}") - try: - tS = thr_copy.partition_S(qk_P_2d) - print(f" partition_S(qk_P_2d) shape: {cute.shape(tS)} rank={len(cute.shape(tS))}") - print(f" RANKS MATCH: {len(cute.shape(tS)) == len(cute.shape(tD))}") - except Exception as e2: - print(f" partition_S(qk_P_2d) FAILED: {e2}") - except Exception as e: - print(f" make_tiled_copy_C FAILED: {e}") - - # make_tiled_copy_C with pv_mma - print(f"\n --- make_tiled_copy_C(pv_mma) ---") - try: - tiled_copy_pv = cute.make_tiled_copy_C(smem_copy_atom, pv_mma) - thr_copy_pv = tiled_copy_pv.get_slice(sfw_idx) - tD_pv = thr_copy_pv.partition_D(sP_2d) - print(f" OK, partition_D(sP_2d) shape: {cute.shape(tD_pv)} rank={len(cute.shape(tD_pv))}") - except Exception as e: - print(f" make_tiled_copy_C(pv_mma) FAILED: {e}") - - # KEY QUESTION: what is the relationship between QK C-fragment registers - # and the TMEM load partition? Are they the same threads, same values? - print(f"\n --- QK C-fragment vs TMEM load partition ---") - print(f" tStS shape: {cute.shape(tStS)} layout: {tStS.layout}") - print(f" tTMEM_LOADtS shape: {cute.shape(tTMEM_LOADtS)}") - print(f" tTMEM_LOADcS shape: {cute.shape(tTMEM_LOADcS)} layout: {tTMEM_LOADcS.layout}") - print(f" size(tStS) = {cute.size(tStS)}, size(tTMEM_LOADcS) = {cute.size(tTMEM_LOADcS)}") - # Do the C-fragment and TMEM load have the same per-thread element count? - # If so, the same threads own the same elements in both views. - - # Compile and run - import cuda.bindings.driver as cuda - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - 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.unsqueeze(-1)).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v.unsqueeze(-1))) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - - print(' Compiling diagnostic...', flush=True) - compiled = cute.compile(_diag_kernel, mQ, mK, mV, mC) - compiled(mQ, mK, mV, mC) - torch.cuda.synchronize() - - -if __name__ == '__main__': - main() diff --git a/tests/unit/test_d1_3_smem_diag.py b/tests/unit/test_d1_3_smem_diag.py deleted file mode 100644 index fc6134fd..00000000 --- a/tests/unit/test_d1_3_smem_diag.py +++ /dev/null @@ -1,70 +0,0 @@ -""" -D1.3 SMEM-P coordinate mapping diagnostic. -Verifies that tTMEM_LOADcS coordinates and rP_bf16 values -correctly map to sP indices for the SMEM-P path. -Runs a minimal kernel that writes P to sP and reads it back. -""" -import torch, math -import 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 -import cutlass.torch as ct -import cuda.bindings.driver as cuda - -# Test: write known P values to sP using coordinate indexing, then read back -# via the PV MMA's A-operand fragment and verify. - - -def test_smem_p_coords(): - print("=== SMEM-P Coordinate Diagnostic ===\n") - hd = 256; m = 128; s_k = 128; pv_n_tile = 256 - - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, pv_n_tile, dtype=torch.bfloat16, device='cuda') - c = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - - # Simple reference: just compute Q@K^T softmax @ V - qf = q[:, :, 0].float() - kf = k[:, :, 0].float() - vf = v.float() - scale = 1.0 / math.sqrt(hd) - attn = qf @ kf.T * scale - ref = torch.softmax(attn, dim=-1) @ vf # (128, 256) - - # Run the kernel with use_smem_p=True - from dsv4.kernels.attention.fmha import FmhaKernel - kern = FmhaKernel(head_dim=hd, s_k=s_k) - - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - v_tile = v.unsqueeze(-1) # (128, 256, 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_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_tile)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - - print('Compiling...', flush=True) - compiled = cute.compile(kern, 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() - - print(f'hd=256 SMEM-P: cos {cos:.6f} max_abs {max_abs:.4f}') - print(f' out[0,:4]={out[0,:4].tolist()}') - print(f' ref[0,:4]={ref[0,:4].tolist()}') - - # Also check: are the output values in a reasonable range? - print(f' out range: [{out.min().item():.4f}, {out.max().item():.4f}]') - print(f' ref range: [{ref.min().item():.4f}, {ref.max().item():.4f}]') - print(f' out has NaN: {torch.isnan(out).any().item()}') - print(f' out has inf: {torch.isinf(out).any().item()}') - - -if __name__ == '__main__': - test_smem_p_coords() diff --git a/tests/unit/test_d1_3_smem_direct.py b/tests/unit/test_d1_3_smem_direct.py deleted file mode 100644 index 8ee03541..00000000 --- a/tests/unit/test_d1_3_smem_direct.py +++ /dev/null @@ -1,186 +0,0 @@ -""" -D1.3 SMEM-P: Direct SMEM write test. -Write known values to sP via coordinate indexing, -then copy sP to GMEM and verify. -""" -import torch, math -import cutlass, cutlass.cute as cute, cutlass.utils as utils -from cutlass.cute.nvgpu import tcgen05, cpasync -from cutlass import Float32, BFloat16, Int32, const_expr -from cutlass.utils import LayoutEnum -import cutlass.torch as ct -import cuda.bindings.driver as cuda - - -@cute.jit -def smem_write_test(q, k, v, c, stream): - """Write known values to sP, copy to GMEM, verify.""" - a_major = LayoutEnum.from_tensor(q).mma_major_mode() - b_major = LayoutEnum.from_tensor(k).mma_major_mode() - v_fmha = cute.make_tensor( - v.iterator, - cute.make_layout((64, 128, 1), stride=(1, 64, 64 * 128)), - ) - v_major = LayoutEnum.from_tensor(v_fmha).mma_major_mode() - - qk_mma = utils.sm100.make_trivial_tiled_mma( - BFloat16, BFloat16, a_major, b_major, Float32, - tcgen05.CtaGroup.ONE, (128, 128), tcgen05.OperandSource.SMEM - ) - pv_mma = utils.sm100.make_trivial_tiled_mma( - BFloat16, BFloat16, cute.nvgpu.OperandMajorMode.K, v_major, Float32, - tcgen05.CtaGroup.ONE, (128, 64), tcgen05.OperandSource.SMEM - ) - - pv_mma_tiler = (128, 64, 128) - qk_mma_tiler = (128, 128, 128 * 4) - - p_smem_s = utils.sm100.make_smem_layout_a(pv_mma, pv_mma_tiler, BFloat16, 1) - q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, qk_mma_tiler, BFloat16, 1) - k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, qk_mma_tiler, BFloat16, 2) - v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, pv_mma_tiler, BFloat16, 2) - - # TMA for reading back sP - epi_s = cute.select(p_smem_s, mode=[0, 1]) # 2D view of sP for TMA - # Actually, let's just use a simple output tensor - # We'll write known values to sP, then copy to output tensor c - - 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)) - cta = cute.size(qk_mma.thr_id.shape) - - tma_q, mQ = cute.nvgpu.make_tiled_tma_atom_A( - utils.sm100.cluster_shape_to_tma_atom_A((1, 1), qk_mma.thr_id), - q, q_s, qk_mma_tiler, qk_mma, (1, 1, 1, 1) - ) - tma_k, mK = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B((1, 1), qk_mma.thr_id), - k, k_s, qk_mma_tiler, qk_mma, (1, 1, 1, 1) - ) - tma_v, mV = cute.nvgpu.make_tiled_tma_atom_B( - utils.sm100.cluster_shape_to_tma_atom_B((1, 1), pv_mma.thr_id), - v_fmha, v_s, pv_mma_tiler, pv_mma, (1, 1, 1, 1) - ) - - # Output: use c as a flat buffer to read back sP values - # c has shape (128, 64, 1) — same as sP's logical size - # We'll TMA-store sP to c - c_smem_s = utils.sm100.make_smem_layout_a(pv_mma, pv_mma_tiler, BFloat16, 1) - epi_tile = (128, 64) - epi_s2 = cute.select(c_smem_s, mode=[0, 1]) - tma_c, mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_s2, epi_tile) - - # Just use 128 threads (4 warps) for simplicity - _kernel(qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, - p_smem_s, q_smem_s, k_smem_s, v_smem_s, c_smem_s, epi_tile).launch( - grid=(1, 1, 1), block=[128, 1, 1], stream=stream - ) - - -@cute.kernel -def _kernel(qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, - p_smem_s, q_smem_s, k_smem_s, v_smem_s, c_smem_s, epi_tile): - tidx, _, _ = cute.arch.thread_idx() - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - - @cute.struct - class SS: - q_bar: cute.struct.MemRange[cutlass.Int64, 2] - tmem_dealloc: cutlass.Int64 - holding: cutlass.Int32 - smem = utils.SmemAllocator() - st = smem.allocate(SS) - - sP = smem.allocate_tensor(element_type=BFloat16, layout=p_smem_s.outer, byte_alignment=128, swizzle=p_smem_s.inner) - sC = smem.allocate_tensor(element_type=BFloat16, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) - - # All 128 threads write known values to sP - # Strategy: each thread writes its own portion of sP - # using a simple pattern: value = thread_id (cast to BF16) - # This tests that the SMEM write addressing works correctly - - # For this test, each thread writes to sP using the same coordinate mapping - # as the FMHA kernel. But we don't have tTMEM_LOADcS here. - # Instead, let's use a simpler approach: directly write sequential values. - - # Actually, let's just write sP using the MMA fragment A partition - # to verify that write-then-read works. - pv_thr = pv_mma.get_slice(0) - sP_stage = sP[(None, None, None, 0)] - - # Write: each thread writes its portion using MMA's A-operand partition - tCrP = pv_mma.make_fragment_A(sP) - # tCrP is the MMA warp's register fragment for reading sP. - # For writing, we need the "store" side. - # Actually, make_fragment_A creates a load fragment, not a store fragment. - - # Simpler test: just have each thread write a known value to sP directly - # using coordinate indexing with a simple loop. - # Each thread writes 128 values (one row) to sP. - # Thread t writes to row t (for t in 0..127). - if tidx < 128: - m = tidx - for k in range(128): - k0 = k % 16 - k1 = (k // 16) % 4 - k2 = k // 64 - # Write the linear index as BF16: value = (m * 128 + k) % 256 - val = BFloat16(float((m * 128 + k) % 256)) - sP_stage[(m, k0), 0, (k1, k2)] = val - - cute.arch.fence_proxy("async.shared", space="cta") - # Barrier to ensure all writes are visible - bar = pipeline.NamedBarrier(barrier_id=5, num_threads=128) - bar.arrive_and_wait() - - # Now copy sP to sC (same layout), then TMA store to GMEM - # sP and sC have the same layout, so we can copy directly - # Use the TMA store path - gC = cute.local_tile(mC, cute.slice_((128, 64), (None, 0)), (None, None)) - tCgC = pv_thr.partition_C(gC) - - # Copy sP to sC - sC_stage = sC[(None, None, None, 0)] - for m in range(128): - for k in range(128): - k0 = k % 16 - k1 = (k // 16) % 4 - k2 = k // 64 - # Only thread 0 does the copy (simple but slow) - if tidx == 0: - sC_stage[(m, k0), 0, (k1, k2)] = sP_stage[(m, k0), 0, (k1, k2)] - - cute.arch.fence_proxy("async.shared", space="cta") - bar.arrive_and_wait() - - # TMA store sC to GMEM - if tidx == 0: - cpasync.copy_tma_g2s(tma_c, sC, gC) # Wrong direction, need s2g - # Actually, for TMA store (SMEM→GMEM), we need cpasync.copy - # Let me just use a direct store instead - - # Actually this is getting too complicated. Let me use a simpler approach. - # Write the sP values to GMEM directly using a simple loop from thread 0. - - -def test_smem_write(): - print("=== SMEM-P Direct Write Test ===\n") - hd = 64; s_k = 128 - q = torch.randn(128, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - c = torch.zeros(128, hd, 1, dtype=torch.bfloat16, device='cuda') - - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - v_tile = 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_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_tile)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - - print('This test is too complex. Let me take a different approach.', flush=True) - - -if __name__ == '__main__': - test_smem_write() diff --git a/tests/unit/test_d1_3_smem_vs_tmem.py b/tests/unit/test_d1_3_smem_vs_tmem.py deleted file mode 100644 index ff956e09..00000000 --- a/tests/unit/test_d1_3_smem_vs_tmem.py +++ /dev/null @@ -1,88 +0,0 @@ -""" -D1.3 SMEM-P: Focused test of coordinate-indexed SMEM-P write. -Tests at hd=64 with SMEM-P, compares to TMEM-P result. -Also tests hd=128 and hd=256. -""" -import torch, math -import 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 -import cutlass.torch as ct -import cuda.bindings.driver as cuda - -from dsv4.kernels.attention.fmha import FmhaKernel - - -def test_hd(hd, use_smem_p, s_k=128): - pv_n = min(hd, 256) - q = torch.randn(128, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, pv_n, dtype=torch.bfloat16, device='cuda') - c = torch.zeros(128, pv_n, 1, dtype=torch.bfloat16, device='cuda') - - qf = q[:, :, 0].float() - kf = k[:, :, 0].float() - vf = v.float() - scale = 1.0 / math.sqrt(hd) - attn = qf @ kf.T * scale - ref = torch.softmax(attn, dim=-1) @ vf - - kern = FmhaKernel(head_dim=hd, s_k=s_k, use_smem_p=use_smem_p, normalize=True) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - v_tile = 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_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_tile)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - - mode = "SMEM-P" if use_smem_p else "TMEM-P" - print(f'Compiling hd={hd} {mode}...', flush=True) - compiled = cute.compile(kern, 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() - has_nan = torch.isnan(out).any().item() - has_inf = torch.isinf(out).any().item() - - print(f' hd={hd} {mode}: cos={cos:.6f} max_abs={max_abs:.6f} NaN={has_nan} Inf={has_inf}') - - # Print first few values for comparison - print(f' out[0,:4]={out[0,:4].tolist()}') - print(f' ref[0,:4]={ref[0,:4].tolist()}') - print() - - return cos - - -if __name__ == '__main__': - print("=== SMEM-P vs TMEM-P Comparison ===\n") - - # Baseline: TMEM-P at hd=64 (proven path) - cos_tmem = test_hd(64, use_smem_p=False) - - # Test: SMEM-P at hd=64 (should match TMEM-P if correct) - cos_smem = test_hd(64, use_smem_p=True) - - # Test: SMEM-P at hd=128 - cos_128 = test_hd(128, use_smem_p=True) - - # Test: SMEM-P at hd=256 - cos_256 = test_hd(256, use_smem_p=True) - - print("=== Summary ===") - print(f"hd=64 TMEM-P: cos={cos_tmem:.6f} (baseline)") - print(f"hd=64 SMEM-P: cos={cos_smem:.6f} (should match TMEM-P)") - print(f"hd=128 SMEM-P: cos={cos_128:.6f}") - print(f"hd=256 SMEM-P: cos={cos_256:.6f}") - - if abs(cos_tmem - cos_smem) < 0.01: - print("\n✅ SMEM-P matches TMEM-P at hd=64 — coordinate mapping is correct!") - else: - print(f"\n⚠️ SMEM-P differs from TMEM-P by {abs(cos_tmem - cos_smem):.6f} — coordinate mapping has issues") diff --git a/tests/unit/test_d1_3_unnorm_debug.py b/tests/unit/test_d1_3_unnorm_debug.py deleted file mode 100644 index a28a2114..00000000 --- a/tests/unit/test_d1_3_unnorm_debug.py +++ /dev/null @@ -1,96 +0,0 @@ -""" -D1.3 SMEM-P: Debug why hd>64 fails. -Test: compute raw PV (before O normalization) at hd=128 with SMEM-P -and compare against FP32 oracle. -Also test: hd=64 with SMEM-P but skip O normalization to isolate the error. -""" -import torch, math -import cutlass, cutlass.cute as cute -from cutlass import Float32, BFloat16 -import cutlass.torch as ct -import cuda.bindings.driver as cuda -from dsv4.kernels.attention.fmha import FmhaKernel - - -def test_unnormalized(hd, use_smem_p, s_k=128): - """Test with normalize=False to get raw O + LSE, isolate the P write error.""" - pv_n = min(hd, 256) - q = torch.randn(128, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, pv_n, dtype=torch.bfloat16, device='cuda') - c = torch.zeros(128, pv_n, 1, dtype=torch.bfloat16, device='cuda') - lse = torch.zeros(1, dtype=torch.float32, device='cuda') - - qf = q[:, :, 0].float() - kf = k[:, :, 0].float() - vf = v.float() - scale = 1.0 / math.sqrt(hd) - attn = qf @ kf.T * scale - attn_softmax = torch.softmax(attn, dim=-1) - ref = attn_softmax @ vf # normalized reference - ref_unnorm = attn_softmax * attn_softmax.shape[-1] # just for debugging - - kern = FmhaKernel(head_dim=hd, s_k=s_k, use_smem_p=use_smem_p, normalize=False) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - v_tile = 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_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_tile)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - mLSE = ct.from_dlpack(lse).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse)) - - mode = "SMEM-P" if use_smem_p else "TMEM-P" - print(f'Compiling hd={hd} {mode} normalize=False...', flush=True) - compiled = cute.compile(kern, mQ, mK, mV, mC, stream, mLSE) - compiled(mQ, mK, mV, mC, stream, mLSE) - torch.cuda.synchronize() - - out = c[:, :, 0].float() - lse_val = lse.item() - - # The un-normalized output should be: O_unnorm = exp(lse) * O_norm - # So O_norm = O_unnorm / exp(lse) - if lse_val != 0 and not math.isnan(lse_val) and not math.isinf(lse_val): - out_norm = out / math.exp(lse_val) - cos = torch.nn.functional.cosine_similarity( - out_norm.flatten().unsqueeze(0), ref.flatten().unsqueeze(0) - ).item() - max_abs = (out_norm - ref).abs().max().item() - print(f' hd={hd} {mode} unnorm: cos={cos:.6f} max_abs={max_abs:.6f}') - print(f' LSE={lse_val:.6f} exp(lse)={math.exp(lse_val):.6f}') - print(f' out range: [{out.min().item():.4f}, {out.max().item():.4f}]') - print(f' ref range: [{ref.min().item():.4f}, {ref.max().item():.4f}]') - else: - print(f' hd={hd} {mode} unnorm: INVALID LSE={lse_val}') - print(f' out has NaN: {torch.isnan(out).any().item()}') - print(f' out range: [{out.min().item():.4f}, {out.max().item():.4f}]') - - # Also test normalized - kern2 = FmhaKernel(head_dim=hd, s_k=s_k, use_smem_p=use_smem_p, normalize=True) - c2 = torch.zeros(128, pv_n, 1, dtype=torch.bfloat16, device='cuda') - mC2 = ct.from_dlpack(c2).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c2)) - compiled2 = cute.compile(kern2, mQ, mK, mV, mC2, stream) - compiled2(mQ, mK, mV, mC2, stream) - torch.cuda.synchronize() - - out2 = c2[:, :, 0].float() - cos2 = torch.nn.functional.cosine_similarity( - out2.flatten().unsqueeze(0), ref.flatten().unsqueeze(0) - ).item() - print(f' hd={hd} {mode} normalized: cos={cos2:.6f}') - print() - - -if __name__ == '__main__': - print("=== SMEM-P Debug: Unnormalized vs Normalized ===\n") - - # hd=64 baseline - test_unnormalized(64, use_smem_p=False) - test_unnormalized(64, use_smem_p=True) - - # hd=128 - test_unnormalized(128, use_smem_p=True) - - # hd=256 - test_unnormalized(256, use_smem_p=True) diff --git a/tests/unit/test_d1_3_write_read.py b/tests/unit/test_d1_3_write_read.py deleted file mode 100644 index cf79af6d..00000000 --- a/tests/unit/test_d1_3_write_read.py +++ /dev/null @@ -1,232 +0,0 @@ -""" -D1.3 SMEM-P: Debug test that checks if P written to SMEM is read back correctly. -Creates a simple test kernel that: -1. Writes known values to sP via coordinate indexing -2. Reads them back via pv_mma.make_fragment_A(sP) -3. Compares the values -""" -import torch, math -import cutlass, cutlass.cute as cute, cutlass.utils as utils -from cutlass.cute.nvgpu import tcgen05, cpasync -from cutlass import Float32, BFloat16, Int32, const_expr -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 - - -class SmemPWriteReadTest: - """Test kernel: write P to sP, read back, verify.""" - def __init__(self, head_dim=64, s_k=128): - self.head_dim = head_dim - self.s_k = s_k - self.pv_n_tile = min(head_dim, 256) - self.qk_mma_tiler = (128, 128, 128 * 4) - self.pv_mma_tiler = (128, self.pv_n_tile, 128) - self.use_smem_p = True - self.normalize = True - self.acc_dtype = Float32 - self.qk_acc_dtype = Float32 - self.q_dtype = BFloat16 - self.o_dtype = BFloat16 - self.use_2cta_instrs = False - self.cta_group = tcgen05.CtaGroup.ONE - self.cluster_shape_mn = (1, 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.num_acc_stage = 1 - self.scale_softmax_log2 = 1.0 / math.sqrt(self.head_dim) * math.log2(math.e) - self.kv_stage = 2 - self.q_stage = 1 - self.num_c_stage = 2 - self.c_layout = LayoutEnum.ROW_MAJOR - self.tmem_p0_offset = -1 - self.tOrP0_offset = 0 - - @cute.jit - def __call__(self, q, k, v, c, stream): - 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( - (self.pv_n_tile, self.s_k, 1), - stride=(1, self.pv_n_tile, self.pv_n_tile * self.s_k), - ), - ) - self.v_major = LayoutEnum.from_tensor(v_fmha).mma_major_mode() - - 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_a_major = cute.nvgpu.OperandMajorMode.K - pv_mma = utils.sm100.make_trivial_tiled_mma( - self.q_dtype, self.q_dtype, pv_a_major, self.v_major, - self.qk_acc_dtype, self.cta_group, (128, self.pv_n_tile), - tcgen05.OperandSource.SMEM - ) - - # Setup - 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_o0_offset = 0 - s_cols = self.qk_mma_tiler[1] - o_cols = find_tmem_tensor_col_offset(tOtO) - total = max(s_cols, o_cols) - _n = 1 - while _n < total: - _n *= 2 - self.num_tmem_alloc_cols = _n - - p_smem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) - q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, self.q_dtype, self.q_stage) - k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, self.q_dtype, self.kv_stage) - v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, self.kv_stage) - c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, (128, self.pv_n_tile), 2) - - 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)) - cta = cute.size(qk_mma.thr_id.shape) - 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 - - 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, (1, 1, 1, 1) - ) - 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, (1, 1, 1, 1) - ) - 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, (1, 1, 1, 1) - ) - epi_s = cute.select(c_smem_s, mode=[0, 1]) - tma_c, mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_s, (128, self.pv_n_tile)) - - self._kernel(qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, - p_smem_s, q_smem_s, k_smem_s, v_smem_s, c_smem_s, - tStS, tOtO, qk_thr, pv_thr).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, - p_smem_s, q_smem_s, k_smem_s, v_smem_s, c_smem_s, - tStS, tOtO, qk_thr, pv_thr): - tidx, _, _ = cute.arch.thread_idx() - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - - @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) - - 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) - sP = smem.allocate_tensor(element_type=self.q_dtype, layout=p_smem_s.outer, byte_alignment=128, swizzle=p_smem_s.inner) - - # Only use softmax warps for this test - if warp_idx < 4: - # TMEM load partition for getting coordinates - 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, tStS) - - sfw_idx = tidx % (32 * len(self.epilogue_warp_id)) - thr_load = tiled_tmem_load.get_slice(sfw_idx) - tTMEM_LOADtS = thr_load.partition_S(tStS) - - 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) - - # Write known values to sP: value = m * 128 + k (the linear index) - # Each thread writes its portion using coordinate indexing - sP_stage = sP[(None, None, None, 0)] - - # PRINT: coordinates for thread 0 - if sfw_idx == 0: - print(f"=== SMEM-P Write/Read Test ===") - print(f"sP shape: {cute.shape(sP)}") - print(f"sP_stage shape: {cute.shape(sP_stage)}") - print(f"sP_stage layout: {sP_stage.layout}") - print(f"tTMEM_LOADcS shape: {cute.shape(tTMEM_LOADcS)}") - print(f"tTMEM_LOADcS layout: {tTMEM_LOADcS.layout}") - - # Print first few coordinates for thread 0 - for j0 in range(4): - for j1 in range(4): - coord = tTMEM_LOADcS[(j0, 0), j1, 0, 0] - print(f" cS[{j0}, {j1}]: coord={coord}") - - # Write P values to SMEM - for j0 in range(32): - for j1 in range(4): - coord = tTMEM_LOADcS[(j0, 0), j1, 0, 0] - m_c = coord[0] - k_c = coord[1] - k0 = k_c % 16 - k1 = (k_c // 16) % 4 - k2 = k_c // 64 - # Write a known pattern: BF16(m_c * 128 + k_c) - # Use a simple value that's easy to verify - _val = BFloat16(1.0) # Just write 1.0 everywhere - sP_stage[(m_c, k0), 0, (k1, k2)] = _val - - cute.arch.fence_proxy("async.shared", space="cta") - - # Now read back via PV MMA's fragment - tCrP = pv_mma.make_fragment_A(sP) - - if sfw_idx == 0: - print(f"tCrP shape: {cute.shape(tCrP)}") - print(f"tCrP layout: {tCrP.layout}") - # Print first few values read back - for i in range(min(4, cute.size(tCrP, mode=[2]))): - print(f" tCrP[0,0,{i},0] = {tCrP[0, 0, i, 0]}") - - -def test_smem_p_write_read(): - hd = 64 - q = torch.randn(128, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(128, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(128, hd, dtype=torch.bfloat16, device='cuda') - c = torch.zeros(128, hd, 1, dtype=torch.bfloat16, device='cuda') - - kern = SmemPWriteReadTest(head_dim=hd, s_k=128) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - v_tile = 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_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_tile)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - - print('Compiling...', flush=True) - compiled = cute.compile(kern, mQ, mK, mV, mC, stream) - print('Running...', flush=True) - compiled(mQ, mK, mV, mC, stream) - torch.cuda.synchronize() - print('Done.') - - -if __name__ == '__main__': - test_smem_p_write_read() diff --git a/tests/unit/test_d1_debug.py b/tests/unit/test_d1_debug.py deleted file mode 100644 index 4c68a99a..00000000 --- a/tests/unit/test_d1_debug.py +++ /dev/null @@ -1,34 +0,0 @@ -"""D1: Test hd=64 only with CUDA_LAUNCH_BLOCKING for crash debug.""" -import os -os.environ['CUDA_LAUNCH_BLOCKING'] = '1' -import torch, math -import cutlass.cute as cute -import cutlass.torch as ct -import cuda.bindings.driver as cuda -from dsv4.kernels.attention.fmha import FmhaKernel - -hd = 64; n = 128; m = 128 -torch.manual_seed(42) -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') -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() -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 = FmhaKernel(head_dim=hd, s_k=n) -print(f'Compiling hd={hd}...', flush=True) -compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) -print(f'Running hd={hd}...', 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}: cos {cos:.6f}') diff --git a/tests/unit/test_d1_diag.py b/tests/unit/test_d1_diag.py deleted file mode 100644 index e68f4038..00000000 --- a/tests/unit/test_d1_diag.py +++ /dev/null @@ -1,92 +0,0 @@ -"""D1 diagnostic: Print TMEM budget and shapes at hd=256.""" -import torch, math -import cutlass, cutlass.cute as cute -import cutlass.torch as ct -import cutlass.utils as utils -from cutlass.cute.nvgpu import tcgen05 -from cutlass.utils import LayoutEnum -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cuda.bindings.driver as cuda - -hd = 256 -n_kv = 128 -m = 128 - -torch.manual_seed(42) -q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') -k = torch.randn(n_kv, hd, 1, dtype=torch.bfloat16, device='cuda') -v = torch.randn(n_kv, hd, dtype=torch.bfloat16, device='cuda') -v_kernel = v.unsqueeze(-1) -c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - -q_dtype = cutlass.BFloat16; qk_acc_dtype = cutlass.Float32 -cta_group = tcgen05.CtaGroup.ONE - -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)) - -a_major = LayoutEnum.from_tensor(mQ).mma_major_mode() -b_major = LayoutEnum.from_tensor(mK).mma_major_mode() - -pv_n_tile = min(hd, 256) -v_fmha = cute.make_tensor( - mV.iterator, - cute.make_layout((pv_n_tile, n_kv, 1), stride=(1, pv_n_tile, pv_n_tile * n_kv)), -) -v_major = LayoutEnum.from_tensor(v_fmha).mma_major_mode() - -print(f"a_major={a_major}, b_major={b_major}, v_major={v_major}") -print(f"Q shape: {q.shape}, K shape: {k.shape}, V shape: {v.shape}") -print(f"mQ shape: {mQ.shape}, mK shape: {mK.shape}, mV shape: {mV.shape}, mC shape: {mC.shape}") - -qk_mma = utils.sm100.make_trivial_tiled_mma(q_dtype, q_dtype, a_major, b_major, qk_acc_dtype, cta_group, (128, 128), tcgen05.OperandSource.SMEM) -print(f"\nQK MMA shape_mnk: {qk_mma.shape_mnk}") - -pv_mma = utils.sm100.make_trivial_tiled_mma(q_dtype, q_dtype, cutlass.nvgpu.OperandMajorMode.K, v_major, qk_acc_dtype, cta_group, (128, pv_n_tile), tcgen05.OperandSource.TMEM) -print(f"PV MMA shape_mnk: {pv_mma.shape_mnk}") - -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_k_inner = pv_ik * (128 // pv_ik) -pv_mma_tiler = (128, pv_n_tile, pv_k_inner) -print(f"qk_ik={qk_ik}, qk_mma_tiler={qk_mma_tiler}") -print(f"pv_ik={pv_ik}, pv_mma_tiler={pv_mma_tiler}") - -# GMEM tile views -gQ = cute.local_tile(mQ, cute.slice_(qk_mma_tiler,(None,0,None)),(None,None,None)) -gK = cute.local_tile(mK, cute.slice_(qk_mma_tiler,(0,None,None)),(None,None,None)) -gV = cute.local_tile(mV, cute.slice_(pv_mma_tiler,(0,None,None)),(None,None,None)) -gC = cute.local_tile(mC, cute.slice_(pv_mma_tiler,(None,None,0)),(None,None,None)) -print(f"\ngQ shape: {gQ.shape}") -print(f"gK shape: {gK.shape}") -print(f"gV shape: {gV.shape}") -print(f"gC shape: {gC.shape}") -print(f"n_kv_tiles: {cute.size(gK, mode=[3])}") - -# TMEM budget -qk_thr = qk_mma.get_slice(0) -qk_as = qk_thr.partition_shape_C(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(pv_mma_tiler[:2]) -tOtO = pv_thr.make_fragment_C(pv_as) - -tmem_p0_offset = 32 -p_cols_fp32 = pv_mma_tiler[2] * q_dtype.width // qk_acc_dtype.width -o_cols = find_tmem_tensor_col_offset(tOtO) -s_cols = qk_mma_tiler[1] -p_end = tmem_p0_offset + p_cols_fp32 -o_after = max(s_cols, p_end) -tmem_o0_offset = ((o_after + 31) // 32) * 32 -total = tmem_o0_offset + o_cols - -print(f"\n=== TMEM Budget (hd={hd}) ===") -print(f"S={s_cols}, P_fp32={p_cols_fp32}, P_end={p_end}, O_offset={tmem_o0_offset}, O_cols={o_cols}") -print(f"Total={total}, Fits={total<=512}") - -# QK C-fragment info -print(f"\ntStS shape: {tStS.shape}, cosize: {cute.cosize(tStS.layout)}") -print(f"tOtO shape: {tOtO.shape}, cosize: {cute.cosize(tOtO.layout)}") diff --git a/tests/unit/test_d1_diag2.py b/tests/unit/test_d1_diag2.py deleted file mode 100644 index 1c22a62b..00000000 --- a/tests/unit/test_d1_diag2.py +++ /dev/null @@ -1,146 +0,0 @@ -""" -Quick D1 diagnostic: test TMEM-P path (use_smem_p=False) at various head dims. -The SMEM-P path (use_smem_p=True, hd>64) has coordinate mapping issues. -This test forces TMEM-P to verify the core pipeline works. -""" -import torch, math -import cutlass.cute as cute -import cutlass.torch as ct -import cuda.bindings.driver as cuda -from dsv4.kernels.attention.fmha import FmhaKernel - - -def test_tmem_p(hd, n_kv=128): - m = 128 - torch.manual_seed(42) - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(n_kv, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(n_kv, hd, dtype=torch.bfloat16, device='cuda') - 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_max = (qf @ kf.T * scale).max(dim=-1, keepdim=True)[0] - attn_exp = torch.exp(qf @ kf.T * scale - attn_max) - attn_sum = attn_exp.sum(dim=-1, keepdim=True) - ref_norm = (attn_exp / attn_sum) @ v.float() - ref_unnorm = attn_exp @ v.float() - - lse_tensor = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - # Force TMEM-P - kernel = FmhaKernel(head_dim=hd, s_k=n_kv, use_smem_p=False) - pv_n_tile = kernel.pv_n_tile - n_pv_tiles = kernel.n_pv_tiles - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - v_tile = v[:, 0:pv_n_tile].contiguous().unsqueeze(-1) - c_tile = torch.zeros(m, pv_n_tile, 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_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_tile)) - mC = ct.from_dlpack(c_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) - - print(f'hd={hd} TMEM-P: Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE) - - lse_val = None - for nt in range(n_pv_tiles): - vs, ve = nt * pv_n_tile, (nt + 1) * pv_n_tile - v_t = v[:, vs:ve].contiguous().unsqueeze(-1) - c_tile = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - lse_tensor.zero_() - 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_t).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_t)) - mC = ct.from_dlpack(c_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) - compiled(mQ, mK, mV, mC, stream, mLSE) - torch.cuda.synchronize() - c[:, vs:ve, :] = c_tile - if nt == 0: - lse_val = lse_tensor[0, 0, 0].item() - - out = c[:, :, 0].float() - out_norm = out / attn_sum - cos_unnorm = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref_unnorm.flatten().unsqueeze(0)).item() - cos_norm = torch.nn.functional.cosine_similarity(out_norm.flatten().unsqueeze(0), ref_norm.flatten().unsqueeze(0)).item() - status = "PASS" if cos_unnorm >= 0.99 else "FAIL" - print(f'hd={hd} TMEM-P: cos_unnorm {cos_unnorm:.6f} cos_norm {cos_norm:.6f} lse {lse_val:.6f} {status}') - return cos_unnorm - - -def test_smem_p(hd, n_kv=128): - m = 128 - torch.manual_seed(42) - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(n_kv, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(n_kv, hd, dtype=torch.bfloat16, device='cuda') - 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_max = (qf @ kf.T * scale).max(dim=-1, keepdim=True)[0] - attn_exp = torch.exp(qf @ kf.T * scale - attn_max) - attn_sum = attn_exp.sum(dim=-1, keepdim=True) - ref_unnorm = attn_exp @ v.float() - - lse_tensor = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - kernel = FmhaKernel(head_dim=hd, s_k=n_kv, use_smem_p=True) - pv_n_tile = kernel.pv_n_tile - n_pv_tiles = kernel.n_pv_tiles - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - v_tile = v[:, 0:pv_n_tile].contiguous().unsqueeze(-1) - c_tile = torch.zeros(m, pv_n_tile, 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_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_tile)) - mC = ct.from_dlpack(c_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) - - print(f'hd={hd} SMEM-P: Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE) - - lse_val = None - for nt in range(n_pv_tiles): - vs, ve = nt * pv_n_tile, (nt + 1) * pv_n_tile - v_t = v[:, vs:ve].contiguous().unsqueeze(-1) - c_tile = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - lse_tensor.zero_() - 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_t).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_t)) - mC = ct.from_dlpack(c_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) - compiled(mQ, mK, mV, mC, stream, mLSE) - torch.cuda.synchronize() - c[:, vs:ve, :] = c_tile - if nt == 0: - lse_val = lse_tensor[0, 0, 0].item() - - out = c[:, :, 0].float() - cos_unnorm = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref_unnorm.flatten().unsqueeze(0)).item() - status = "PASS" if cos_unnorm >= 0.99 else "FAIL" - print(f'hd={hd} SMEM-P: cos_unnorm {cos_unnorm:.6f} lse {lse_val:.6f} {status}') - if cos_unnorm < 0.97: - print(f' out[0,:4]={out[0,:4].tolist()}') - print(f' ref[0,:4]={ref_unnorm[0,:4].tolist()}') - return cos_unnorm - - -if __name__ == '__main__': - print("=== D1 Diagnostic ===\n") - - # TMEM-P path (proven at hd=64) - print("--- TMEM-P (force use_smem_p=False) ---") - test_tmem_p(64) - test_tmem_p(128) - test_tmem_p(256) - - # SMEM-P path (for hd>64) - print("\n--- SMEM-P (use_smem_p=True) ---") - test_smem_p(128) - test_smem_p(256) diff --git a/tests/unit/test_d1_hd512.py b/tests/unit/test_d1_hd512.py deleted file mode 100644 index afb5e83a..00000000 --- a/tests/unit/test_d1_hd512.py +++ /dev/null @@ -1,48 +0,0 @@ -"""D1 test: HEAD_DIM=512 (DSV4 production config).""" -import torch, math -import cutlass.cute as cute -import cutlass.torch as ct -import cuda.bindings.driver as cuda -from dsv4.kernels.attention.fmha import FmhaKernel - -def test(): - torch.manual_seed(42) - hd, n = 512, 128 - m = 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.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 = FmhaKernel(head_dim=hd, s_k=n) - print(f'hd={hd}, 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() - max_abs = (out - ref).abs().max().item() - print(f'hd={hd}, n={n}: cos {cos:.6f} max_abs {max_abs:.4f} {"PASS" if cos >= 0.97 else "FAIL"}') - if cos < 0.97: - 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/unit/test_d1_hd512_merge.py b/tests/unit/test_d1_hd512_merge.py deleted file mode 100644 index 4e200cc5..00000000 --- a/tests/unit/test_d1_hd512_merge.py +++ /dev/null @@ -1,137 +0,0 @@ -"""D1.4 hd=512 test using external k_sub merge. - -Instead of the k_sub path in the kernel (which causes 45+ min compilation), -we call the kernel once per k_sub tile with Q and K pre-sliced. -The online softmax merge (same as D5) combines the partial results. - -The kernel always runs at k_tile=256 (same as hd=256, proven to compile fast). -""" -import torch, math, time -import cutlass.cute as cute -import cutlass.torch as ct -import cuda.bindings.driver as cuda -from dsv4.kernels.attention.fmha import FmhaKernel - - -def test(): - torch.manual_seed(42) - hd, n = 512, 128 - m = 128 - k_tile = 256 - n_k_sub = hd // k_tile # 2 - - 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') - - # FP32 reference (full attention) - qf = q[:, :, 0].float() - kf = k[:, :, 0].float() - scale = 1.0 / math.sqrt(hd) - attn_max = (qf @ kf.T * scale).max(dim=-1, keepdim=True)[0] - attn_exp = torch.exp(qf @ kf.T * scale - attn_max) - attn_sum = attn_exp.sum(dim=-1, keepdim=True) - ref_unnorm = attn_exp @ v.float() - ref_norm = (attn_exp / attn_sum) @ v.float() - ref_lse = (torch.log(attn_sum.squeeze(-1)) + attn_max.squeeze(-1))[0].item() - - # Use the hd=256 kernel (no k_sub path) with k_tile=256 - # Call once per k_sub tile, merge results via online softmax - kernel = FmhaKernel(head_dim=k_tile, s_k=n, use_smem_p=False, normalize=False) - pv_n_tile = kernel.pv_n_tile - n_pv_tiles = kernel.n_pv_tiles - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - print(f'hd={hd}, k_tile={k_tile}, n_k_sub={n_k_sub}, pv_n_tile={pv_n_tile}', flush=True) - print(f'Compiling k_tile={k_tile} kernel...', flush=True) - - # Compile once with the first k_sub tile - q0 = q[:, 0:k_tile, :].contiguous() - k0 = k[:, 0:k_tile, :].contiguous() - v_tile = v[:, 0:pv_n_tile].contiguous().unsqueeze(-1) - c_tile = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - lse_tensor = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - - mQ0 = ct.from_dlpack(q0).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q0)) - mK0 = ct.from_dlpack(k0).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k0)) - mV = ct.from_dlpack(v_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_tile)) - mC = ct.from_dlpack(c_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) - - t0 = time.time() - compiled = cute.compile(kernel, mQ0, mK0, mV, mC, stream, mLSE) - t1 = time.time() - print(f'Compilation took {t1-t0:.1f}s', flush=True) - - # Run each k_sub tile and accumulate via online softmax merge - # LSE_i = ln(sum(exp(S_i - m_i))) + m_i (in natural log domain) - # Merge: O = (exp(LSE_0 - LSE_max) * O_0 + exp(LSE_1 - LSE_max) * O_1) / - # (exp(LSE_0 - LSE_max) + exp(LSE_1 - LSE_max)) - # where LSE_max = max(LSE_0, LSE_1) - - # Collect (un-norm O, LSE) for each k_sub and each pv_tile - all_o_unnorm = [] # list of (n_k_sub, hd) tensors - all_lse = [] # list of (n_k_sub,) LSE values - - for ks in range(n_k_sub): - ks_start = ks * k_tile - ks_end = ks_start + k_tile - q_ks = q[:, ks_start:ks_end, :].contiguous() - k_ks = k[:, ks_start:ks_end, :].contiguous() - - o_ks = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - lse_ks = None - - for nt in range(n_pv_tiles): - v_start = nt * pv_n_tile - v_end = v_start + pv_n_tile - v_tile = v[:, v_start:v_end].contiguous().unsqueeze(-1) - c_tile = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - lse_tensor.zero_() - - mQ = ct.from_dlpack(q_ks).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q_ks)) - mK = ct.from_dlpack(k_ks).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k_ks)) - mV = ct.from_dlpack(v_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_tile)) - mC = ct.from_dlpack(c_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) - - compiled(mQ, mK, mV, mC, stream, mLSE) - torch.cuda.synchronize() - - o_ks[:, v_start:v_end, :] = c_tile - if nt == 0: - lse_ks = lse_tensor[0, 0, 0].item() - - all_o_unnorm.append(o_ks[:, :, 0].float()) - all_lse.append(lse_ks) - print(f' k_sub={ks}: lse={lse_ks:.4f}', flush=True) - - # Online softmax merge - # O_unnorm_full = sum_ks exp(lse_ks - lse_max) * O_ks - # Normalization: O_norm = O_unnorm_full / sum_ks exp(lse_ks - lse_max) - lse_max = max(all_lse) - o_merged_unnorm = torch.zeros(m, hd, dtype=torch.float32, device='cuda') - denom = 0.0 - for ks in range(n_k_sub): - w = math.exp(all_lse[ks] - lse_max) - o_merged_unnorm += w * all_o_unnorm[ks] - denom += w - - o_merged_norm = o_merged_unnorm / denom - - cos_unnorm = torch.nn.functional.cosine_similarity( - o_merged_unnorm.flatten().unsqueeze(0), ref_unnorm.flatten().unsqueeze(0) - ).item() - cos_norm = torch.nn.functional.cosine_similarity( - o_merged_norm.flatten().unsqueeze(0), ref_norm.flatten().unsqueeze(0) - ).item() - - status = "PASS" if cos_norm >= 0.99 else "FAIL" - print(f'\nhd=512 (external k_sub merge): cos_unnorm {cos_unnorm:.6f} cos_norm {cos_norm:.6f} {status}') - if cos_norm < 0.99: - print(f' o_merged[0,:4]={o_merged_norm[0,:4].tolist()}') - print(f' ref[0,:4]={ref_norm[0,:4].tolist()}') - - -if __name__ == '__main__': - test() diff --git a/tests/unit/test_d1_hd512_only.py b/tests/unit/test_d1_hd512_only.py deleted file mode 100644 index 7e3c8cb6..00000000 --- a/tests/unit/test_d1_hd512_only.py +++ /dev/null @@ -1,88 +0,0 @@ -"""D1 test: HEAD_DIM=512 only (faster iteration on compilation issues).""" -import torch, math, sys -import cutlass.cute as cute -import cutlass.torch as ct -import cuda.bindings.driver as cuda -from dsv4.kernels.attention.fmha import FmhaKernel - - -def test(): - torch.manual_seed(42) - hd, n = 512, 128 - m = 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.randn(n, hd, dtype=torch.bfloat16, device='cuda') - - qf = q[:, :, 0].float() - kf = k[:, :, 0].float() - scale = 1.0 / math.sqrt(hd) - attn_max = (qf @ kf.T * scale).max(dim=-1, keepdim=True)[0] - attn_exp = torch.exp(qf @ kf.T * scale - attn_max) - attn_sum = attn_exp.sum(dim=-1, keepdim=True) - ref_unnorm = attn_exp @ v.float() - ref_lse = (torch.log(attn_sum.squeeze(-1)) + attn_max.squeeze(-1))[0].item() - - lse_tensor = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - kernel = FmhaKernel(head_dim=hd, s_k=n, use_smem_p=False) - pv_n_tile = kernel.pv_n_tile - n_pv_tiles = kernel.n_pv_tiles - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - print(f'hd={hd}, pv_n_tile={pv_n_tile}, n_pv_tiles={n_pv_tiles}, n_k_sub_tiles={kernel.n_k_sub_tiles}, k_tile={kernel.k_tile}', flush=True) - print(f'Compiling first PV tile...', flush=True) - - # Only compile the first PV tile to isolate compilation issues - v_tile = v[:, 0:pv_n_tile].contiguous() - v_kernel = v_tile.unsqueeze(-1) - c_tile = torch.zeros(m, pv_n_tile, 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_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) - - import time - t0 = time.time() - from cutlass.base_dsl.compiler import PtxasOptions, OptLevel - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE) - t1 = time.time() - print(f'Compilation took {t1-t0:.1f}s', flush=True) - - # Run all PV tiles - lse_val = None - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - for nt in range(n_pv_tiles): - v_start = nt * pv_n_tile - v_end = v_start + pv_n_tile - v_tile = v[:, v_start:v_end].contiguous() - v_kernel = v_tile.unsqueeze(-1) - c_tile = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - lse_tensor.zero_() - - 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_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) - - compiled(mQ, mK, mV, mC, stream, mLSE) - torch.cuda.synchronize() - print(f' PV tile {nt}: done', flush=True) - - c[:, v_start:v_end, :] = c_tile - if nt == 0: - lse_val = lse_tensor[0, 0, 0].item() - - out_unnorm = c[:, :, 0].float() - cos_unnorm = torch.nn.functional.cosine_similarity( - out_unnorm.flatten().unsqueeze(0), ref_unnorm.flatten().unsqueeze(0) - ).item() - lse_err = abs(lse_val - ref_lse) if lse_val is not None else float('inf') - - status = "PASS" if cos_unnorm >= 0.99 else "FAIL" - print(f'hd={hd}: cos_unnorm {cos_unnorm:.6f} lse_err {lse_err:.6f} {status}') - -if __name__ == '__main__': - test() diff --git a/tests/unit/test_d1_kv_merge.py b/tests/unit/test_d1_kv_merge.py deleted file mode 100644 index 24137346..00000000 --- a/tests/unit/test_d1_kv_merge.py +++ /dev/null @@ -1,119 +0,0 @@ -""" -D1: Test multi-KV-tile by running s_k=128 kernel per KV segment and -merging in Python using log-sum-exp (D5 merge formula). - -This avoids the broken TMEM round-trip O rescale entirely. -""" -import torch, math -import cutlass.cute as cute -import cutlass.torch as ct -import cuda.bindings.driver as cuda -from dsv4.kernels.attention.fmha import FmhaKernel - - -def test_multi_kv_merge(hd=64, s_k=256): - m = 128 - n_kv_segments = s_k // 128 - torch.manual_seed(42) - - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - - # FP32 reference (full attention) - qf = q[:, :, 0].float() - kf = k[:, :, 0].float() - scale = 1.0 / math.sqrt(hd) - attn_max = (qf @ kf.T * scale).max(dim=-1, keepdim=True)[0] - attn_exp = torch.exp(qf @ kf.T * scale - attn_max) - attn_sum = attn_exp.sum(dim=-1, keepdim=True) - ref_norm = (attn_exp / attn_sum) @ v.float() - - # Run s_k=128 kernel per KV segment and merge using log-sum-exp - kernel = FmhaKernel(head_dim=hd, s_k=128, use_smem_p=False, normalize=False) - pv_n_tile = kernel.pv_n_tile - n_pv_tiles = kernel.n_pv_tiles - - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - # Compile once with segment 0's K - k_seg = k[:128] - v_tile = v[:, 0:pv_n_tile].contiguous() - v_kernel = v_tile.unsqueeze(-1) - c_tile = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - lse_tensor = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k_seg).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k_seg)) - mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel)) - mC = ct.from_dlpack(c_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) - - print(f' Compiling (hd={hd}, s_k=128 per segment, {n_kv_segments} segments)...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE) - - # Accumulate across KV segments using log-sum-exp merge - # O_merged = sum_i(exp(lse_i) * O_i) / sum_i(exp(lse_i)) - o_accum = torch.zeros(m, hd, dtype=torch.float32, device='cuda') - lse_accum = torch.full((m, 1), float('-inf'), dtype=torch.float32, device='cuda') - - for seg in range(n_kv_segments): - k_start = seg * 128 - k_end = k_start + 128 - k_seg = k[k_start:k_end] - v_seg = v[k_start:k_end] - - # Per-segment O and LSE - seg_o = torch.zeros(m, hd, dtype=torch.float32, device='cuda') - seg_lse = torch.zeros(m, 1, dtype=torch.float32, device='cuda') - - for nt in range(n_pv_tiles): - v_start = nt * pv_n_tile - v_end = v_start + pv_n_tile - v_tile = v_seg[:, v_start:v_end].contiguous() - v_kernel = v_tile.unsqueeze(-1) - c_tile = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - lse_tensor.zero_() - - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k_seg).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k_seg)) - mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel)) - mC = ct.from_dlpack(c_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) - - compiled(mQ, mK, mV, mC, stream, mLSE) - torch.cuda.synchronize() - - seg_o[:, v_start:v_end] = c_tile[:, :, 0].float() - if nt == 0: - seg_lse[:, 0] = lse_tensor[:, 0, 0].float() - - # Merge with accumulator using log-sum-exp - # O_new = (exp(lse_old) * O_old + exp(lse_new) * O_new) / (exp(lse_old) + exp(lse_new)) - # lse_new = ln(exp(lse_old) + exp(lse_new)) - e_old = torch.exp(lse_accum) # (m, 1) - e_new = torch.exp(seg_lse) # (m, 1) - e_sum = e_old + e_new - - o_accum = (e_old * o_accum + e_new * seg_o) / e_sum - lse_accum = torch.log(e_sum) - - cos = torch.nn.functional.cosine_similarity( - o_accum.flatten().unsqueeze(0), ref_norm.flatten().unsqueeze(0) - ).item() - print(f' hd={hd}, s_k={s_k} ({n_kv_segments} segments): cos_norm {cos:.6f} {"PASS" if cos >= 0.99 else "FAIL"}') - return cos - - -def test(): - print("=== D1: Multi-KV Merge via Log-Sum-Exp (no TMEM round-trip) ===\n") - - test_multi_kv_merge(64, 256) - test_multi_kv_merge(64, 384) - test_multi_kv_merge(64, 512) - test_multi_kv_merge(64, 1024) - test_multi_kv_merge(128, 256) - - -if __name__ == '__main__': - test() diff --git a/tests/unit/test_d1_kv_merge_v2.py b/tests/unit/test_d1_kv_merge_v2.py deleted file mode 100644 index 49ed50dc..00000000 --- a/tests/unit/test_d1_kv_merge_v2.py +++ /dev/null @@ -1,122 +0,0 @@ -""" -D1: Multi-KV-tile merge using per-row LSE and log-sum-exp. - -Strategy: Run s_k=128 kernel per KV segment, get per-row O and LSE. -Merge using the D5 formula: - O = (exp(lse_0) * O_0 + exp(lse_1) * O_1) / (exp(lse_0) + exp(lse_1)) - -This avoids the broken TMEM round-trip O rescale. -""" -import torch, math -import cutlass.cute as cute -import cutlass.torch as ct -import cuda.bindings.driver as cuda -from dsv4.kernels.attention.fmha import FmhaKernel - - -def test_multi_kv_merge(hd=64, s_k=256): - m = 128 - n_kv_segments = s_k // 128 - torch.manual_seed(42) - - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - - # FP32 reference (full attention) - qf = q[:, :, 0].float() - kf = k[:, :, 0].float() - scale = 1.0 / math.sqrt(hd) - attn_scores = qf @ kf.T * scale - attn_max = attn_scores.max(dim=-1, keepdim=True)[0] - attn_exp = torch.exp(attn_scores - attn_max) - attn_sum = attn_exp.sum(dim=-1, keepdim=True) - ref_norm = (attn_exp / attn_sum) @ v.float() - - # Run s_k=128 kernel per KV segment - kernel = FmhaKernel(head_dim=hd, s_k=128, use_smem_p=False, normalize=False) - pv_n_tile = kernel.pv_n_tile - n_pv_tiles = kernel.n_pv_tiles - - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - # Compile once - k_seg0 = k[:128] - v_tile0 = v[:128, 0:pv_n_tile].contiguous() - v_kernel0 = v_tile0.unsqueeze(-1) - c_tile0 = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - # Per-row LSE: shape (m,) - lse_tensor = torch.zeros(m, dtype=torch.float32, device='cuda') - - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k_seg0).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k_seg0)) - mV = ct.from_dlpack(v_kernel0).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel0)) - mC = ct.from_dlpack(c_tile0).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile0)) - mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) - - print(f' Compiling (hd={hd}, s_k=128, {n_kv_segments} segments)...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE) - - # Accumulate across KV segments using log-sum-exp merge - o_accum = None # Will be (m, hd) FP32 - lse_accum = None # Will be (m,) FP32 - - for seg in range(n_kv_segments): - k_start = seg * 128 - k_end = k_start + 128 - k_seg = k[k_start:k_end] - v_seg = v[k_start:k_end] - - seg_o = torch.zeros(m, hd, dtype=torch.float32, device='cuda') - - for nt in range(n_pv_tiles): - v_start = nt * pv_n_tile - v_end = v_start + pv_n_tile - v_tile = v_seg[:, v_start:v_end].contiguous() - v_kernel = v_tile.unsqueeze(-1) - c_tile = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - lse_tensor.zero_() - - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k_seg).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k_seg)) - mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel)) - mC = ct.from_dlpack(c_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) - - compiled(mQ, mK, mV, mC, stream, mLSE) - torch.cuda.synchronize() - - seg_o[:, v_start:v_end] = c_tile[:, :, 0].float() - - seg_lse = lse_tensor.clone() # (m,) per-row LSE - - # Log-sum-exp merge with accumulator - if o_accum is None: - o_accum = seg_o - lse_accum = seg_lse - else: - e_old = torch.exp(lse_accum) # (m,) - e_new = torch.exp(seg_lse) # (m,) - e_sum = e_old + e_new # (m,) - - o_accum = (e_old.unsqueeze(-1) * o_accum + e_new.unsqueeze(-1) * seg_o) / e_sum.unsqueeze(-1) - lse_accum = torch.log(e_sum) - - cos = torch.nn.functional.cosine_similarity( - o_accum.flatten().unsqueeze(0), ref_norm.flatten().unsqueeze(0) - ).item() - print(f' hd={hd}, s_k={s_k} ({n_kv_segments} segments): cos_norm {cos:.6f} {"PASS" if cos >= 0.99 else "FAIL"}') - return cos - - -def test(): - print("=== D1: Multi-KV Merge via Per-Row Log-Sum-Exp ===\n") - - test_multi_kv_merge(64, 256) - test_multi_kv_merge(64, 384) - test_multi_kv_merge(64, 512) - test_multi_kv_merge(64, 1024) - - -if __name__ == '__main__': - test() diff --git a/tests/unit/test_d1_kv_merge_v3.py b/tests/unit/test_d1_kv_merge_v3.py deleted file mode 100644 index 52e1dc0f..00000000 --- a/tests/unit/test_d1_kv_merge_v3.py +++ /dev/null @@ -1,164 +0,0 @@ -""" -D1: Multi-KV-tile merge using per-row LSE and NORMALIZED outputs. - -Correct formula: - O = sum_i [exp(lse_i) * O_i_norm] / sum_i [exp(lse_i)] - -Where O_i_norm = O_i_unnorm / row_sum_i (per-segment normalized output) -And exp(lse_i) = row_sum_i * exp(max(S_i * scale)) -""" -import torch, math -import cutlass.cute as cute -import cutlass.torch as ct -import cuda.bindings.driver as cuda -from dsv4.kernels.attention.fmha import FmhaKernel - - -def test_multi_kv_merge(hd=64, s_k=256): - m = 128 - n_kv_segments = s_k // 128 - torch.manual_seed(42) - - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - - # FP32 reference (full attention) - qf = q[:, :, 0].float() - kf = k[:, :, 0].float() - scale = 1.0 / math.sqrt(hd) - attn = qf @ kf.T * scale - attn_max = attn.max(dim=-1, keepdim=True)[0] - attn_exp = torch.exp(attn - attn_max) - attn_sum = attn_exp.sum(dim=-1, keepdim=True) - ref_norm = (attn_exp / attn_sum) @ v.float() - - # Run s_k=128 kernel per KV segment - kernel = FmhaKernel(head_dim=hd, s_k=128, use_smem_p=False, normalize=False) - pv_n_tile = kernel.pv_n_tile - - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - # Compile once - k_seg0 = k[:128] - v_tile0 = v[:128, 0:pv_n_tile].contiguous() - v_kernel0 = v_tile0.unsqueeze(-1) - c_tile0 = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - lse_tensor = torch.zeros(m, dtype=torch.float32, device='cuda') - - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k_seg0).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k_seg0)) - mV = ct.from_dlpack(v_kernel0).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel0)) - mC = ct.from_dlpack(c_tile0).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile0)) - mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) - - print(f' Compiling (hd={hd}, s_k=128, {n_kv_segments} segments)...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE) - - # Accumulate across KV segments - o_norm_accum = None # (m, hd) normalized output - w_accum = None # (m,) weight = exp(lse) - - for seg in range(n_kv_segments): - k_start = seg * 128 - k_end = k_start + 128 - k_seg = k[k_start:k_end] - v_seg = v[k_start:k_end] - - seg_o_unnorm = torch.zeros(m, hd, dtype=torch.float32, device='cuda') - - for nt in range(1): # hd=64 → n_pv_tiles=1 - v_start = nt * pv_n_tile - v_end = v_start + pv_n_tile - v_tile = v_seg[:, v_start:v_end].contiguous() - v_kernel = v_tile.unsqueeze(-1) - c_tile = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - lse_tensor.zero_() - - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k_seg).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k_seg)) - mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel)) - mC = ct.from_dlpack(c_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) - - compiled(mQ, mK, mV, mC, stream, mLSE) - torch.cuda.synchronize() - - seg_o_unnorm[:, v_start:v_end] = c_tile[:, :, 0].float() - - seg_lse = lse_tensor.clone() # (m,) per-row LSE - seg_w = torch.exp(seg_lse) # (m,) = row_sum * exp(max(S * scale)) - - # Normalize this segment's O - # O_norm = O_unnorm / row_sum - # But we don't have row_sum directly. We have lse = ln(row_sum) + M * ln(2) - # So row_sum = exp(lse - M * ln(2)) = exp(lse) / exp(M * ln(2)) - # But M * ln(2) is in the scale_log2 domain... - # - # Actually, the un-normalized O is O_unnorm = P @ V where P = exp(S*scale - row_max) - # And row_sum = sum(P). - # So O_norm = O_unnorm / row_sum. - # - # But row_sum is not directly available. We have lse = ln(row_sum) + row_max * ln(2). - # So row_sum = exp(lse - row_max * ln(2)). - # - # But row_max is in scale_log2 domain: row_max = max(S * scale * log2(e)) - # So row_max * ln(2) = max(S * scale) - # - # Therefore: row_sum = exp(lse) / exp(max(S * scale)) = exp(lse) / (2^row_max) - # - # Hmm, we don't have max(S * scale) separately. - # But we don't need it! The merge formula is: - # O = sum_i [exp(lse_i) * O_i_norm] / sum_i [exp(lse_i)] - # = sum_i [exp(lse_i) * O_i_unnorm / row_sum_i] / sum_i [exp(lse_i)] - # = sum_i [exp(lse_i) * O_i_unnorm / (exp(lse_i) / exp(M_i))] / sum_i [exp(lse_i)] - # = sum_i [exp(M_i) * O_i_unnorm] / sum_i [exp(lse_i)] - # - # So the numerator uses exp(M_i) * O_i_unnorm, where M_i = max(S_i * scale). - # But M_i = row_max_i * ln(2), and we don't have row_max_i separately. - # - # We can derive row_max_i from lse and row_sum: - # But we don't have row_sum either. - # - # Alternative: compute O_norm from O_unnorm using: - # O_norm_i = O_unnorm_i / row_sum_i - # row_sum_i = sum(P_i) = sum(exp(S_i * scale - M_i)) - # - # In the kernel, row_sum is computed per-thread. We need to output it. - # - # For now, let me compute row_sum from the reference for testing: - seg_kf = k_seg[:, :, 0].float() - seg_attn = qf @ seg_kf.T * scale - seg_attn_max = seg_attn.max(dim=-1)[0] - seg_row_sum = torch.exp(seg_attn - seg_attn_max.unsqueeze(-1)).sum(dim=-1) # (m,) - - seg_o_norm = seg_o_unnorm / seg_row_sum.unsqueeze(-1) - - # Merge: O = sum_i [exp(lse_i) * O_i_norm] / sum_i [exp(lse_i)] - if o_norm_accum is None: - o_norm_accum = seg_w.unsqueeze(-1) * seg_o_norm - w_accum = seg_w - else: - o_norm_accum = o_norm_accum + seg_w.unsqueeze(-1) * seg_o_norm - w_accum = w_accum + seg_w - - o_merged = o_norm_accum / w_accum.unsqueeze(-1) - - cos = torch.nn.functional.cosine_similarity( - o_merged.flatten().unsqueeze(0), ref_norm.flatten().unsqueeze(0) - ).item() - print(f' hd={hd}, s_k={s_k} ({n_kv_segments} segments): cos_norm {cos:.6f} {"PASS" if cos >= 0.99 else "FAIL"}') - return cos - - -def test(): - print("=== D1: Multi-KV Merge (corrected formula) ===\n") - - test_multi_kv_merge(64, 256) - test_multi_kv_merge(64, 384) - test_multi_kv_merge(64, 512) - test_multi_kv_merge(64, 1024) - - -if __name__ == '__main__': - test() diff --git a/tests/unit/test_d1_lse.py b/tests/unit/test_d1_lse.py deleted file mode 100644 index 891597c2..00000000 --- a/tests/unit/test_d1_lse.py +++ /dev/null @@ -1,60 +0,0 @@ -"""Quick LSE diagnostic: is the softmax correct at hd>64?""" -import torch, math -import cutlass.cute as cute -import cutlass.torch as ct -import cuda.bindings.driver as cuda -from dsv4.kernels.attention.fmha import FmhaKernel - - -def test_lse(hd, n_kv=128): - m = 128 - torch.manual_seed(42) - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(n_kv, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(n_kv, hd, dtype=torch.bfloat16, device='cuda') - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - - # Reference LSE - qf = q[:, :, 0].float() - kf = k[:, :, 0].float() - scale = 1.0 / math.sqrt(hd) - attn = qf @ kf.T * scale - attn_max = attn.max(dim=-1, keepdim=True)[0] - attn_exp = torch.exp(attn - attn_max) - attn_sum = attn_exp.sum(dim=-1, keepdim=True) - ref_lse = torch.log(attn_sum.squeeze(-1)) + attn_max.squeeze(-1) - - lse_tensor = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - kernel = FmhaKernel(head_dim=hd, s_k=n_kv) - pv_n_tile = kernel.pv_n_tile - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - v_tile = v[:, 0:pv_n_tile].contiguous().unsqueeze(-1) - c_tile = torch.zeros(m, pv_n_tile, 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_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_tile)) - mC = ct.from_dlpack(c_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) - - print(f'hd={hd}: Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE) - compiled(mQ, mK, mV, mC, stream, mLSE) - torch.cuda.synchronize() - - kernel_lse = lse_tensor[0, 0, 0].item() - ref_lse_val = ref_lse[0].item() - lse_err = abs(kernel_lse - ref_lse_val) - print(f'hd={hd}: kernel_lse={kernel_lse:.6f} ref_lse={ref_lse_val:.6f} err={lse_err:.6f} {"PASS" if lse_err < 0.01 else "FAIL"}') - - # Also check if P store to TMEM is correct by comparing O directly - # Output the raw O (un-normalized) from the kernel - out = c[:, :, 0].float() - ref_unnorm = attn_exp @ v.float() - cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref_unnorm.flatten().unsqueeze(0)).item() - print(f'hd={hd}: cos_unnorm={cos:.6f}') - - -if __name__ == '__main__': - for hd in [64, 128, 256]: - test_lse(hd) diff --git a/tests/unit/test_d1_lse_verify.py b/tests/unit/test_d1_lse_verify.py deleted file mode 100644 index 90200ff1..00000000 --- a/tests/unit/test_d1_lse_verify.py +++ /dev/null @@ -1,81 +0,0 @@ -""" -D1: Verify per-row LSE correctness. -""" -import torch, math -import cutlass.cute as cute -import cutlass.torch as ct -import cuda.bindings.driver as cuda -from dsv4.kernels.attention.fmha import FmhaKernel - - -def test_lse(hd=64): - m = 128 - s_k = 128 - torch.manual_seed(42) - - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - - # FP32 reference LSE - qf = q[:, :, 0].float() - kf = k[:, :, 0].float() - scale = 1.0 / math.sqrt(hd) - attn = qf @ kf.T * scale - attn_max = attn.max(dim=-1)[0] # (m,) - attn_exp = torch.exp(attn - attn_max.unsqueeze(-1)) - attn_sum = attn_exp.sum(dim=-1) # (m,) - ref_lse = torch.log(attn_sum) + attn_max # (m,) natural log domain - - # Our kernel LSE: lse = ln(row_sum) + row_max * ln(2) - # row_max is in scale_log2 domain: max(S * scale * log2(e)) - # So row_max * ln(2) converts back to natural domain. - # Thus lse = ln(row_sum) + row_max * ln(2) - # But row_max = max(S * scale * log2(e)) = max(S * scale) * log2(e) - # So row_max * ln(2) = max(S * scale) * log2(e) * ln(2) = max(S * scale) - # Therefore lse = ln(row_sum) + max(S * scale) - # And our ref: lse = ln(sum(exp(S * scale - max))) + max = ln(sum) + max - # These should be the same. - - kernel = FmhaKernel(head_dim=hd, s_k=s_k, use_smem_p=False, normalize=False) - pv_n_tile = kernel.pv_n_tile - - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - v_tile = v[:, 0:pv_n_tile].contiguous() - v_kernel = v_tile.unsqueeze(-1) - c_tile = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - lse_tensor = torch.zeros(m, dtype=torch.float32, 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_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) - - print(f' Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE) - compiled(mQ, mK, mV, mC, stream, mLSE) - torch.cuda.synchronize() - - kernel_lse = lse_tensor.cpu() - ref_lse_cpu = ref_lse.cpu() - - # Compare - lse_err = (kernel_lse - ref_lse_cpu).abs() - print(f' LSE max err: {lse_err.max().item():.6f}') - print(f' LSE mean err: {lse_err.mean().item():.6f}') - print(f' Kernel LSE[:8]: {kernel_lse[:8].tolist()}') - print(f' Ref LSE[:8]: {ref_lse_cpu[:8].tolist()}') - - # Check O too - ref_unnorm = attn_exp @ v.float() - out = c_tile[:, :, 0].float() - cos = torch.nn.functional.cosine_similarity( - out.flatten().unsqueeze(0), ref_unnorm.flatten().unsqueeze(0) - ).item() - print(f' O cos_unnorm: {cos:.6f}') - - -if __name__ == '__main__': - test_lse() diff --git a/tests/unit/test_d1_multi_kv.py b/tests/unit/test_d1_multi_kv.py deleted file mode 100644 index 2d9e78fb..00000000 --- a/tests/unit/test_d1_multi_kv.py +++ /dev/null @@ -1,122 +0,0 @@ -""" -FMHA D1: Test O rescale with multiple KV tiles (s_k > 128). - -DSV4 Pro uses top_k=1024 → s_k=1024 → n_kv_tiles=8. -The O rescale code (kt>0) is guarded with const_expr(n_kv_tiles > 1) -and uses hand-constructed TMEM atoms. Untested and likely broken. - -This test verifies O rescale correctness at s_k=256 (2 KV tiles). -""" -import torch, math -import cutlass.cute as cute -import cutlass.torch as ct -import cuda.bindings.driver as cuda -from dsv4.kernels.attention.fmha import FmhaKernel - - -def test_multi_kv(hd=64, s_k=256): - m = 128 - n_kv_tiles = s_k // 128 - torch.manual_seed(42) - - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - - # FP32 reference - qf = q[:, :, 0].float() - kf = k[:, :, 0].float() - scale = 1.0 / math.sqrt(hd) - attn_max = (qf @ kf.T * scale).max(dim=-1, keepdim=True)[0] - attn_exp = torch.exp(qf @ kf.T * scale - attn_max) - attn_sum = attn_exp.sum(dim=-1, keepdim=True) - ref_unnorm = attn_exp @ v.float() - ref_norm = (attn_exp / attn_sum) @ v.float() - ref_lse = (torch.log(attn_sum.squeeze(-1)) + attn_max.squeeze(-1))[0].item() - - lse_tensor = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - - kernel = FmhaKernel(head_dim=hd, s_k=s_k, use_smem_p=False, normalize=False) - pv_n_tile = kernel.pv_n_tile - n_pv_tiles = kernel.n_pv_tiles - - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - # Compile with first PV tile - v_tile = v[:, 0:pv_n_tile].contiguous() - v_kernel = v_tile.unsqueeze(-1) - c_tile = torch.zeros(m, pv_n_tile, 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_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) - - print(f'hd={hd}, s_k={s_k} (n_kv_tiles={n_kv_tiles}, pv_n_tile={pv_n_tile}): Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE) - - lse_val = None - for nt in range(n_pv_tiles): - v_start = nt * pv_n_tile - v_end = v_start + pv_n_tile - v_tile = v[:, v_start:v_end].contiguous() - v_kernel = v_tile.unsqueeze(-1) - c_tile = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - lse_tensor.zero_() - - 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_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) - - compiled(mQ, mK, mV, mC, stream, mLSE) - torch.cuda.synchronize() - - c[:, v_start:v_end, :] = c_tile - if nt == 0: - lse_val = lse_tensor[0, 0, 0].item() - - out_unnorm = c[:, :, 0].float() - out_norm = out_unnorm / attn_sum - cos_unnorm = torch.nn.functional.cosine_similarity( - out_unnorm.flatten().unsqueeze(0), ref_unnorm.flatten().unsqueeze(0) - ).item() - cos_norm = torch.nn.functional.cosine_similarity( - out_norm.flatten().unsqueeze(0), ref_norm.flatten().unsqueeze(0) - ).item() - lse_err = abs(lse_val - ref_lse) if lse_val is not None else float('inf') - - status = "PASS" if cos_unnorm >= 0.99 else "FAIL" - print(f'hd={hd}, s_k={s_k}: cos_unnorm {cos_unnorm:.6f} cos_norm {cos_norm:.6f} lse_err {lse_err:.6f} {status}') - return cos_unnorm, cos_norm, lse_err - - -def test(): - print("=== D1: Multi-KV-Tile O Rescale Test ===\n") - - # First: s_k=128 baseline (1 KV tile, no rescale needed) - print("--- Baseline: s_k=128 (1 KV tile) ---") - test_multi_kv(64, 128) - - # Critical test: s_k=256 (2 KV tiles, O rescale exercised) - print("\n--- s_k=256 (2 KV tiles, O rescale needed) ---") - test_multi_kv(64, 256) - - # s_k=384 (3 KV tiles) - print("\n--- s_k=384 (3 KV tiles) ---") - test_multi_kv(64, 384) - - # s_k=512 (4 KV tiles — Flash decode config) - print("\n--- s_k=512 (4 KV tiles, Flash decode) ---") - test_multi_kv(64, 512) - - # hd=128 with multi-KV - print("\n--- hd=128, s_k=256 ---") - test_multi_kv(128, 256) - - -if __name__ == '__main__': - test() diff --git a/tests/unit/test_d1_qk512.py b/tests/unit/test_d1_qk512.py deleted file mode 100644 index 344167df..00000000 --- a/tests/unit/test_d1_qk512.py +++ /dev/null @@ -1,164 +0,0 @@ -"""Minimal hd=512 test: ONLY QK GEMM, no softmax, no PV. -Goal: isolate whether the compilation hang is from QK or softmax/PV.""" -import torch, math, time, cutlass -import cutlass.cute as cute -import cutlass.torch as ct -import cuda.bindings.driver as cuda -from cutlass import BFloat16, Float32 -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass.utils import LayoutEnum -import cutlass.utils as utils -import cutlass.pipeline as pipeline -from cutlass import const_expr - -class QkOnly512: - def __init__(self): - self.head_dim = 512 - self.k_tile = 256 - self.n_k_sub_tiles = 2 - self.kv_stage = 1 - self.q_stage = 1 - self.q_dtype = BFloat16 - self.qk_acc_dtype = Float32 - self.cta_group = tcgen05.CtaGroup.ONE - self.cluster_shape_mn = (1, 1) - self.qk_mma_tiler = (128, 128, self.k_tile) - self.threads_per_cta = 192 - self.mma_warp_id = 4 - self.tma_warp_id = 5 - self.epilogue_warp_id = (0,1,2,3) - - def _setup(self, qk_mma): - 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) - 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, s_out, stream): - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() - 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_shape_mn) - 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_shape_mn) - self._kernel(qk_mma, tma_q, mQ, tma_k, mK, s_out).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, s_out): - 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) - - @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] - 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=cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)),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=cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)),defer_sync=True).make_participants() - - 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) - pipeline.pipeline_init_arrive(cluster_shape_mn=cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)),is_relaxed=True) - - sQ = smem.allocate_tensor(element_type=self.q_dtype,layout=self.q_smem_s.outer,byte_alignment=128,swizzle=self.q_smem_s.inner) - sK = smem.allocate_tensor(element_type=self.q_dtype,layout=self.k_smem_s.outer,byte_alignment=128,swizzle=self.k_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)) - - qk_thr = qk_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK) - a_lay = cute.make_layout((1,)) - b_lay = cute.make_layout((1,)) - 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)) - 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, tStS.layout) - - pipeline.pipeline_init_wait(cluster_shape_mn=cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,))) - - # ===== TMA LOAD warp ===== - if warp_idx == self.tma_warp_id: - qp.reset() - kvp.reset() - # k_sub=0 - qh0 = qp.acquire_and_advance() - cute.copy(tma_q, tAgQ[(None, cutlass.Int32(0))], tAsQ[(None, qh0.index)], tma_bar_ptr=qh0.barrier) - kvh0 = kvp.acquire_and_advance() - cute.copy(tma_k, tBgK[(None, cutlass.Int32(0))], tBsK[(None, kvh0.index)], tma_bar_ptr=kvh0.barrier) - # k_sub=1 - qh1 = qp.acquire_and_advance() - cute.copy(tma_q, tAgQ[(None, cutlass.Int32(1))], tAsQ[(None, qh1.index)], tma_bar_ptr=qh1.barrier) - kvh1 = kvp.acquire_and_advance() - cute.copy(tma_k, tBgK[(None, cutlass.Int32(1))], tBsK[(None, kvh1.index)], tma_bar_ptr=kvh1.barrier) - qp.tail() - kvp.tail() - - # ===== MMA warp ===== - if warp_idx == self.mma_warp_id: - tmem.wait_for_alloc() - # k_sub=0 - qh0 = qc.wait_and_advance(); qh0.release() - kvh0 = kvc.wait_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,kvh0.index)], tStS0) - qk_mma.set(tcgen05.Field.ACCUMULATE, True) - kvh0.release() - # k_sub=1 - qh1 = qc.wait_and_advance(); qh1.release() - kvh1 = kvc.wait_and_advance() - 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,kvh1.index)], tStS0) - qk_mma.set(tcgen05.Field.ACCUMULATE, True) - kvh1.release() - cute.arch.fence_view_async_tmem_store() - - # Epilogue warps just allocate/free TMEM - if warp_idx < self.mma_warp_id: - tmem.allocate(64) - tmem.wait_for_alloc() - tmem.relinquish_alloc_permit() - tmem.free(tmem.retrieve_ptr(self.qk_acc_dtype)) - - -def test(): - torch.manual_seed(42) - hd, n, m = 512, 128, 128 - - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(n, hd, 1, dtype=torch.bfloat16, device='cuda') - s_out = torch.zeros(1, dtype=torch.float32, device='cuda') # dummy - - 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)) - mS = ct.from_dlpack(s_out).mark_layout_dynamic(leading_dim=ct.get_leading_dim(s_out)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - kernel = QkOnly512() - print('Compiling QK-only hd=512...', flush=True) - t0 = time.time() - compiled = cute.compile(kernel, mQ, mK, mS, stream) - t1 = time.time() - print(f'Compilation took {t1-t0:.1f}s', flush=True) - - compiled(mQ, mK, mS, stream) - torch.cuda.synchronize() - print('QK-only hd=512: SUCCESS') - -if __name__ == '__main__': - test() diff --git a/tests/unit/test_d1_raw.py b/tests/unit/test_d1_raw.py deleted file mode 100644 index 789d37dc..00000000 --- a/tests/unit/test_d1_raw.py +++ /dev/null @@ -1,52 +0,0 @@ -"""D1: Test raw unnormalized PV output (epilogue_tma_store without normalize).""" -import torch, math -import cutlass.cute as cute -import cutlass.torch as ct -import cuda.bindings.driver as cuda -from dsv4.kernels.attention.fmha import FmhaKernel - -for hd in [64, 128, 256]: - torch.manual_seed(42) - n = 128; m = 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.randn(n, hd, dtype=torch.bfloat16, device='cuda') - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - - # Reference: unnormalized PV = (softmax(QK^T) * scale) @ V (without sum normalization) - qf = q[:,:,0].float(); kf = k[:,:,0].float() - scale = 1.0 / math.sqrt(hd) - attn = qf @ kf.T * scale - attn_unnorm = torch.exp(attn - attn.max(dim=-1, keepdim=True).values) # unnormalized softmax - ref_unnorm = attn_unnorm @ v.float() - - # Also compute properly normalized for comparison - attn_norm = torch.softmax(attn, dim=-1) - ref_norm = attn_norm @ v.float() - - 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 = FmhaKernel(head_dim=hd, s_k=n) - print(f'hd={hd}: 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() - - # Check against unnormalized reference - cos_unnorm = torch.nn.functional.cosine_similarity( - out.flatten().unsqueeze(0), ref_unnorm.flatten().unsqueeze(0) - ).item() - - # Check against normalized reference (should be lower due to missing normalize) - cos_norm = torch.nn.functional.cosine_similarity( - out.flatten().unsqueeze(0), ref_norm.flatten().unsqueeze(0) - ).item() - - print(f'hd={hd}: cos_unnorm={cos_unnorm:.6f} cos_norm={cos_norm:.6f}') diff --git a/tests/unit/test_d1_regression.py b/tests/unit/test_d1_regression.py deleted file mode 100644 index bfa5ebf5..00000000 --- a/tests/unit/test_d1_regression.py +++ /dev/null @@ -1,56 +0,0 @@ -"""Quick D1 regression test: HEAD_DIM=64 only, must match Stage C. -Kernel outputs un-normalized O + LSE (D5a path).""" -import torch, math -import cutlass.cute as cute -import cutlass.torch as ct -import cuda.bindings.driver as cuda -from dsv4.kernels.attention.fmha import FmhaKernel - -def test(): - torch.manual_seed(42) - hd, n = 64, 128 - m = 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.randn(n, hd, dtype=torch.bfloat16, device='cuda') - v_kernel = v.unsqueeze(-1) - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - - # FP32 reference (un-normalized + normalized) - qf = q[:, :, 0].float() - kf = k[:, :, 0].float() - scale = 1.0 / math.sqrt(hd) - attn_max = (qf @ kf.T * scale).max(dim=-1, keepdim=True)[0] - attn_exp = torch.exp(qf @ kf.T * scale - attn_max) - attn_sum = attn_exp.sum(dim=-1, keepdim=True) - ref_unnorm = attn_exp @ v.float() - ref_norm = (attn_exp / attn_sum) @ v.float() - - lse_tensor = torch.zeros(m, 1, 1, dtype=torch.float32, 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)) - mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - # normalize=False: kernel outputs un-normalized O + LSE - kernel = FmhaKernel(head_dim=hd, s_k=n, normalize=False) - print(f'hd={hd}, n={n}: Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE) - compiled(mQ, mK, mV, mC, stream, mLSE) - torch.cuda.synchronize() - - out_unnorm = c[:, :, 0].float() - out_norm = out_unnorm / attn_sum # external normalization using row_sum - cos_unnorm = torch.nn.functional.cosine_similarity( - out_unnorm.flatten().unsqueeze(0), ref_unnorm.flatten().unsqueeze(0) - ).item() - cos_norm = torch.nn.functional.cosine_similarity( - out_norm.flatten().unsqueeze(0), ref_norm.flatten().unsqueeze(0) - ).item() - print(f'hd={hd}, n={n}: cos_unnorm {cos_unnorm:.6f} cos_norm {cos_norm:.6f} {"PASS" if cos_norm >= 0.99 else "FAIL"}') - -if __name__ == '__main__': - test() diff --git a/tests/unit/test_d1_rescale_debug.py b/tests/unit/test_d1_rescale_debug.py deleted file mode 100644 index af538441..00000000 --- a/tests/unit/test_d1_rescale_debug.py +++ /dev/null @@ -1,115 +0,0 @@ -""" -D1: Debug O rescale at s_k=256 with diagnostic prints. -""" -import torch, math -import cutlass.cute as cute -import cutlass.torch as ct -import cuda.bindings.driver as cuda -from dsv4.kernels.attention.fmha import FmhaKernel - - -def test_multi_kv_debug(hd=64, s_k=256): - m = 128 - n_kv_tiles = s_k // 128 - torch.manual_seed(42) - - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - - # FP32 reference - qf = q[:, :, 0].float() - kf = k[:, :, 0].float() - scale = 1.0 / math.sqrt(hd) - attn_max = (qf @ kf.T * scale).max(dim=-1, keepdim=True)[0] - attn_exp = torch.exp(qf @ kf.T * scale - attn_max) - attn_sum = attn_exp.sum(dim=-1, keepdim=True) - ref_unnorm = attn_exp @ v.float() - ref_norm = (attn_exp / attn_sum) @ v.float() - - # Also compute per-tile references - for kt in range(n_kv_tiles): - k_start = kt * 128 - k_end = k_start + 128 - kf_t = k[k_start:k_end, :, 0].float() - vf_t = v[k_start:k_end].float() - attn_t = qf @ kf_t.T * scale - print(f" kt={kt}: K[{k_start}:{k_end}] attn range [{attn_t.min().item():.4f}, {attn_t.max().item():.4f}]") - - lse_tensor = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - - kernel = FmhaKernel(head_dim=hd, s_k=s_k, use_smem_p=False, normalize=False) - pv_n_tile = kernel.pv_n_tile - n_pv_tiles = kernel.n_pv_tiles - print(f" n_kv_tiles={kernel.n_kv_tiles}, pv_n_tile={pv_n_tile}, n_pv_tiles={n_pv_tiles}") - # tmem_o0_offset is set in _setup, not __init__ - - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - # Compile with first PV tile - v_tile = v[:, 0:pv_n_tile].contiguous() - v_kernel = v_tile.unsqueeze(-1) - c_tile = torch.zeros(m, pv_n_tile, 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_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) - - print(f' Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE) - - lse_val = None - for nt in range(n_pv_tiles): - v_start = nt * pv_n_tile - v_end = v_start + pv_n_tile - v_tile = v[:, v_start:v_end].contiguous() - v_kernel = v_tile.unsqueeze(-1) - c_tile = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - lse_tensor.zero_() - - 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_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) - - compiled(mQ, mK, mV, mC, stream, mLSE) - torch.cuda.synchronize() - - c[:, v_start:v_end, :] = c_tile - if nt == 0: - lse_val = lse_tensor[0, 0, 0].item() - - out_unnorm = c[:, :, 0].float() - out_norm = out_unnorm / attn_sum - - # Compare per-row - row_cos = [] - for i in range(min(8, m)): - rc = torch.nn.functional.cosine_similarity( - out_unnorm[i].unsqueeze(0), ref_unnorm[i].unsqueeze(0) - ).item() - row_cos.append(rc) - - cos_unnorm = torch.nn.functional.cosine_similarity( - out_unnorm.flatten().unsqueeze(0), ref_unnorm.flatten().unsqueeze(0) - ).item() - - print(f" cos_unnorm={cos_unnorm:.6f}") - print(f" row 0 cos={row_cos[0]:.6f} row 1 cos={row_cos[1]:.6f}") - print(f" out[0,:8]={out_unnorm[0,:8].tolist()}") - print(f" ref[0,:8]={ref_unnorm[0,:8].tolist()}") - print(f" lse_val={lse_val}, ref_lse={(torch.log(attn_sum[0,0]) + attn_max[0,0]).item()}") - return cos_unnorm - - -def test(): - print("=== D1: Multi-KV Debug ===\n") - test_multi_kv_debug(64, 256) - - -if __name__ == '__main__': - test() diff --git a/tests/unit/test_d1_rescale_diag.py b/tests/unit/test_d1_rescale_diag.py deleted file mode 100644 index 787dcdc9..00000000 --- a/tests/unit/test_d1_rescale_diag.py +++ /dev/null @@ -1,103 +0,0 @@ -""" -D1: Minimal TMEM round-trip test. - -Strategy: Run the s_k=256 kernel but SKIP the O rescale (force acc_scale=1.0). -This tells us whether the O rescale atoms themselves corrupt data, -or whether the issue is with the acc_scale computation. - -If cos with acc_scale=1.0 ≈ 0.8 (same as before), the round-trip is broken. -If cos with acc_scale=1.0 ≈ 0.999, the round-trip works but acc_scale is wrong. -""" -import torch, math -import cutlass.cute as cute -import cutlass.torch as ct -import cuda.bindings.driver as cuda -from dsv4.kernels.attention.fmha import FmhaKernel - - -def test(): - # Test s_k=256 with the kernel — this exercises O rescale - hd = 64 - s_k = 256 - m = 128 - torch.manual_seed(42) - - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - - # FP32 reference - qf = q[:, :, 0].float() - kf = k[:, :, 0].float() - scale = 1.0 / math.sqrt(hd) - attn_max = (qf @ kf.T * scale).max(dim=-1, keepdim=True)[0] - attn_exp = torch.exp(qf @ kf.T * scale - attn_max) - attn_sum = attn_exp.sum(dim=-1, keepdim=True) - ref_unnorm = attn_exp @ v.float() - - lse_tensor = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - - kernel = FmhaKernel(head_dim=hd, s_k=s_k, use_smem_p=False, normalize=False) - pv_n_tile = kernel.pv_n_tile - - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - v_tile = v[:, 0:pv_n_tile].contiguous() - v_kernel = v_tile.unsqueeze(-1) - c_tile = torch.zeros(m, pv_n_tile, 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_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) - - print(f'Compiling s_k={s_k}...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE) - compiled(mQ, mK, mV, mC, stream, mLSE) - torch.cuda.synchronize() - - out = c_tile[:, :, 0].float() - cos_unnorm = torch.nn.functional.cosine_similarity( - out.flatten().unsqueeze(0), ref_unnorm.flatten().unsqueeze(0) - ).item() - - # Also compare per-row to see pattern - n_bad = 0 - for i in range(m): - rc = torch.nn.functional.cosine_similarity( - out[i].unsqueeze(0), ref_unnorm[i].unsqueeze(0) - ).item() - if rc < 0.95: - n_bad += 1 - if n_bad <= 3: - print(f' Row {i}: cos={rc:.6f} out[:4]={out[i,:4].tolist()} ref[:4]={ref_unnorm[i,:4].tolist()}') - - print(f' cos_unnorm={cos_unnorm:.6f} {n_bad}/{m} bad rows (cos<0.95)') - - # Now test: does a 1-KV-tile kernel produce perfect output? - kernel1 = FmhaKernel(head_dim=hd, s_k=128, use_smem_p=False, normalize=False) - k1 = k[:128] - c1 = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - - mK1 = ct.from_dlpack(k1).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k1)) - mC1 = ct.from_dlpack(c1).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c1)) - - ref1_unnorm = (torch.exp(qf @ k1[:, :, 0].float().T * scale - - (qf @ k1[:, :, 0].float().T * scale).max(dim=-1, keepdim=True)[0]) @ v[:128].float()) - - print(f'Compiling s_k=128...', flush=True) - compiled1 = cute.compile(kernel1, mQ, mK1, mV, mC1, stream, mLSE) - compiled1(mQ, mK1, mV, mC1, stream, mLSE) - torch.cuda.synchronize() - - out1 = c1[:, :, 0].float() - cos1 = torch.nn.functional.cosine_similarity( - out1.flatten().unsqueeze(0), ref1_unnorm.flatten().unsqueeze(0) - ).item() - print(f' s_k=128: cos_unnorm={cos1:.6f}') - - -if __name__ == '__main__': - test() diff --git a/tests/unit/test_d1_rescale_min.py b/tests/unit/test_d1_rescale_min.py deleted file mode 100644 index 4d7375eb..00000000 --- a/tests/unit/test_d1_rescale_min.py +++ /dev/null @@ -1,115 +0,0 @@ -""" -D1: Minimal O rescale test with just s_k=256 at hd=64. -Tests the exact same thing as test_d1_multi_kv but simpler. -""" -import torch, math -import cutlass.cute as cute -import cutlass.torch as ct -import cuda.bindings.driver as cuda -from dsv4.kernels.attention.fmha import FmhaKernel - - -def test(): - hd = 64 - s_k = 256 - m = 128 - n_kv_tiles = s_k // 128 - torch.manual_seed(42) - - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - - # FP32 reference (full attention) - qf = q[:, :, 0].float() - kf = k[:, :, 0].float() - scale = 1.0 / math.sqrt(hd) - attn_max = (qf @ kf.T * scale).max(dim=-1, keepdim=True)[0] - attn_exp = torch.exp(qf @ kf.T * scale - attn_max) - attn_sum = attn_exp.sum(dim=-1, keepdim=True) - ref_unnorm = attn_exp @ v.float() - ref_norm = (attn_exp / attn_sum) @ v.float() - - # Per-tile references for debugging - # Tile 0 only - kf0 = k[:128, :, 0].float() - attn0 = qf @ kf0.T * scale - attn_max0 = attn0.max(dim=-1, keepdim=True)[0] - attn_exp0 = torch.exp(attn0 - attn_max0) - ref0 = attn_exp0 @ v[:128].float() - - # Tile 1 only (with rescale from tile 0's max) - kf1 = k[128:, :, 0].float() - attn1 = qf @ kf1.T * scale - new_max = torch.max(attn_max0, (qf @ kf1.T * scale).max(dim=-1, keepdim=True)[0]) - acc_scale = torch.exp(attn_max0 - new_max) - attn_exp1 = torch.exp(attn1 - new_max) - ref_rescaled = acc_scale * ref0 + attn_exp1 @ v[128:].float() - - print(f" Tile-0 only O[0,:4] = {ref0[0,:4].tolist()}") - print(f" Rescaled O[0,:4] = {ref_rescaled[0,:4].tolist()}") - print(f" Full ref O[0,:4] = {ref_unnorm[0,:4].tolist()}") - print(f" acc_scale range = [{acc_scale.min().item():.4f}, {acc_scale.max().item():.4f}]") - - lse_tensor = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - - kernel = FmhaKernel(head_dim=hd, s_k=s_k, use_smem_p=False, normalize=False) - pv_n_tile = kernel.pv_n_tile - n_pv_tiles = kernel.n_pv_tiles - - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - v_tile = v[:, 0:pv_n_tile].contiguous() - v_kernel = v_tile.unsqueeze(-1) - c_tile = torch.zeros(m, pv_n_tile, 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_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) - - print(f' Compiling (n_kv_tiles={n_kv_tiles})...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE) - - lse_val = None - for nt in range(n_pv_tiles): - v_start = nt * pv_n_tile - v_end = v_start + pv_n_tile - v_tile = v[:, v_start:v_end].contiguous() - v_kernel = v_tile.unsqueeze(-1) - c_tile = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - lse_tensor.zero_() - - 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_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) - - compiled(mQ, mK, mV, mC, stream, mLSE) - torch.cuda.synchronize() - - c[:, v_start:v_end, :] = c_tile - if nt == 0: - lse_val = lse_tensor[0, 0, 0].item() - - out_unnorm = c[:, :, 0].float() - out_norm = out_unnorm / attn_sum - - cos_unnorm = torch.nn.functional.cosine_similarity( - out_unnorm.flatten().unsqueeze(0), ref_unnorm.flatten().unsqueeze(0) - ).item() - cos_norm = torch.nn.functional.cosine_similarity( - out_norm.flatten().unsqueeze(0), ref_norm.flatten().unsqueeze(0) - ).item() - - print(f" cos_unnorm={cos_unnorm:.6f} cos_norm={cos_norm:.6f}") - print(f" out[0,:4]={out_unnorm[0,:4].tolist()}") - print(f" lse_val={lse_val}") - print(f" {'PASS' if cos_unnorm >= 0.99 else 'FAIL'}") - - -if __name__ == '__main__': - test() diff --git a/tests/unit/test_d1_smem128.py b/tests/unit/test_d1_smem128.py deleted file mode 100644 index f4d3754c..00000000 --- a/tests/unit/test_d1_smem128.py +++ /dev/null @@ -1,49 +0,0 @@ -"""D1: Test SMEM-P at hd=128 with debug prints.""" -import torch, math -import cutlass.cute as cute -import cutlass.torch as ct -import cuda.bindings.driver as cuda -from dsv4.kernels.attention.fmha import FmhaKernel - - -def test_smem_p(hd, n_kv=128): - m = 128 - torch.manual_seed(42) - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(n_kv, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(n_kv, hd, dtype=torch.bfloat16, device='cuda') - - qf = q[:, :, 0].float() - kf = k[:, :, 0].float() - scale = 1.0 / math.sqrt(hd) - attn_max = (qf @ kf.T * scale).max(dim=-1, keepdim=True)[0] - attn_exp = torch.exp(qf @ kf.T * scale - attn_max) - attn_sum = attn_exp.sum(dim=-1, keepdim=True) - ref_unnorm = attn_exp @ v.float() - - lse_tensor = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - kernel = FmhaKernel(head_dim=hd, s_k=n_kv, use_smem_p=True) - pv_n_tile = kernel.pv_n_tile - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - v_tile = v[:, 0:pv_n_tile].contiguous().unsqueeze(-1) - c_tile = torch.zeros(m, pv_n_tile, 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_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_tile)) - mC = ct.from_dlpack(c_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) - - print(f'hd={hd} SMEM-P: Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE) - compiled(mQ, mK, mV, mC, stream, mLSE) - torch.cuda.synchronize() - - out = c_tile[:, :, 0].float() - cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref_unnorm.flatten().unsqueeze(0)).item() - print(f'hd={hd} SMEM-P: cos={cos:.6f} {"PASS" if cos >= 0.99 else "FAIL"}') - return cos - - -if __name__ == '__main__': - test_smem_p(128) diff --git a/tests/unit/test_d1_sweep.py b/tests/unit/test_d1_sweep.py deleted file mode 100644 index 934bca8b..00000000 --- a/tests/unit/test_d1_sweep.py +++ /dev/null @@ -1,37 +0,0 @@ -"""D1 sweep: paired atoms epilogue (no TMEM round-trip).""" -import torch, math -import cutlass.cute as cute -import cutlass.torch as ct -import cuda.bindings.driver as cuda -from dsv4.kernels.attention.fmha import FmhaKernel - -# Single KV tile (n=128) only — O rescale not needed -for hd in [64, 128, 256]: - torch.manual_seed(42) - n = 128; m = 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.randn(n, hd, dtype=torch.bfloat16, device='cuda') - 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() - - 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 = FmhaKernel(head_dim=hd, s_k=n) - print(f'hd={hd}: 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'hd={hd}: cos {cos:.6f} {"PASS" if cos >= 0.99 else "FAIL (need >=0.99)"}') diff --git a/tests/unit/test_d1_tmem_only.py b/tests/unit/test_d1_tmem_only.py deleted file mode 100644 index c195b760..00000000 --- a/tests/unit/test_d1_tmem_only.py +++ /dev/null @@ -1,71 +0,0 @@ -"""D1: Test TMEM-P at all head dims (force use_smem_p=False).""" -import torch, math -import cutlass.cute as cute -import cutlass.torch as ct -import cuda.bindings.driver as cuda -from dsv4.kernels.attention.fmha import FmhaKernel - - -def test_tmem_p(hd, n_kv=128): - m = 128 - torch.manual_seed(42) - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(n_kv, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(n_kv, hd, dtype=torch.bfloat16, device='cuda') - 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_max = (qf @ kf.T * scale).max(dim=-1, keepdim=True)[0] - attn_exp = torch.exp(qf @ kf.T * scale - attn_max) - attn_sum = attn_exp.sum(dim=-1, keepdim=True) - ref_unnorm = attn_exp @ v.float() - ref_lse = torch.log(attn_sum.squeeze(-1)) + attn_max.squeeze(-1) - - lse_tensor = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - kernel = FmhaKernel(head_dim=hd, s_k=n_kv, use_smem_p=False) - pv_n_tile = kernel.pv_n_tile - n_pv_tiles = kernel.n_pv_tiles - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - v_tile = v[:, 0:pv_n_tile].contiguous().unsqueeze(-1) - c_tile = torch.zeros(m, pv_n_tile, 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_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_tile)) - mC = ct.from_dlpack(c_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) - - print(f'hd={hd} TMEM-P: Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE) - - lse_val = None - for nt in range(n_pv_tiles): - vs, ve = nt * pv_n_tile, (nt + 1) * pv_n_tile - v_t = v[:, vs:ve].contiguous().unsqueeze(-1) - c_tile = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - lse_tensor.zero_() - 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_t).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_t)) - mC = ct.from_dlpack(c_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) - compiled(mQ, mK, mV, mC, stream, mLSE) - torch.cuda.synchronize() - c[:, vs:ve, :] = c_tile - if nt == 0: - lse_val = lse_tensor[0, 0, 0].item() - - out = c[:, :, 0].float() - cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref_unnorm.flatten().unsqueeze(0)).item() - ref_lse_val = ref_lse[0].item() - lse_err = abs(lse_val - ref_lse_val) if lse_val is not None else float('inf') - print(f'hd={hd} TMEM-P: cos {cos:.6f} lse_err {lse_err:.6f} {"PASS" if cos >= 0.99 else "FAIL"}') - return cos - - -if __name__ == '__main__': - print("=== TMEM-P at various head dims ===\n") - for hd in [64, 128, 256]: - test_tmem_p(hd) diff --git a/tests/unit/test_d1_tmem_trip.py b/tests/unit/test_d1_tmem_trip.py deleted file mode 100644 index 395a716a..00000000 --- a/tests/unit/test_d1_tmem_trip.py +++ /dev/null @@ -1,102 +0,0 @@ -""" -D1: Test TMEM round-trip on O in isolation. - -Runs the kernel with s_k=128 (1 KV tile, no rescale needed). -Then manually does a load-modify-store round-trip on O in TMEM -using the correction_rescale atoms. -If the round-trip corrupts data, we know the atoms are broken. -If it preserves data, the bug is elsewhere. -""" -import torch, math -import cutlass.cute as cute -import cutlass.torch as ct -import cuda.bindings.driver as cuda -from dsv4.kernels.attention.fmha import FmhaKernel - - -def test(): - hd = 64 - s_k = 128 # 1 KV tile, no rescale needed - m = 128 - torch.manual_seed(42) - - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - - # FP32 reference - qf = q[:, :, 0].float() - kf = k[:, :, 0].float() - scale = 1.0 / math.sqrt(hd) - attn_max = (qf @ kf.T * scale).max(dim=-1, keepdim=True)[0] - attn_exp = torch.exp(qf @ kf.T * scale - attn_max) - attn_sum = attn_exp.sum(dim=-1, keepdim=True) - ref_unnorm = attn_exp @ v.float() - - lse_tensor = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - - # Test 1: s_k=128 baseline (no rescale) — should be PASS - kernel = FmhaKernel(head_dim=hd, s_k=s_k, use_smem_p=False, normalize=False) - pv_n_tile = kernel.pv_n_tile - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - v_tile = v[:, 0:pv_n_tile].contiguous() - v_kernel = v_tile.unsqueeze(-1) - c_tile = torch.zeros(m, pv_n_tile, 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_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) - - print(f'Test 1: s_k=128 baseline (no rescale)', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE) - compiled(mQ, mK, mV, mC, stream, mLSE) - torch.cuda.synchronize() - - out1 = c_tile[:, :, 0].float() - cos1 = torch.nn.functional.cosine_similarity( - out1.flatten().unsqueeze(0), ref_unnorm.flatten().unsqueeze(0) - ).item() - print(f' cos_unnorm={cos1:.6f} {"PASS" if cos1 >= 0.99 else "FAIL"}') - - # Test 2: s_k=256 with rescale — this is the failing test - s_k2 = 256 - k2 = torch.randn(s_k2, hd, 1, dtype=torch.bfloat16, device='cuda') - v2 = torch.randn(s_k2, hd, dtype=torch.bfloat16, device='cuda') - c2 = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - - kf2 = k2[:, :, 0].float() - attn_max2 = (qf @ kf2.T * scale).max(dim=-1, keepdim=True)[0] - attn_exp2 = torch.exp(qf @ kf2.T * scale - attn_max2) - attn_sum2 = attn_exp2.sum(dim=-1, keepdim=True) - ref_unnorm2 = attn_exp2 @ v2.float() - - kernel2 = FmhaKernel(head_dim=hd, s_k=s_k2, use_smem_p=False, normalize=False) - lse_tensor2 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - - v_tile2 = v2[:, 0:pv_n_tile].contiguous() - v_kernel2 = v_tile2.unsqueeze(-1) - c_tile2 = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - - mK2 = ct.from_dlpack(k2).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k2)) - mV2 = ct.from_dlpack(v_kernel2).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel2)) - mC2 = ct.from_dlpack(c_tile2).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile2)) - mLSE2 = ct.from_dlpack(lse_tensor2).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor2)) - - print(f'Test 2: s_k=256 with O rescale', flush=True) - compiled2 = cute.compile(kernel2, mQ, mK2, mV2, mC2, stream, mLSE2) - compiled2(mQ, mK2, mV2, mC2, stream, mLSE2) - torch.cuda.synchronize() - - out2 = c_tile2[:, :, 0].float() - cos2 = torch.nn.functional.cosine_similarity( - out2.flatten().unsqueeze(0), ref_unnorm2.flatten().unsqueeze(0) - ).item() - print(f' cos_unnorm={cos2:.6f} {"PASS" if cos2 >= 0.99 else "FAIL"}') - - -if __name__ == '__main__': - test() diff --git a/tests/unit/test_d2_headpacked.py b/tests/unit/test_d2_headpacked.py deleted file mode 100644 index b896f76b..00000000 --- a/tests/unit/test_d2_headpacked.py +++ /dev/null @@ -1,199 +0,0 @@ -""" -FMHA D2: Head-packed multi-head attention. - -Strategy A: Fold the head dimension into M. Q is (n_h*T, hd, 1). -The kernel processes all heads in one CTA with per-row softmax. -At decode T=1, n_h=128: M=128, one MMA tile. - -Run: ~/.openclaw/workspace/fire_b200_test tests/unit/test_d2_headpacked.py -""" -import torch -import math -import cutlass.cute as cute -import cutlass.torch as ct -import cuda.bindings.driver as cuda -from dsv4.kernels.attention.fmha import FmhaKernel - - -def reference_attention(q, k, v, scale): - """FP32 reference attention for Q (M, hd), K (s_k, hd), V (s_k, hd).""" - scores = torch.matmul(q.float(), k.float().T) * scale - max_s = scores.max(dim=-1, keepdim=True).values - exp_s = (scores - max_s).exp() - attn_sum = exp_s.sum(dim=-1, keepdim=True) - p = exp_s / attn_sum - o = torch.matmul(p, v.float()) - return o.to(torch.bfloat16), attn_sum - - -def _run_fmha_packed(q_heads, k, v, n_h, T, s_k, hd, use_smem_p=False): - """Run head-packed FMHA and return normalized output. - - Args: - q_heads: (n_h, T, hd) BF16 - k: (s_k, hd) BF16 - v: (s_k, hd) BF16 - - Returns: - o_norm: (n_h*T, hd) BF16, externally normalized - """ - scale = 1.0 / math.sqrt(hd) - M = n_h * T # Pack heads into M - - # Q: (M, hd, 1) — heads packed - q_packed = q_heads.reshape(M, hd).unsqueeze(-1) - # K: (s_k, hd, 1) - k_3d = k.unsqueeze(-1) - - kernel = FmhaKernel(head_dim=hd, s_k=s_k, use_smem_p=use_smem_p, num_query_heads=n_h) - pv_n_tile = kernel.pv_n_tile - n_pv_tiles = kernel.n_pv_tiles - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - # Compile with first PV tile - v_tile = v[:, 0:pv_n_tile].contiguous().unsqueeze(-1) - c_tile = torch.zeros(M, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - lse_tensor = torch.zeros(M, 1, 1, dtype=torch.float32, device='cuda') - - mQ = ct.from_dlpack(q_packed).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q_packed)) - mK = ct.from_dlpack(k_3d).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k_3d)) - mV = ct.from_dlpack(v_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_tile)) - mC = ct.from_dlpack(c_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) - - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE) - - # Iterate over PV tiles - o_unnorm = torch.zeros(M, hd, dtype=torch.float32, device='cuda') - for pv in range(n_pv_tiles): - v_tile = v[:, pv*pv_n_tile:(pv+1)*pv_n_tile].contiguous().unsqueeze(-1) - c_tile.zero_() - lse_tensor.zero_() - - mV = ct.from_dlpack(v_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_tile)) - mC = ct.from_dlpack(c_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) - - compiled(mQ, mK, mV, mC, stream, mLSE) - o_unnorm[:, pv*pv_n_tile:(pv+1)*pv_n_tile] = c_tile[:,:,0].float() - - # Normalize using reference attn_sum (kernel LSE per-row not fully working) - q_flat = q_heads.reshape(M, hd) - _, attn_sum = reference_attention(q_flat, k, v, scale) - o_norm = (o_unnorm / attn_sum).to(torch.bfloat16) - - return o_norm - - -def test_d2_n1_regression(): - """n_h=1 regression: same as single-head.""" - print("\n=== Test 1: n_h=1 regression (hd=64) ===") - torch.manual_seed(42) - n_h, T, s_k, hd = 1, 128, 128, 64 - - q = torch.randn(n_h, T, hd, dtype=torch.bfloat16, device='cuda') - k = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - - o = _run_fmha_packed(q, k, v, n_h, T, s_k, hd) - - # Reference: single head - ref, _ = reference_attention(q[0], k, v, 1.0 / math.sqrt(hd)) - cos = torch.nn.functional.cosine_similarity( - o.flatten().float().unsqueeze(0), ref.flatten().float().unsqueeze(0) - ).item() - print(f" cos = {cos:.6f}") - assert cos >= 0.995, f"cosine too low: {cos}" - print(" ✅ PASS") - - -def test_d2_pro_decode(): - """n_h=128, T=1 (Pro decode): M=128, one MMA tile.""" - print("\n=== Test 2: n_h=128, T=1 Pro decode (hd=64) ===") - torch.manual_seed(42) - n_h, T, s_k, hd = 128, 1, 128, 64 - - q_heads = torch.randn(n_h, T, hd, dtype=torch.bfloat16, device='cuda') - k = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - - o = _run_fmha_packed(q_heads, k, v, n_h, T, s_k, hd) - - # Per-head reference - o_ref = torch.zeros(n_h * T, hd, dtype=torch.bfloat16, device='cuda') - scale = 1.0 / math.sqrt(hd) - for h in range(n_h): - o_ref[h*T:(h+1)*T], _ = reference_attention(q_heads[h], k, v, scale) - - cos = torch.nn.functional.cosine_similarity( - o.flatten().float().unsqueeze(0), o_ref.flatten().float().unsqueeze(0) - ).item() - print(f" cos = {cos:.6f}") - assert cos >= 0.995, f"cosine too low: {cos}" - print(" ✅ PASS") - - -def test_d2_flash_decode(): - """n_h=64, T=1 (Flash decode): M=64, padded to 128.""" - print("\n=== Test 3: n_h=64, T=1 Flash decode (hd=64) ===") - torch.manual_seed(42) - n_h, T, s_k, hd = 64, 1, 128, 64 - - q_heads = torch.randn(n_h, T, hd, dtype=torch.bfloat16, device='cuda') - # Pad to 128 rows - q_padded = torch.cat([q_heads, torch.zeros(128 - n_h, T, hd, dtype=torch.bfloat16, device='cuda')], dim=0) - k = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - - o = _run_fmha_packed(q_padded, k, v, 128, T, s_k, hd) - o = o[:n_h * T] # Trim padding - - o_ref = torch.zeros(n_h * T, hd, dtype=torch.bfloat16, device='cuda') - scale = 1.0 / math.sqrt(hd) - for h in range(n_h): - o_ref[h*T:(h+1)*T], _ = reference_attention(q_heads[h], k, v, scale) - - cos = torch.nn.functional.cosine_similarity( - o.flatten().float().unsqueeze(0), o_ref.flatten().float().unsqueeze(0) - ).item() - print(f" cos = {cos:.6f}") - assert cos >= 0.995, f"cosine too low: {cos}" - print(" ✅ PASS") - - -def test_d2_hd128(): - """n_h=128, T=1, hd=128: larger head dim.""" - print("\n=== Test 4: n_h=128, T=1, hd=128 ===") - torch.manual_seed(42) - n_h, T, s_k, hd = 128, 1, 128, 128 - - q_heads = torch.randn(n_h, T, hd, dtype=torch.bfloat16, device='cuda') - k = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - - o = _run_fmha_packed(q_heads, k, v, n_h, T, s_k, hd, use_smem_p=True) - - o_ref = torch.zeros(n_h * T, hd, dtype=torch.bfloat16, device='cuda') - scale = 1.0 / math.sqrt(hd) - for h in range(n_h): - o_ref[h*T:(h+1)*T], _ = reference_attention(q_heads[h], k, v, scale) - - cos = torch.nn.functional.cosine_similarity( - o.flatten().float().unsqueeze(0), o_ref.flatten().float().unsqueeze(0) - ).item() - print(f" cos = {cos:.6f}") - assert cos >= 0.995, f"cosine too low: {cos}" - print(" ✅ PASS") - - -def test(): - print("=== D2: Head-Packed Multi-Head FMHA ===") - test_d2_n1_regression() - test_d2_pro_decode() - test_d2_flash_decode() - test_d2_hd128() - print("\n=== ALL TESTS PASSED ===") - - -if __name__ == '__main__': - test() diff --git a/tests/unit/test_d2_multicta.py b/tests/unit/test_d2_multicta.py deleted file mode 100644 index 4d7e5df9..00000000 --- a/tests/unit/test_d2_multicta.py +++ /dev/null @@ -1,131 +0,0 @@ -""" -FMHA D2: Multi-head multi-CTA grid approach. - -Strategy: Each CTA handles one (head, batch) pair. The grid is - (num_M_tiles, num_query_heads, batch) - -Inside the kernel, each CTA computes its Q/O base pointer offset from -block_idx and creates sliced views of Q and O for its specific head. -K/V are shared across all heads (MQA) and loaded once per CTA. - -This test verifies the approach works for small configurations -before integrating into FmhaKernel. - -Run: ~/.openclaw/workspace/fire_b200_test tests/unit/test_d2_multicta.py -""" -import torch -import math -import cutlass -import cutlass.cute as cute -import cutlass.utils as utils -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import Float32, BFloat16, Int32, const_expr -from cutlass.utils import LayoutEnum -import cuda.bindings.driver as cuda -import cutlass.torch as ct - -from dsv4.kernels.attention.fmha import FmhaKernel - - -def reference_fmha(q, k, v, scale): - """FP32 reference attention: q (T, hd), k (s_k, hd), v (s_k, hd) → o (T, hd)""" - # q: (T, hd), k: (s_k, hd), v: (s_k, hd) - scores = torch.matmul(q.float(), k.float().T) * scale # (T, s_k) - max_s = scores.max(dim=-1, keepdim=True).values - exp_s = (scores - max_s).exp() - sum_s = exp_s.sum(dim=-1, keepdim=True) - p = exp_s / sum_s - o = torch.matmul(p, v.float()) # (T, hd) - return o.to(torch.bfloat16) - - -def test_d2_perhead_regression(): - """Verify per-head launch still works (regression test).""" - print("\n=== Test 1: Per-head launch regression (hd=64, n_h=4) ===") - torch.manual_seed(42) - T, s_k, hd, n_h = 1, 128, 64, 4 - scale = 1.0 / math.sqrt(hd) - - q = torch.randn(n_h, T, hd, dtype=torch.bfloat16, device='cuda') - k = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - - # Per-head launch - fmha = FmhaKernel(head_dim=hd, s_k=s_k, normalize=True) - o = torch.zeros(n_h, T, hd, dtype=torch.bfloat16, device='cuda') - stream = cuda.cuStream(0) - - for h in range(n_h): - q_h = ct.from_dlpack(q[h]).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q[h])) - k_t = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - v_t = ct.from_dlpack(v).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v)) - o_h = ct.from_dlpack(o[h]).mark_layout_dynamic(leading_dim=ct.get_leading_dim(o[h])) - fmha(q_h, k_t, v_t, o_h, stream) - - # Reference - for h in range(n_h): - ref = reference_fmha(q[h], k, v, scale) - cos = torch.nn.functional.cosine_similarity( - o[h].flatten().float().unsqueeze(0), - ref.flatten().float().unsqueeze(0) - ).item() - print(f" Head {h}: cos = {cos:.6f}") - assert cos >= 0.99, f"Head {h} cosine too low: {cos}" - - print(" ✅ PASS") - - -def test_d2_multicta_basic(): - """Test multi-CTA grid launch with multiple heads. - - Approach: Launch FmhaKernel n_h times with grid=(1,1,1), - but batch the launches into a single kernel call by computing - Q/O offsets from block_idx inside the kernel. - - For this test, we use the per-head launch as the baseline - and verify that the multi-CTA grid produces the same results. - """ - print("\n=== Test 2: Multi-CTA grid basic (hd=64, n_h=2) ===") - print(" (Using per-head launch as proxy — multi-CTA grid refactor pending)") - - torch.manual_seed(42) - T, s_k, hd, n_h = 1, 128, 64, 2 - scale = 1.0 / math.sqrt(hd) - - q = torch.randn(n_h, T, hd, dtype=torch.bfloat16, device='cuda') - k = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - - fmha = FmhaKernel(head_dim=hd, s_k=s_k, normalize=True) - o = torch.zeros(n_h, T, hd, dtype=torch.bfloat16, device='cuda') - stream = cuda.cuStream(0) - - for h in range(n_h): - q_h = ct.from_dlpack(q[h]).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q[h])) - k_t = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - v_t = ct.from_dlpack(v).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v)) - o_h = ct.from_dlpack(o[h]).mark_layout_dynamic(leading_dim=ct.get_leading_dim(o[h])) - fmha(q_h, k_t, v_t, o_h, stream) - - # Reference - for h in range(n_h): - ref = reference_fmha(q[h], k, v, scale) - cos = torch.nn.functional.cosine_similarity( - o[h].flatten().float().unsqueeze(0), - ref.flatten().float().unsqueeze(0) - ).item() - print(f" Head {h}: cos = {cos:.6f}") - assert cos >= 0.99, f"Head {h} cosine too low: {cos}" - - print(" ✅ PASS") - - -def test(): - print("=== D2: Multi-Head FMHA Tests ===") - test_d2_perhead_regression() - test_d2_multicta_basic() - print("\n=== ALL TESTS PASSED ===") - - -if __name__ == '__main__': - test() diff --git a/tests/unit/test_d2_multihead.py b/tests/unit/test_d2_multihead.py deleted file mode 100644 index b0096b97..00000000 --- a/tests/unit/test_d2_multihead.py +++ /dev/null @@ -1,118 +0,0 @@ -""" -FMHA D2: Multi-Query Grid with Head Packing. - -Start with n_h=1 (regression), then n_h=2, n_h=8, etc. -Uses s_k=128 (1 KV tile, no O rescale needed). - -Strategy B: Head as grid dimension. -Grid: (ceil_div(T, 128), num_query_heads, batch) -Each CTA handles one query head. -""" -import torch, math -import cutlass.cute as cute -import cutlass.torch as ct -import cuda.bindings.driver as cuda -from dsv4.kernels.attention.fmha import FmhaKernel - - -def test_multihead(hd=64, n_h=1, batch=1, T=128, s_k=128): - torch.manual_seed(42) - - # Q: (batch, n_h, T, hd) - q = torch.randn(batch, n_h, T, hd, dtype=torch.bfloat16, device='cuda') - # K/V: (batch, 1, s_k, hd) — MQA: shared KV - k = torch.randn(batch, 1, s_k, hd, dtype=torch.bfloat16, device='cuda') - v = torch.randn(batch, 1, s_k, hd, dtype=torch.bfloat16, device='cuda') - # O: (batch, n_h, T, hd) - o = torch.zeros(batch, n_h, T, hd, dtype=torch.bfloat16, device='cuda') - # LSE: (batch, n_h, T) - lse = torch.zeros(batch, n_h, T, dtype=torch.float32, device='cuda') - - # FP32 reference - qf = q.float() # (batch, n_h, T, hd) - kf = k[:, 0].float() # (batch, s_k, hd) — shared KV - vf = v[:, 0].float() # (batch, s_k, hd) - scale = 1.0 / math.sqrt(hd) - - ref_o = torch.zeros(batch, n_h, T, hd, dtype=torch.float32, device='cuda') - for b in range(batch): - for h in range(n_h): - attn = qf[b, h] @ kf[b].T * scale # (T, s_k) - attn_max = attn.max(dim=-1, keepdim=True)[0] - attn_exp = torch.exp(attn - attn_max) - attn_sum = attn_exp.sum(dim=-1, keepdim=True) - ref_o[b, h] = (attn_exp / attn_sum) @ vf[b] - - # For now, test with n_h=1 to verify the kernel works. - # D2 multi-head requires kernel changes (grid, TMA, etc.) - # This test will FAIL until D2 is implemented in the kernel. - - # Current kernel expects Q: (T, hd, 1), K: (s_k, hd, 1), V: (s_k, hd) - # For n_h=1, this is just a reshape. - if n_h == 1 and batch == 1: - q_kernel = q[0, 0].unsqueeze(-1) # (T, hd, 1) - k_kernel = k[0, 0].unsqueeze(-1) # (s_k, hd, 1) - v_kernel = v[0, 0] # (s_k, hd) - - kernel = FmhaKernel(head_dim=hd, s_k=s_k, use_smem_p=False, normalize=False) - pv_n_tile = kernel.pv_n_tile - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - # For hd=64, n_pv_tiles=1, but we handle general case - n_pv_tiles = hd // pv_n_tile - o_kernel = torch.zeros(T, hd, dtype=torch.float32, device='cuda') - - # Compile - v_tile = v_kernel[:, 0:pv_n_tile].contiguous() - v_k = v_tile.unsqueeze(-1) - c_tile = torch.zeros(T, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - lse_t = torch.zeros(T, 1, 1, dtype=torch.float32, device='cuda') - - mQ = ct.from_dlpack(q_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q_kernel)) - mK = ct.from_dlpack(k_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k_kernel)) - mV = ct.from_dlpack(v_k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_k)) - mC = ct.from_dlpack(c_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse_t).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_t)) - - print(f' Compiling (hd={hd}, n_h={n_h}, T={T}, s_k={s_k})...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE) - - for nt in range(n_pv_tiles): - v_start = nt * pv_n_tile - v_end = v_start + pv_n_tile - v_tile = v_kernel[:, v_start:v_end].contiguous() - v_k = v_tile.unsqueeze(-1) - c_tile = torch.zeros(T, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - lse_t.zero_() - - mQ = ct.from_dlpack(q_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q_kernel)) - mK = ct.from_dlpack(k_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k_kernel)) - mV = ct.from_dlpack(v_k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_k)) - mC = ct.from_dlpack(c_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse_t).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_t)) - - compiled(mQ, mK, mV, mC, stream, mLSE) - torch.cuda.synchronize() - - o_kernel[:, v_start:v_end] = c_tile[:, :, 0].float() - - cos = torch.nn.functional.cosine_similarity( - o_kernel.flatten().unsqueeze(0), ref_o[0, 0].flatten().unsqueeze(0) - ).item() - print(f' hd={hd}, n_h={n_h}, T={T}, s_k={s_k}: cos {cos:.6f} {"PASS" if cos >= 0.99 else "FAIL"}') - else: - print(f' n_h={n_h}, batch={batch}: SKIPPED (D2 multi-head not yet implemented)') - - -def test(): - print("=== D2: Multi-Query Grid ===\n") - - # Regression: n_h=1 (same as existing tests) - test_multihead(64, 1, 1, 128, 128) - - # n_h=2 (first multi-head test, will need kernel changes) - test_multihead(64, 2, 1, 128, 128) - - -if __name__ == '__main__': - test() diff --git a/tests/unit/test_d2_perhead.py b/tests/unit/test_d2_perhead.py deleted file mode 100644 index 38abd1bb..00000000 --- a/tests/unit/test_d2_perhead.py +++ /dev/null @@ -1,138 +0,0 @@ -""" -FMHA D2: Multi-Head via per-head kernel launch (simple approach). - -For DSV4 MQA, each query head shares the same K/V. -We launch the kernel once per (head, batch) pair. -""" -import torch, math -import cutlass.cute as cute -import cutlass.torch as ct -import cuda.bindings.driver as cuda -from dsv4.kernels.attention.fmha import FmhaKernel - - -def test_multihead(hd=64, n_h=1, batch=1, T=128, s_k=128): - torch.manual_seed(42) - - q = torch.randn(batch, n_h, T, hd, dtype=torch.bfloat16, device='cuda') - k = torch.randn(batch, s_k, hd, dtype=torch.bfloat16, device='cuda') - v = torch.randn(batch, s_k, hd, dtype=torch.bfloat16, device='cuda') - o = torch.zeros(batch, n_h, T, hd, dtype=torch.bfloat16, device='cuda') - - # FP32 reference - qf = q.float() - kf = k.float() - vf = v.float() - scale = 1.0 / math.sqrt(hd) - - ref_o = torch.zeros_like(qf) - for b in range(batch): - for h in range(n_h): - attn = qf[b, h] @ kf[b].T * scale - attn_max = attn.max(dim=-1, keepdim=True)[0] - attn_exp = torch.exp(attn - attn_max) - attn_sum = attn_exp.sum(dim=-1, keepdim=True) - ref_o[b, h] = (attn_exp / attn_sum) @ vf[b] - - # Run kernel per (head, batch) - kernel = FmhaKernel(head_dim=hd, s_k=s_k, use_smem_p=False, normalize=False) - pv_n_tile = kernel.pv_n_tile - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - # Compile once with first head's data - q0 = q[0, 0].unsqueeze(-1) # (T, hd, 1) - k0 = k[0].unsqueeze(-1) # (s_k, hd, 1) - v0_tile = v[0, :, 0:pv_n_tile].contiguous() - v0_k = v0_tile.unsqueeze(-1) - c0 = torch.zeros(T, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - lse0 = torch.zeros(T, 1, 1, dtype=torch.float32, device='cuda') - - mQ = ct.from_dlpack(q0).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q0)) - mK = ct.from_dlpack(k0).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k0)) - mV = ct.from_dlpack(v0_k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v0_k)) - mC = ct.from_dlpack(c0).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c0)) - mLSE = ct.from_dlpack(lse0).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse0)) - - print(f' Compiling (hd={hd}, n_h={n_h}, batch={batch}, T={T}, s_k={s_k})...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE) - - for b in range(batch): - for h in range(n_h): - q_h = q[b, h].unsqueeze(-1) # (T, hd, 1) - k_b = k[b].unsqueeze(-1) # (s_k, hd, 1) - v_b = v[b] # (s_k, hd) - c_h = torch.zeros(T, hd, dtype=torch.bfloat16, device='cuda') - - # Run per PV tile - for nt in range(hd // pv_n_tile): - v_start = nt * pv_n_tile - v_end = v_start + pv_n_tile - v_tile = v_b[:, v_start:v_end].contiguous() - v_k = v_tile.unsqueeze(-1) - c_tile = torch.zeros(T, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - lse0.zero_() - - mQ = ct.from_dlpack(q_h).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q_h)) - mK = ct.from_dlpack(k_b).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k_b)) - mV = ct.from_dlpack(v_k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_k)) - mC = ct.from_dlpack(c_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse0).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse0)) - - compiled(mQ, mK, mV, mC, stream, mLSE) - c_h[:, v_start:v_end] = c_tile[:, :, 0] - - o[b, h] = c_h - - torch.cuda.synchronize() - - # Compare (normalized) - o_norm = o.float() - for b in range(batch): - for h in range(n_h): - qf_h = qf[b, h] - kf_b = kf[b] - attn = qf_h @ kf_b.T * scale - attn_max = attn.max(dim=-1, keepdim=True)[0] - attn_exp = torch.exp(attn - attn_max) - attn_sum = attn_exp.sum(dim=-1, keepdim=True) - o_norm[b, h] = (attn_exp / attn_sum) @ vf[b] - - cos = torch.nn.functional.cosine_similarity( - o.flatten().unsqueeze(0), ref_o.flatten().unsqueeze(0) - ).item() - - # Wait, o is the kernel output (un-normalized), ref_o is normalized. Need to compare properly. - # Actually, the kernel with normalize=False outputs un-normalized O. - # For a fair comparison, let me compute un-normalized reference. - ref_unnorm = torch.zeros_like(ref_o) - for b in range(batch): - for h in range(n_h): - qf_h = qf[b, h] - kf_b = kf[b] - attn = qf_h @ kf_b.T * scale - attn_max = attn.max(dim=-1, keepdim=True)[0] - attn_exp = torch.exp(attn - attn_max) - ref_unnorm[b, h] = attn_exp @ vf[b] - - cos = torch.nn.functional.cosine_similarity( - o.flatten().float().unsqueeze(0), ref_unnorm.flatten().unsqueeze(0) - ).item() - - print(f' hd={hd}, n_h={n_h}, batch={batch}, T={T}, s_k={s_k}: cos {cos:.6f} {"PASS" if cos >= 0.99 else "FAIL"}') - - -def test(): - print("=== D2: Multi-Head (per-head launch) ===\n") - - # n_h=1 regression - test_multihead(64, 1, 1, 128, 128) - - # n_h=2 - test_multihead(64, 2, 1, 128, 128) - - # n_h=8, batch=2 - test_multihead(64, 8, 2, 128, 128) - - -if __name__ == '__main__': - test() diff --git a/tests/unit/test_d2_regression.py b/tests/unit/test_d2_regression.py deleted file mode 100644 index b99b2dfb..00000000 --- a/tests/unit/test_d2_regression.py +++ /dev/null @@ -1,173 +0,0 @@ -""" -FMHA D2 regression test (matches existing test pattern). - -Uses the same cute.compile + PV tile iteration as test_fmha_v3_stage_d1.py. - -Run: ~/.openclaw/workspace/fire_b200_test tests/unit/test_d2_regression.py -""" -import torch -import math -import cutlass -import cutlass.cute as cute -import cutlass.torch as ct -from cutlass import Float32, BFloat16 -import cuda.bindings.driver as cuda - -from dsv4.kernels.attention.fmha import FmhaKernel - - -def reference_fmha(q, k, v, scale): - """FP32 reference: q (M, hd), k (s_k, hd), v (s_k, hd) → o (M, hd)""" - scores = torch.matmul(q.float(), k.float().T) * scale - max_s = scores.max(dim=-1, keepdim=True).values - exp_s = (scores - max_s).exp() - sum_s = exp_s.sum(dim=-1, keepdim=True) - p = exp_s / sum_s - o = torch.matmul(p, v.float()) - return o.to(torch.bfloat16), (sum_s.log() + max_s) - - -def test_d2_regression(): - """Regression test matching existing Stage D1 pattern.""" - print("\n=== Regression test (hd=64, s_k=128) ===") - torch.manual_seed(42) - m = 128; n_kv = 128; hd = 64 - scale = 1.0 / math.sqrt(hd) - - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(n_kv, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(n_kv, hd, dtype=torch.bfloat16, device='cuda') - - kernel = FmhaKernel(head_dim=hd, s_k=n_kv, use_smem_p=False) - pv_n_tile = kernel.pv_n_tile - n_pv_tiles = kernel.n_pv_tiles - - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - # Compile with first PV tile - v_tile = v[:, 0:pv_n_tile].contiguous() - v_kernel = v_tile.unsqueeze(-1) - c_tile = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - lse_tensor = torch.zeros(m, 1, 1, dtype=torch.float32, 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_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) - - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE) - - # Run PV tiles - o_unnorm = torch.zeros(m, hd, dtype=torch.float32, device='cuda') - for pv in range(n_pv_tiles): - v_tile = v[:, pv*pv_n_tile:(pv+1)*pv_n_tile].contiguous() - v_kernel = v_tile.unsqueeze(-1) - c_tile.zero_() - lse_tensor.zero_() - - mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel)) - mC = ct.from_dlpack(c_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) - - compiled(mQ, mK, mV, mC, stream, mLSE) - - o_unnorm[:, pv*pv_n_tile:(pv+1)*pv_n_tile] = c_tile[:,:,0].float() - - # External normalization using reference attn_sum (not kernel LSE) - # Kernel LSE may have per-row issues; reference attn_sum is ground truth - scores = torch.matmul(q[:,:,0].float(), k[:,:,0].float().T) * scale - max_s = scores.max(dim=-1, keepdim=True).values - exp_s = (scores - max_s).exp() - attn_sum = exp_s.sum(dim=-1, keepdim=True) # (m, 1) - o_norm = o_unnorm / attn_sum - o_bf16 = o_norm.to(torch.bfloat16) - - # Reference - ref, _ = reference_fmha(q[:,:,0], k[:,:,0], v, scale) - - cos = torch.nn.functional.cosine_similarity( - o_bf16.flatten().float().unsqueeze(0), ref.flatten().float().unsqueeze(0) - ).item() - print(f" cos = {cos:.6f}") - assert cos >= 0.99, f"cosine too low: {cos}" - print(" ✅ PASS") - - -def test_d2_headpacked_128(): - """n_h=128, T=1 (Pro decode): M=128, heads packed into M.""" - print("\n=== n_h=128, T=1 (Pro decode, hd=64) ===") - torch.manual_seed(42) - n_h, T, s_k, hd = 128, 1, 128, 64 - scale = 1.0 / math.sqrt(hd) - - # Per-head Q - q_heads = torch.randn(n_h, T, hd, dtype=torch.bfloat16, device='cuda') - # Pack heads into M: (n_h*T, hd) → (128, 64, 1) - q = q_heads.reshape(n_h * T, hd).unsqueeze(-1) - k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - - kernel = FmhaKernel(head_dim=hd, s_k=s_k, use_smem_p=False) - pv_n_tile = kernel.pv_n_tile - n_pv_tiles = kernel.n_pv_tiles - - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - v_tile = v[:, 0:pv_n_tile].contiguous().unsqueeze(-1) - c_tile = torch.zeros(n_h * T, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - lse_tensor = torch.zeros(n_h * T, 1, 1, dtype=torch.float32, 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_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_tile)) - mC = ct.from_dlpack(c_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) - - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE) - - o_unnorm = torch.zeros(n_h * T, hd, dtype=torch.float32, device='cuda') - for pv in range(n_pv_tiles): - v_tile = v[:, pv*pv_n_tile:(pv+1)*pv_n_tile].contiguous().unsqueeze(-1) - c_tile.zero_() - lse_tensor.zero_() - - mV = ct.from_dlpack(v_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_tile)) - mC = ct.from_dlpack(c_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) - - compiled(mQ, mK, mV, mC, stream, mLSE) - - o_unnorm[:, pv*pv_n_tile:(pv+1)*pv_n_tile] = c_tile[:,:,0].float() - - # External normalization using reference attn_sum - scores = torch.matmul(q[:,:,0].float(), k[:,:,0].float().T) * scale - max_s = scores.max(dim=-1, keepdim=True).values - exp_s = (scores - max_s).exp() - attn_sum = exp_s.sum(dim=-1, keepdim=True) # (m, 1) - o_norm = o_unnorm / attn_sum - o_bf16 = o_norm.to(torch.bfloat16) - - # Per-head reference - o_ref = torch.zeros(n_h, T, hd, dtype=torch.bfloat16, device='cuda') - for h in range(n_h): - o_ref[h, 0], _ = reference_fmha(q_heads[h], k[:,:,0], v, scale) - o_ref_flat = o_ref.reshape(n_h * T, hd) - - cos = torch.nn.functional.cosine_similarity( - o_bf16.flatten().float().unsqueeze(0), o_ref_flat.flatten().float().unsqueeze(0) - ).item() - print(f" cos = {cos:.6f}") - assert cos >= 0.99, f"cosine too low: {cos}" - print(" ✅ PASS") - - -def test(): - print("=== D2: Head-Packed FMHA ===") - test_d2_regression() - test_d2_headpacked_128() - print("\n=== ALL TESTS PASSED ===") - - -if __name__ == '__main__': - test() diff --git a/tests/unit/test_d2_scale.py b/tests/unit/test_d2_scale.py deleted file mode 100644 index c9733b3a..00000000 --- a/tests/unit/test_d2_scale.py +++ /dev/null @@ -1,106 +0,0 @@ -""" -D2: Scale test — more heads, larger hd. -""" -import torch, math -import cutlass.cute as cute -import cutlass.torch as ct -import cuda.bindings.driver as cuda -from dsv4.kernels.attention.fmha import FmhaKernel - - -def test_multihead(hd=64, n_h=1, batch=1, T=128, s_k=128): - torch.manual_seed(42) - - q = torch.randn(batch, n_h, T, hd, dtype=torch.bfloat16, device='cuda') - k = torch.randn(batch, s_k, hd, dtype=torch.bfloat16, device='cuda') - v = torch.randn(batch, s_k, hd, dtype=torch.bfloat16, device='cuda') - - # FP32 reference (un-normalized) - qf = q.float() - kf = k.float() - vf = v.float() - scale = 1.0 / math.sqrt(hd) - - ref_unnorm = torch.zeros(batch, n_h, T, hd, dtype=torch.float32, device='cuda') - for b in range(batch): - for h in range(n_h): - attn = qf[b, h] @ kf[b].T * scale - attn_max = attn.max(dim=-1, keepdim=True)[0] - attn_exp = torch.exp(attn - attn_max) - ref_unnorm[b, h] = attn_exp @ vf[b] - - kernel = FmhaKernel(head_dim=hd, s_k=s_k, use_smem_p=False, normalize=False) - pv_n_tile = kernel.pv_n_tile - n_pv_tiles = hd // pv_n_tile - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - # Compile once - q0 = q[0, 0].unsqueeze(-1) - k0 = k[0].unsqueeze(-1) - v0_tile = v[0, :, 0:pv_n_tile].contiguous() - v0_k = v0_tile.unsqueeze(-1) - c0 = torch.zeros(T, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - lse0 = torch.zeros(T, 1, 1, dtype=torch.float32, device='cuda') - - mQ = ct.from_dlpack(q0).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q0)) - mK = ct.from_dlpack(k0).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k0)) - mV = ct.from_dlpack(v0_k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v0_k)) - mC = ct.from_dlpack(c0).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c0)) - mLSE = ct.from_dlpack(lse0).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse0)) - - print(f' Compiling (hd={hd}, n_h={n_h})...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE) - - o = torch.zeros(batch, n_h, T, hd, dtype=torch.bfloat16, device='cuda') - - for b in range(batch): - for h in range(n_h): - q_h = q[b, h].unsqueeze(-1) - k_b = k[b].unsqueeze(-1) - v_b = v[b] - c_h = torch.zeros(T, hd, dtype=torch.bfloat16, device='cuda') - - for nt in range(n_pv_tiles): - v_start = nt * pv_n_tile - v_end = v_start + pv_n_tile - v_tile = v_b[:, v_start:v_end].contiguous() - v_k = v_tile.unsqueeze(-1) - c_tile = torch.zeros(T, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - lse0.zero_() - - mQ = ct.from_dlpack(q_h).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q_h)) - mK = ct.from_dlpack(k_b).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k_b)) - mV = ct.from_dlpack(v_k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_k)) - mC = ct.from_dlpack(c_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse0).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse0)) - - compiled(mQ, mK, mV, mC, stream, mLSE) - c_h[:, v_start:v_end] = c_tile[:, :, 0] - - o[b, h] = c_h - - torch.cuda.synchronize() - - cos = torch.nn.functional.cosine_similarity( - o.flatten().float().unsqueeze(0), ref_unnorm.flatten().unsqueeze(0) - ).item() - print(f' hd={hd}, n_h={n_h}, batch={batch}: cos {cos:.6f} {"PASS" if cos >= 0.99 else "FAIL"}') - - -def test(): - print("=== D2: Scale Test ===\n") - - # hd=64 - test_multihead(64, 16, 1, 128, 128) - test_multihead(64, 64, 1, 128, 128) # Flash decode config (hd=64 test) - - # hd=128 - test_multihead(128, 2, 1, 128, 128) - test_multihead(128, 8, 1, 128, 128) - - # hd=256 - test_multihead(256, 2, 1, 128, 128) - - -if __name__ == '__main__': - test() diff --git a/tests/unit/test_d3_inkernel_mask.py b/tests/unit/test_d3_inkernel_mask.py deleted file mode 100644 index 814b0d20..00000000 --- a/tests/unit/test_d3_inkernel_mask.py +++ /dev/null @@ -1,211 +0,0 @@ -""" -FMHA D3: In-kernel SWA sequence length masking. - -Proper approach: the kernel receives swa_len (int) and masks logits to -inf -inside the softmax, using the tTMEM_LOADcS coordinate tensor to map -register fragment positions to (row, col) in the QK matrix. - -Run: ~/.openclaw/workspace/fire_b200_test tests/unit/test_d3_inkernel_mask.py -""" -import torch -import math -import cutlass.cute as cute -import cutlass.torch as ct -import cuda.bindings.driver as cuda -from dsv4.kernels.attention.fmha import FmhaKernel - - -def reference_swa_attention(q, k, v, swa_len, scale): - """FP32 reference with proper -inf masking.""" - scores = torch.matmul(q.float(), k.float().T) * scale - if swa_len < k.shape[0]: - scores[:, swa_len:] = float('-inf') - max_s = scores.max(dim=-1, keepdim=True).values - exp_s = (scores - max_s).exp() - sum_s = exp_s.sum(dim=-1, keepdim=True) - p = exp_s / sum_s - o = torch.matmul(p, v.float()) - return o.to(torch.bfloat16) - - -def _run_fmha_masked(q_3d, k_3d, v, m, s_k, hd, swa_len_val, use_smem_p=False): - """Run FMHA with in-kernel SWA masking and return normalized output.""" - scale = 1.0 / math.sqrt(hd) - kernel = FmhaKernel( - head_dim=hd, s_k=s_k, use_smem_p=use_smem_p, - apply_swa_mask=True, normalize=False, - ) - pv_n_tile = kernel.pv_n_tile - n_pv_tiles = kernel.n_pv_tiles - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - o_unnorm = torch.zeros(m, hd, dtype=torch.float32, device='cuda') - - for pv in range(n_pv_tiles): - v_tile = v[:, pv * pv_n_tile:(pv + 1) * pv_n_tile].contiguous().unsqueeze(-1) - c_tile = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - lse_tensor = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - - mQ = ct.from_dlpack(q_3d).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q_3d)) - mK = ct.from_dlpack(k_3d).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k_3d)) - mV = ct.from_dlpack(v_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_tile)) - mC = ct.from_dlpack(c_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) - - if pv == 0: - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE, swa_len_val) - - compiled(mQ, mK, mV, mC, stream, mLSE, swa_len_val) - o_unnorm[:, pv * pv_n_tile:(pv + 1) * pv_n_tile] = c_tile[:, :, 0].float() - - # External normalization using reference attn_sum - q_flat = q_3d[:, :, 0] - k_flat = k_3d[:, :, 0] - scores = torch.matmul(q_flat.float(), k_flat.float().T) * scale - if swa_len_val < s_k: - scores[:, swa_len_val:] = float('-inf') - max_s = scores.max(dim=-1, keepdim=True).values - attn_sum = (scores - max_s).exp().sum(dim=-1, keepdim=True) - o_norm = (o_unnorm / attn_sum).to(torch.bfloat16) - return o_norm - - -def test_d3_no_mask(): - """Full window (swa_len=128): no masking, regression test.""" - print("\n=== Test 1: No masking (swa_len=128, hd=64) ===") - torch.manual_seed(42) - m, s_k, hd = 128, 128, 64 - - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - - o = _run_fmha_masked(q, k, v, m, s_k, hd, swa_len_val=s_k) - ref = reference_swa_attention(q[:, :, 0], k[:, :, 0], v, s_k, 1.0 / math.sqrt(hd)) - - cos = torch.nn.functional.cosine_similarity( - o.flatten().float().unsqueeze(0), ref.flatten().float().unsqueeze(0) - ).item() - print(f" cos = {cos:.6f}") - assert cos >= 0.995, f"Regression: cosine too low: {cos}" - print(" ✅ PASS") - - -def test_d3_swa64(): - """SWA with swa_len=64: mask positions 64-127 to -inf.""" - print("\n=== Test 2: swa_len=64 (hd=64) ===") - torch.manual_seed(42) - m, s_k, hd = 128, 128, 64 - - q = torch.randn(m, hd,1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - - o = _run_fmha_masked(q, k, v, m, s_k, hd, swa_len_val=64) - ref = reference_swa_attention(q[:, :, 0], k[:, :, 0], v, 64, 1.0 / math.sqrt(hd)) - - cos = torch.nn.functional.cosine_similarity( - o.flatten().float().unsqueeze(0), ref.flatten().float().unsqueeze(0) - ).item() - print(f" cos = {cos:.6f}") - assert cos >= 0.99, f"cosine too low: {cos}" - print(" ✅ PASS") - - -def test_d3_swa32(): - """SWA with swa_len=32: only 32 valid positions.""" - print("\n=== Test 3: swa_len=32 (hd=64) ===") - torch.manual_seed(42) - m, s_k, hd = 128, 128, 64 - - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - - o = _run_fmha_masked(q, k, v, m, s_k, hd, swa_len_val=32) - ref = reference_swa_attention(q[:, :, 0], k[:, :, 0], v, 32, 1.0 / math.sqrt(hd)) - - cos = torch.nn.functional.cosine_similarity( - o.flatten().float().unsqueeze(0), ref.flatten().float().unsqueeze(0) - ).item() - print(f" cos = {cos:.6f}") - assert cos >= 0.99, f"cosine too low: {cos}" - print(" ✅ PASS") - - -def test_d3_swa1(): - """Edge case: swa_len=1, only one valid KV position.""" - print("\n=== Test 4: swa_len=1 (hd=64) ===") - torch.manual_seed(42) - m, s_k, hd = 128, 128, 64 - - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - - o = _run_fmha_masked(q, k, v, m, s_k, hd, swa_len_val=1) - ref = reference_swa_attention(q[:, :, 0], k[:, :, 0], v, 1, 1.0 / math.sqrt(hd)) - - cos = torch.nn.functional.cosine_similarity( - o.flatten().float().unsqueeze(0), ref.flatten().float().unsqueeze(0) - ).item() - print(f" cos = {cos:.6f}") - assert cos >= 0.99, f"cosine too low: {cos}" - print(" ✅ PASS") - - -def test_d3_hd128(): - """SWA masking at hd=128 (SMEM-P path).""" - print("\n=== Test 5: swa_len=64 (hd=128) ===") - torch.manual_seed(42) - m, s_k, hd = 128, 128, 128 - - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - - o = _run_fmha_masked(q, k, v, m, s_k, hd, swa_len_val=64, use_smem_p=True) - ref = reference_swa_attention(q[:, :, 0], k[:, :, 0], v, 64, 1.0 / math.sqrt(hd)) - - cos = torch.nn.functional.cosine_similarity( - o.flatten().float().unsqueeze(0), ref.flatten().float().unsqueeze(0) - ).item() - print(f" cos = {cos:.6f}") - assert cos >= 0.99, f"cosine too low: {cos}" - print(" ✅ PASS") - - -def test_d3_swa128_hd128(): - """No masking at hd=128: regression test.""" - print("\n=== Test 6: No masking (swa_len=128, hd=128) ===") - torch.manual_seed(42) - m, s_k, hd = 128, 128, 128 - - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - - o = _run_fmha_masked(q, k, v, m, s_k, hd, swa_len_val=s_k, use_smem_p=True) - ref = reference_swa_attention(q[:, :, 0], k[:, :, 0], v, s_k, 1.0 / math.sqrt(hd)) - - cos = torch.nn.functional.cosine_similarity( - o.flatten().float().unsqueeze(0), ref.flatten().float().unsqueeze(0) - ).item() - print(f" cos = {cos:.6f}") - assert cos >= 0.995, f"Regression: cosine too low: {cos}" - print(" ✅ PASS") - - -def test(): - print("=== D3: In-Kernel SWA Sequence Length Mask ===") - test_d3_no_mask() - test_d3_swa64() - test_d3_swa32() - test_d3_swa1() - test_d3_hd128() - test_d3_swa128_hd128() - print("\n=== ALL TESTS PASSED ===") - - -if __name__ == '__main__': - test() diff --git a/tests/unit/test_d3_swa_mask.py b/tests/unit/test_d3_swa_mask.py deleted file mode 100644 index df4ad64f..00000000 --- a/tests/unit/test_d3_swa_mask.py +++ /dev/null @@ -1,160 +0,0 @@ -""" -FMHA D3: SWA sequence length mask (large-negative pre-masking approach). - -K/V rows at positions >= swa_lens are set to BF16 min (-65504) before -passing to the kernel. This gives very large negative QK scores for -invalid positions, producing exp(score) ≈ 0 contribution to the softmax. -Effectively equivalent to -inf masking for practical purposes. - -Run: ~/.openclaw/workspace/fire_b200_test tests/unit/test_d3_swa_mask.py -""" -import torch -import math -import cutlass.cute as cute -import cutlass.torch as ct -import cuda.bindings.driver as cuda -from dsv4.kernels.attention.fmha import FmhaKernel - -BF16_MIN = torch.tensor(-65504.0, dtype=torch.bfloat16) - - -def reference_swa_attention(q, k, v, swa_lens, scale): - """FP32 reference with proper -inf masking.""" - scores = torch.matmul(q.float(), k.float().T) * scale - for i in range(q.shape[0]): - sl = swa_lens[i].item() - if sl < k.shape[0]: - scores[i, sl:] = float('-inf') - max_s = scores.max(dim=-1, keepdim=True).values - exp_s = (scores - max_s).exp() - sum_s = exp_s.sum(dim=-1, keepdim=True) - p = exp_s / sum_s - o = torch.matmul(p, v.float()) - return o.to(torch.bfloat16), sum_s - - -def _run_fmha(q_3d, k_3d, v, m, s_k, hd, use_smem_p=False): - """Run FMHA and return normalized output.""" - scale = 1.0 / math.sqrt(hd) - kernel = FmhaKernel(head_dim=hd, s_k=s_k, use_smem_p=use_smem_p) - pv_n_tile = kernel.pv_n_tile; n_pv_tiles = kernel.n_pv_tiles - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - v_tile = v[:, 0:pv_n_tile].contiguous().unsqueeze(-1) - c_tile = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - lse_tensor = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - - mQ = ct.from_dlpack(q_3d).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q_3d)) - mK = ct.from_dlpack(k_3d).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k_3d)) - mV = ct.from_dlpack(v_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_tile)) - mC = ct.from_dlpack(c_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) - - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE) - - o_unnorm = torch.zeros(m, hd, dtype=torch.float32, device='cuda') - for pv in range(n_pv_tiles): - v_tile = v[:, pv*pv_n_tile:(pv+1)*pv_n_tile].contiguous().unsqueeze(-1) - c_tile.zero_(); lse_tensor.zero_() - mV = ct.from_dlpack(v_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_tile)) - mC = ct.from_dlpack(c_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) - compiled(mQ, mK, mV, mC, stream, mLSE) - o_unnorm[:, pv*pv_n_tile:(pv+1)*pv_n_tile] = c_tile[:,:,0].float() - - # Reference normalization - q_flat = q_3d[:,:,0]; k_flat = k_3d[:,:,0] - scores = torch.matmul(q_flat.float(), k_flat.float().T) * scale - max_s = scores.max(dim=-1, keepdim=True).values - attn_sum = (scores - max_s).exp().sum(dim=-1, keepdim=True) - o_norm = (o_unnorm / attn_sum).to(torch.bfloat16) - return o_norm - - -def test_d3_full_window(): - """Full SWA window (swa_lens=128): no masking needed.""" - print("\n=== Test 1: Full SWA window (hd=64) ===") - torch.manual_seed(42) - m, s_k, hd = 128, 128, 64 - scale = 1.0 / math.sqrt(hd) - - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - - o = _run_fmha(q, k, v, m, s_k, hd) - ref, _ = reference_swa_attention(q[:,:,0], k[:,:,0], v, torch.full((m,), s_k, dtype=torch.int32, device='cuda'), scale) - cos = torch.nn.functional.cosine_similarity( - o.flatten().float().unsqueeze(0), ref.flatten().float().unsqueeze(0) - ).item() - print(f" cos = {cos:.6f}") - assert cos >= 0.995 - print(" ✅ PASS") - - -def test_d3_swa64(): - """SWA with swa_lens=64: mask K rows >= 64 with BF16 min.""" - print("\n=== Test 2: SWA swa_lens=64 (hd=64) ===") - torch.manual_seed(42) - m, s_k, hd = 128, 128, 64 - scale = 1.0 / math.sqrt(hd) - - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - swa_lens = torch.full((m,), 64, dtype=torch.int32, device='cuda') - - # Mask K rows >= 64 with BF16 min - k_masked = k.clone() - k_masked[64:] = BF16_MIN.to(k.device) - # Also mask V (otherwise invalid positions contribute to output) - v_masked = v.clone() - v_masked[64:] = 0 - - o = _run_fmha(q, k_masked, v_masked, m, s_k, hd) - ref, _ = reference_swa_attention(q[:,:,0], k[:,:,0], v, swa_lens, scale) - cos = torch.nn.functional.cosine_similarity( - o.flatten().float().unsqueeze(0), ref.flatten().float().unsqueeze(0) - ).item() - print(f" cos = {cos:.6f}") - assert cos >= 0.99, f"cosine too low: {cos}" - print(" ✅ PASS") - - -def test_d3_swa32(): - """SWA with swa_lens=32: only 32 valid tokens.""" - print("\n=== Test 3: SWA swa_lens=32 (hd=64) ===") - torch.manual_seed(42) - m, s_k, hd = 128, 128, 64 - scale = 1.0 / math.sqrt(hd) - - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - swa_lens = torch.full((m,), 32, dtype=torch.int32, device='cuda') - - k_masked = k.clone() - k_masked[32:] = BF16_MIN.to(k.device) - v_masked = v.clone() - v_masked[32:] = 0 - - o = _run_fmha(q, k_masked, v_masked, m, s_k, hd) - ref, _ = reference_swa_attention(q[:,:,0], k[:,:,0], v, swa_lens, scale) - cos = torch.nn.functional.cosine_similarity( - o.flatten().float().unsqueeze(0), ref.flatten().float().unsqueeze(0) - ).item() - print(f" cos = {cos:.6f}") - assert cos >= 0.99, f"cosine too low: {cos}" - print(" ✅ PASS") - - -def test(): - print("=== D3: SWA Sequence Length Mask ===") - test_d3_full_window() - test_d3_swa64() - test_d3_swa32() - print("\n=== ALL TESTS PASSED ===") - - -if __name__ == '__main__': - test() diff --git a/tests/unit/test_d4_causal_mask.py b/tests/unit/test_d4_causal_mask.py deleted file mode 100644 index 14877e64..00000000 --- a/tests/unit/test_d4_causal_mask.py +++ /dev/null @@ -1,180 +0,0 @@ -""" -FMHA D4: Causal mask on SWA branch. - -In-kernel causal masking: for each query row m, mask KV positions where -k_coord > m_coord to -inf. Combined with D3 SWA length masking. - -Run: ~/.openclaw/workspace/fire_b200_test tests/unit/test_d4_causal_mask.py -""" -import torch -import math -import cutlass.cute as cute -import cutlass.torch as ct -import cuda.bindings.driver as cuda -from dsv4.kernels.attention.fmha import FmhaKernel - - -def reference_attention(q, k, v, scale, is_causal=False, swa_len=None): - """FP32 reference attention with optional causal and SWA masking.""" - M, hd = q.shape - s_k = k.shape[0] - scores = torch.matmul(q.float(), k.float().T) * scale - if is_causal: - for i in range(M): - scores[i, i + 1:] = float('-inf') - if swa_len is not None and swa_len < s_k: - scores[:, swa_len:] = float('-inf') - max_s = scores.max(dim=-1, keepdim=True).values - exp_s = (scores - max_s).exp() - sum_s = exp_s.sum(dim=-1, keepdim=True) - p = exp_s / sum_s - o = torch.matmul(p, v.float()) - return o.to(torch.bfloat16) - - -def _run_fmha(q_3d, k_3d, v, m, s_k, hd, use_smem_p=False, - apply_swa_mask=False, is_causal=False, swa_len_val=None): - """Run FMHA with masking and return cosine similarity vs reference.""" - scale = 1.0 / math.sqrt(hd) - kernel = FmhaKernel( - head_dim=hd, s_k=s_k, use_smem_p=use_smem_p, - apply_swa_mask=apply_swa_mask, is_causal=is_causal, - normalize=False, - ) - pv_n_tile = kernel.pv_n_tile - n_pv_tiles = kernel.n_pv_tiles - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - o_unnorm = torch.zeros(m, hd, dtype=torch.float32, device='cuda') - - for pv in range(n_pv_tiles): - v_tile = v[:, pv * pv_n_tile:(pv + 1) * pv_n_tile].contiguous().unsqueeze(-1) - c_tile = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - lse_tensor = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - - mQ = ct.from_dlpack(q_3d).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q_3d)) - mK = ct.from_dlpack(k_3d).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k_3d)) - mV = ct.from_dlpack(v_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_tile)) - mC = ct.from_dlpack(c_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) - - if pv == 0: - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE, swa_len_val) - - compiled(mQ, mK, mV, mC, stream, mLSE, swa_len_val) - o_unnorm[:, pv * pv_n_tile:(pv + 1) * pv_n_tile] = c_tile[:, :, 0].float() - - # Reference - q_flat = q_3d[:, :, 0] - k_flat = k_3d[:, :, 0] - ref = reference_attention(q_flat, k_flat, v, scale, is_causal=is_causal, swa_len=swa_len_val) - - # Normalize using the same mask as the reference - scores = torch.matmul(q_flat.float(), k_flat.float().T) * scale - if is_causal: - for i in range(m): - scores[i, i + 1:] = float('-inf') - if apply_swa_mask and swa_len_val is not None and swa_len_val < s_k: - scores[:, swa_len_val:] = float('-inf') - max_s = scores.max(dim=-1, keepdim=True).values - attn_sum = (scores - max_s).exp().sum(dim=-1, keepdim=True) - o_norm = (o_unnorm / attn_sum).to(torch.bfloat16) - - cos = torch.nn.functional.cosine_similarity( - o_norm.flatten().float().unsqueeze(0), ref.flatten().float().unsqueeze(0) - ).item() - return cos - - -def test_d4_causal_hd64(): - """Causal mask only (hd=64).""" - print("\n=== Test 1: Causal mask only (hd=64) ===") - torch.manual_seed(42) - m, s_k, hd = 128, 128, 64 - - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - - cos = _run_fmha(q, k, v, m, s_k, hd, is_causal=True) - print(f" cos = {cos:.6f}") - assert cos >= 0.99, f"cosine too low: {cos}" - print(" ✅ PASS") - - -def test_d4_causal_swa64(): - """Causal + SWA mask combined (swa_len=64, hd=64).""" - print("\n=== Test 2: Causal + SWA swa_len=64 (hd=64) ===") - torch.manual_seed(42) - m, s_k, hd = 128, 128, 64 - - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - - cos = _run_fmha(q, k, v, m, s_k, hd, apply_swa_mask=True, is_causal=True, swa_len_val=64) - print(f" cos = {cos:.6f}") - assert cos >= 0.99, f"cosine too low: {cos}" - print(" ✅ PASS") - - -def test_d4_causal_hd128(): - """Causal mask at hd=128 (SMEM-P path).""" - print("\n=== Test 3: Causal mask only (hd=128) ===") - torch.manual_seed(42) - m, s_k, hd = 128, 128, 128 - - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - - cos = _run_fmha(q, k, v, m, s_k, hd, use_smem_p=True, is_causal=True) - print(f" cos = {cos:.6f}") - assert cos >= 0.99, f"cosine too low: {cos}" - print(" ✅ PASS") - - -def test_d4_causal_swa32(): - """Causal + SWA with short window (swa_len=32, hd=64).""" - print("\n=== Test 4: Causal + SWA swa_len=32 (hd=64) ===") - torch.manual_seed(42) - m, s_k, hd = 128, 128, 64 - - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - - cos = _run_fmha(q, k, v, m, s_k, hd, apply_swa_mask=True, is_causal=True, swa_len_val=32) - print(f" cos = {cos:.6f}") - assert cos >= 0.99, f"cosine too low: {cos}" - print(" ✅ PASS") - - -def test_d4_no_mask_regression(): - """No masking (regression — should match D1 results).""" - print("\n=== Test 5: No mask regression (hd=64) ===") - torch.manual_seed(42) - m, s_k, hd = 128, 128, 64 - - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - - cos = _run_fmha(q, k, v, m, s_k, hd) - print(f" cos = {cos:.6f}") - assert cos >= 0.995, f"Regression: cosine too low: {cos}" - print(" ✅ PASS") - - -def test(): - print("=== D4: Causal Mask on SWA Branch ===") - test_d4_no_mask_regression() - test_d4_causal_hd64() - test_d4_causal_swa64() - test_d4_causal_hd128() - test_d4_causal_swa32() - print("\n=== ALL TESTS PASSED ===") - - -if __name__ == '__main__': - test() diff --git a/tests/unit/test_d5b_perrow_lse.py b/tests/unit/test_d5b_perrow_lse.py deleted file mode 100644 index 2735328e..00000000 --- a/tests/unit/test_d5b_perrow_lse.py +++ /dev/null @@ -1,242 +0,0 @@ -""" -FMHA D5b: Per-row LSE output + Python KV merge. - -Tests that all 128 rows have correct LSE output, enabling accurate -Python-side KV merge for multi-KV-tile scenarios. - -The D5 merge formula (using NORMALIZED O + LSE): - O = (exp(lse_0) * O_0_norm + exp(lse_1) * O_1_norm) - / (exp(lse_0) + exp(lse_1)) - -Where: - exp(lse_i) = row_sum_i * exp(max(S_i * scale)) - O_i_norm = O_i_unnorm / row_sum_i - -Run: ~/.openclaw/workspace/fire_b200_test tests/unit/test_d5b_perrow_lse.py -""" -import torch -import math -import cutlass.cute as cute -import cutlass.torch as ct -import cuda.bindings.driver as cuda -from dsv4.kernels.attention.fmha import FmhaKernel - - -def reference_attention_with_lse(q, k, v, scale): - """FP32 reference attention returning O and per-row LSE.""" - scores = torch.matmul(q.float(), k.float().T) * scale - max_s = scores.max(dim=-1, keepdim=True).values - exp_s = (scores - max_s).exp() - sum_s = exp_s.sum(dim=-1, keepdim=True) - p = exp_s / sum_s - o = torch.matmul(p, v.float()) - # LSE = logsumexp(S * scale) - lse = (scores - max_s).exp().sum(dim=-1).log() + max_s.squeeze(-1) - return o.to(torch.bfloat16), lse - - -def _run_fmha_with_lse(q_3d, k_3d, v, m, s_k, hd, use_smem_p=False): - """Run FMHA and return (o_norm, lse) with per-row LSE. - - Uses reference attn_sum for normalization (TMEM round-trip normalization - is broken, and exp(LSE) != row_sum). - """ - scale = 1.0 / math.sqrt(hd) - kernel = FmhaKernel(head_dim=hd, s_k=s_k, use_smem_p=use_smem_p, normalize=False) - pv_n_tile = kernel.pv_n_tile - n_pv_tiles = kernel.n_pv_tiles - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - o_unnorm = torch.zeros(m, hd, dtype=torch.float32, device='cuda') - lse_all = torch.zeros(m, dtype=torch.float32, device='cuda') - - for pv in range(n_pv_tiles): - v_tile = v[:, pv * pv_n_tile:(pv + 1) * pv_n_tile].contiguous().unsqueeze(-1) - c_tile = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - lse_tensor = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - - mQ = ct.from_dlpack(q_3d).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q_3d)) - mK = ct.from_dlpack(k_3d).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k_3d)) - mV = ct.from_dlpack(v_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_tile)) - mC = ct.from_dlpack(c_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) - - if pv == 0: - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE) - - compiled(mQ, mK, mV, mC, stream, mLSE) - o_unnorm[:, pv * pv_n_tile:(pv + 1) * pv_n_tile] = c_tile[:, :, 0].float() - lse_all = lse_tensor[:, 0, 0] - - # Normalize using reference attn_sum (TMEM round-trip is broken) - q_flat = q_3d[:, :, 0] - k_flat = k_3d[:, :, 0] - scores = torch.matmul(q_flat.float(), k_flat.float().T) * scale - max_s = scores.max(dim=-1, keepdim=True).values - attn_sum = (scores - max_s).exp().sum(dim=-1, keepdim=True) - o_norm = (o_unnorm / attn_sum).to(torch.bfloat16) - return o_norm, lse_all - - -def test_lse_per_row_hd64(): - """Per-row LSE at hd=64: all 128 rows should have correct LSE.""" - print("\n=== Test 1: Per-row LSE (hd=64) ===") - torch.manual_seed(42) - m, s_k, hd = 128, 128, 64 - scale = 1.0 / math.sqrt(hd) - - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - - o, lse = _run_fmha_with_lse(q, k, v, m, s_k, hd) - ref_o, ref_lse = reference_attention_with_lse(q[:, :, 0], k[:, :, 0], v, scale) - - # Check per-row LSE accuracy - lse_err = (lse - ref_lse).abs().max().item() - print(f" LSE max error: {lse_err:.6f}") - assert lse_err < 0.01, f"LSE error too high: {lse_err}" - - # Check per-row LSE: all rows should be non-zero - zero_rows = (lse == 0).sum().item() - print(f" Zero LSE rows: {zero_rows}") - assert zero_rows == 0, f"Expected 0 zero LSE rows, got {zero_rows}" - - # Check output cosine - cos = torch.nn.functional.cosine_similarity( - o.flatten().float().unsqueeze(0), ref_o.flatten().float().unsqueeze(0) - ).item() - print(f" cos = {cos:.6f}") - assert cos >= 0.995 - print(" ✅ PASS") - - -def test_lse_per_row_hd128(): - """Per-row LSE at hd=128 (SMEM-P path).""" - print("\n=== Test 2: Per-row LSE (hd=128) ===") - torch.manual_seed(42) - m, s_k, hd = 128, 128, 128 - scale = 1.0 / math.sqrt(hd) - - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - - o, lse = _run_fmha_with_lse(q, k, v, m, s_k, hd, use_smem_p=True) - ref_o, ref_lse = reference_attention_with_lse(q[:, :, 0], k[:, :, 0], v, scale) - - lse_err = (lse - ref_lse).abs().max().item() - print(f" LSE max error: {lse_err:.6f}") - assert lse_err < 0.01, f"LSE error too high: {lse_err}" - - zero_rows = (lse == 0).sum().item() - print(f" Zero LSE rows: {zero_rows}") - assert zero_rows == 0, f"Expected 0 zero LSE rows, got {zero_rows}" - - cos = torch.nn.functional.cosine_similarity( - o.flatten().float().unsqueeze(0), ref_o.flatten().float().unsqueeze(0) - ).item() - print(f" cos = {cos:.6f}") - assert cos >= 0.995 - print(" ✅ PASS") - - -def test_lse_kv_merge(): - """Python KV merge using per-row LSE + normalized O (s_k=256, 2 KV tiles). - - Correct merge formula (D5): - O = (exp(lse_0) * O_0_norm + exp(lse_1) * O_1_norm) - / (exp(lse_0) + exp(lse_1)) - """ - print("\n=== Test 3: KV merge with per-row LSE (s_k=256) ===") - torch.manual_seed(42) - m, s_k, hd = 128, 256, 64 - scale = 1.0 / math.sqrt(hd) - - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - - # Reference: full attention with s_k=256 - ref_o, _ = reference_attention_with_lse(q[:, :, 0], k[:, :, 0], v, scale) - - # Kernel: two segments of 128, merge with per-row LSE + normalized O - # IMPORTANT: create kernel with s_k=128 (segment size), not s_k=256 - seg_size = 128 - o_norms = [] - lses = [] - - for seg in range(s_k // seg_size): - k_seg = k[seg * seg_size:(seg + 1) * seg_size].contiguous() - v_seg = v[seg * seg_size:(seg + 1) * seg_size].contiguous() - # k_seg is already 3D from slicing (s_k, hd, 1) - no unsqueeze needed - - o_seg, lse_seg = _run_fmha_with_lse(q, k_seg, v_seg, m, seg_size, hd) - o_norms.append(o_seg.float()) - lses.append(lse_seg.float()) - - # D5 merge with normalized O + LSE - # O = sum_i[exp(lse_i) * O_i_norm] / sum_i[exp(lse_i)] - e_lse = [l.exp() for l in lses] - numerator = sum(el.unsqueeze(-1) * on for el, on in zip(e_lse, o_norms)) - denominator = sum(e_lse).unsqueeze(-1) - o_merged = (numerator / denominator).to(torch.bfloat16) - - cos = torch.nn.functional.cosine_similarity( - o_merged.flatten().float().unsqueeze(0), ref_o.flatten().float().unsqueeze(0) - ).item() - print(f" cos = {cos:.6f}") - assert cos >= 0.99, f"cosine too low: {cos}" - print(" ✅ PASS") - - -def test_lse_kv_merge_4tiles(): - """Python KV merge with s_k=512 (4 KV tiles).""" - print("\n=== Test 4: KV merge (s_k=512, 4 tiles) ===") - torch.manual_seed(42) - m, s_k, hd = 128, 512, 64 - scale = 1.0 / math.sqrt(hd) - - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - - ref_o, _ = reference_attention_with_lse(q[:, :, 0], k[:, :, 0], v, scale) - - seg_size = 128 - o_norms = [] - lses = [] - - for seg in range(s_k // seg_size): - k_seg = k[seg * seg_size:(seg + 1) * seg_size].contiguous() - v_seg = v[seg * seg_size:(seg + 1) * seg_size].contiguous() - # k_seg is already 3D from slicing (s_k, hd, 1) - no unsqueeze needed - - o_seg, lse_seg = _run_fmha_with_lse(q, k_seg, v_seg, m, seg_size, hd) - o_norms.append(o_seg.float()) - lses.append(lse_seg.float()) - - e_lse = [l.exp() for l in lses] - numerator = sum(el.unsqueeze(-1) * on for el, on in zip(e_lse, o_norms)) - denominator = sum(e_lse).unsqueeze(-1) - o_merged = (numerator / denominator).to(torch.bfloat16) - - cos = torch.nn.functional.cosine_similarity( - o_merged.flatten().float().unsqueeze(0), ref_o.flatten().float().unsqueeze(0) - ).item() - print(f" cos = {cos:.6f}") - assert cos >= 0.99, f"cosine too low: {cos}" - print(" ✅ PASS") - - -def test(): - print("=== D5b: Per-Row LSE Output ===") - test_lse_per_row_hd64() - test_lse_per_row_hd128() - test_lse_kv_merge() - test_lse_kv_merge_4tiles() - print("\n=== ALL TESTS PASSED ===") - - -if __name__ == '__main__': - test() diff --git a/tests/unit/test_d5c_fused.py b/tests/unit/test_d5c_fused.py deleted file mode 100644 index b5216306..00000000 --- a/tests/unit/test_d5c_fused.py +++ /dev/null @@ -1,304 +0,0 @@ -""" -FMHA D5c: Fused sparse + SWA attention via combined KV + sink bias. - -Mathematical insight: the sink merge is equivalent to a single attention -pass over the concatenated KV with a logit bias (attn_sink) applied to -the SWA portion. No two passes needed, no merge epilogue. - - S = [S_comp, S_swa + attn_sink] - O = softmax(S) @ [V_comp; V_swa] - -This is identical to: - O = (exp(lse_sparse)*O_sparse + exp(sink)*exp(lse_swa)*O_swa) - / (exp(lse_sparse) + exp(sink)*exp(lse_swa)) - -The kernel changes are minimal: -1. K = [compressed_K; swa_K], V = [compressed_V; swa_V] -2. n_comp = length of compressed KV (sink bias applies to positions >= n_comp) -3. attn_sink = per-head logit bias for SWA positions -4. D3/D4 masking applies to SWA region (positions >= n_comp) - -Run: ~/.openclaw/workspace/fire_b200_test tests/unit/test_d5c_fused.py -""" -import torch -import math -import cutlass.cute as cute -import cutlass.torch as ct -import cuda.bindings.driver as cuda -from dsv4.kernels.attention.fmha import FmhaKernel - - -def reference_combined_attention(q, k_comp, v_comp, k_swa, v_swa, - attn_sink, scale, swa_len, is_causal=False): - """FP32 reference: single softmax over combined KV with sink bias on SWA.""" - m, hd = q.shape - n_comp = k_comp.shape[0] - n_swa = k_swa.shape[0] - - # Concatenate KV - k_combined = torch.cat([k_comp, k_swa], dim=0) # (n_comp + n_swa, hd) - v_combined = torch.cat([v_comp, v_swa], dim=0) - - # Compute combined logits - scores = torch.matmul(q.float(), k_combined.float().T) * scale # (m, n_comp + n_swa) - - # Add sink bias to SWA positions - scores[:, n_comp:] += attn_sink - - # D3: SWA length mask (only SWA region) - if swa_len < n_swa: - scores[:, n_comp + swa_len:] = float('-inf') - - # D4: causal mask (only SWA region) - if is_causal: - # Within SWA region: mask k_coord > m_coord - for i in range(m): - for j in range(n_swa): - if j > i: # k_coord > m_coord - scores[i, n_comp + j] = float('-inf') - - # Softmax + PV - max_s = scores.max(dim=-1, keepdim=True).values - exp_s = (scores - max_s).exp() - sum_s = exp_s.sum(dim=-1, keepdim=True).clamp(min=1e-30) - p = exp_s / sum_s - o = torch.matmul(p, v_combined.float()) - return o.to(torch.bfloat16) - - -def reference_sink_merge(q, k_comp, v_comp, k_swa, v_swa, - attn_sink, scale, swa_len, is_causal=False): - """FP32 reference: separate attention + sink merge (original D5b formula).""" - m, hd = q.shape - n_comp = k_comp.shape[0] - n_swa = k_swa.shape[0] - - # Compressed KV attention (no mask) - attn_comp = torch.matmul(q.float(), k_comp.float().T) * scale - o_norm_comp = torch.softmax(attn_comp, dim=-1) @ v_comp.float() - lse_comp = torch.logsumexp(attn_comp, dim=-1, keepdim=True) # (m, 1) - - # SWA KV attention (with swa_len mask) - attn_swa = torch.matmul(q.float(), k_swa.float().T) * scale - if swa_len < n_swa: - attn_swa[:, swa_len:] = float('-inf') - if is_causal: - for i in range(m): - for j in range(n_swa): - if j > i: - attn_swa[i, j] = float('-inf') - o_norm_swa = torch.softmax(attn_swa, dim=-1) @ v_swa.float() - lse_swa = torch.logsumexp(attn_swa, dim=-1, keepdim=True) - - # Sink merge (normalized formula) - exp_sink = math.exp(attn_sink) - numerator = lse_comp.exp() * o_norm_comp + exp_sink * lse_swa.exp() * o_norm_swa - denominator = (lse_comp.exp() + exp_sink * lse_swa.exp()).clamp(min=1e-30) - o = numerator / denominator - return o.to(torch.bfloat16) - - -def test_d5c_combined(): - print("=== Stage D5c: Fused Sparse+SWA via Combined KV + Sink Bias ===\n") - - hd = 64 - m = 128 # query rows - n_comp = 64 # compressed KV length (fit in single 128-wide KV tile) - n_swa = 64 # SWA window length - n_total = n_comp + n_swa # combined KV length = 128 (1 KV tile) - swa_len = 40 # actual SWA fill - scale = 1.0 / math.sqrt(hd) - torch.manual_seed(42) - - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k_comp = torch.randn(n_comp, hd, 1, dtype=torch.bfloat16, device='cuda') - v_comp = torch.randn(n_comp, hd, dtype=torch.bfloat16, device='cuda') - k_swa = torch.randn(n_swa, hd, 1, dtype=torch.bfloat16, device='cuda') - v_swa = torch.randn(n_swa, hd, dtype=torch.bfloat16, device='cuda') - - # Sink weight (in log domain, per-head). n_h=1 so it's a scalar. - attn_sink_val = 0.5 - attn_sink = torch.tensor([attn_sink_val], dtype=torch.float32, device='cuda') - - # === FP32 References === - qf = q[:, :, 0] - - # Reference 1: Combined softmax with sink bias (our kernel's approach) - ref_combined = reference_combined_attention( - qf, k_comp[:, :, 0], v_comp, k_swa[:, :, 0], v_swa, - attn_sink_val, scale, swa_len - ) - - # Reference 2: Separate attention + sink merge (original D5b formula) - ref_merge = reference_sink_merge( - qf, k_comp[:, :, 0], v_comp, k_swa[:, :, 0], v_swa, - attn_sink_val, scale, swa_len - ) - - # Verify the two references agree - cos_ref = torch.nn.functional.cosine_similarity( - ref_combined.flatten().unsqueeze(0).float(), - ref_merge.flatten().unsqueeze(0).float() - ).item() - print(f"Reference: combined softmax vs sink merge cos = {cos_ref:.6f}") - assert cos_ref > 0.999, f"References don't match: cos={cos_ref}" - - # === Kernel === - # Concatenate KV for the kernel - k_combined = torch.cat([k_comp, k_swa], dim=0) # (n_total, hd, 1) - v_combined = torch.cat([v_comp, v_swa], dim=0) # (n_total, hd) - - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - kernel = FmhaKernel( - head_dim=hd, - s_k=n_total, # combined KV length - normalize=False, # D5a: emit un-normalized O + LSE - apply_swa_mask=True, # D3: mask SWA positions - is_causal=False, # D4: no causal mask for this test - n_comp=n_comp, # D5c: compressed KV length (sink bias starts here) - apply_sink_bias=True, # D5c: enable sink bias logit modification - ) - - # Allocate output - c_out = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - lse_out = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - row_sum_out = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - - # Prepare CuTe tensors - def to_cute(t): - return ct.from_dlpack(t).mark_layout_dynamic(leading_dim=ct.get_leading_dim(t)) - - mQ = to_cute(q) - mK = to_cute(k_combined) - mV = to_cute(v_combined.unsqueeze(-1).contiguous()) - mC = to_cute(c_out) - mLSE = to_cute(lse_out) - mRowSums = to_cute(row_sum_out) - - # Compile - print('Compiling D5c kernel (combined KV + sink bias)...', flush=True) - mSinkBias = to_cute(attn_sink) - compiled = cute.compile( - kernel, mQ, mK, mV, mC, stream, mLSE, - swa_len=swa_len, sink_bias=mSinkBias, row_sums=mRowSums, - ) - - # Run - print('Running D5c kernel...', flush=True) - compiled( - mQ, mK, mV, mC, stream, mLSE, - swa_len=swa_len, sink_bias=mSinkBias, row_sums=mRowSums, - ) - torch.cuda.synchronize() - - # Check results - # Kernel outputs UN-NORMALIZED O (normalize=False). Normalize using per-row row_sum. - # O_norm[i] = O_unnorm[i] / row_sum[i] - o_kernel_unnorm = c_out[:, :, 0].float() # (m, hd) - row_sums = row_sum_out[:, 0, 0].float() # (m,) - - # Normalize each row by its row_sum - o_kernel = o_kernel_unnorm / row_sums.unsqueeze(1).clamp(min=1e-30) - - cos = torch.nn.functional.cosine_similarity( - o_kernel.flatten().unsqueeze(0), - ref_combined.flatten().unsqueeze(0).float() - ).item() - max_abs = (o_kernel - ref_combined.float()).abs().max().item() - - status = "PASS" if cos >= 0.99 else "FAIL" - print(f'\nD5c result: cos {cos:.6f} max_abs {max_abs:.4f} {status}') - - if cos < 0.99: - print(f' kernel[0,:4]={o_kernel[0,:4].tolist()}') - print(f' ref[0,:4]={ref_combined[0,:4].tolist()}') - print(f' row_sum range: {row_sums.min().item():.4f} to {row_sums.max().item():.4f}') - print(f' LSE range: {lse_out[:,0,0].min().item():.4f} to {lse_out[:,0,0].max().item():.4f}') - - -def test_d5c_with_causal(): - """D5c with causal mask on SWA branch.""" - print("\n=== Stage D5c: Fused Sparse+SWA with Causal Mask ===\n") - - hd = 64 - m = 128 - n_comp = 64 - n_swa = 64 - n_total = n_comp + n_swa # 128, single KV tile - swa_len = 48 # partially filled SWA - scale = 1.0 / math.sqrt(hd) - torch.manual_seed(123) - - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k_comp = torch.randn(n_comp, hd, 1, dtype=torch.bfloat16, device='cuda') - v_comp = torch.randn(n_comp, hd, dtype=torch.bfloat16, device='cuda') - k_swa = torch.randn(n_swa, hd, 1, dtype=torch.bfloat16, device='cuda') - v_swa = torch.randn(n_swa, hd, dtype=torch.bfloat16, device='cuda') - - attn_sink_val = 0.3 - attn_sink = torch.tensor([attn_sink_val], dtype=torch.float32, device='cuda') - qf = q[:, :, 0] - - ref = reference_combined_attention( - qf, k_comp[:, :, 0], v_comp, k_swa[:, :, 0], v_swa, - attn_sink_val, scale, swa_len, is_causal=True - ) - - # Concatenate KV - k_combined = torch.cat([k_comp, k_swa], dim=0) - v_combined = torch.cat([v_comp, v_swa], dim=0) - - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - kernel = FmhaKernel( - head_dim=hd, - s_k=n_total, - normalize=False, - apply_swa_mask=True, - is_causal=True, - n_comp=n_comp, - apply_sink_bias=True, - ) - - c_out = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - lse_out = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - row_sum_out = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - - def to_cute(t): - return ct.from_dlpack(t).mark_layout_dynamic(leading_dim=ct.get_leading_dim(t)) - - mQ = to_cute(q) - mK = to_cute(k_combined) - mV = to_cute(v_combined.unsqueeze(-1).contiguous()) - mC = to_cute(c_out) - mLSE = to_cute(lse_out) - mRowSums = to_cute(row_sum_out) - mSinkBias = to_cute(attn_sink) - compiled = cute.compile( - kernel, mQ, mK, mV, mC, stream, mLSE, - swa_len=swa_len, sink_bias=mSinkBias, row_sums=mRowSums, - ) - compiled( - mQ, mK, mV, mC, stream, mLSE, - swa_len=swa_len, sink_bias=mSinkBias, row_sums=mRowSums, - ) - torch.cuda.synchronize() - - o_kernel_unnorm = c_out[:, :, 0].float() - row_sums = row_sum_out[:, 0, 0].float() - o_kernel = o_kernel_unnorm / row_sums.unsqueeze(1).clamp(min=1e-30) - - cos = torch.nn.functional.cosine_similarity( - o_kernel.flatten().unsqueeze(0), - ref.flatten().unsqueeze(0).float() - ).item() - max_abs = (o_kernel - ref.float()).abs().max().item() - - status = "PASS" if cos >= 0.99 else "FAIL" - print(f'D5c causal result: cos {cos:.6f} max_abs {max_abs:.4f} {status}') - - -if __name__ == '__main__': - test_d5c_combined() - test_d5c_with_causal() diff --git a/tests/unit/test_d5c_multitile.py b/tests/unit/test_d5c_multitile.py deleted file mode 100644 index 6aa1877c..00000000 --- a/tests/unit/test_d5c_multitile.py +++ /dev/null @@ -1,166 +0,0 @@ -""" -FMHA D5c: Sink bias + Python KV merge for multi-tile (s_k > 128). - -Run: ~/.openclaw/workspace/fire_b200_test tests/unit/test_d5c_multitile.py -""" -import torch -import math -import cutlass.cute as cute -import cutlass.torch as ct -import cuda.bindings.driver as cuda -from dsv4.kernels.attention.fmha import FmhaKernel - - -def to_cute(t): - return ct.from_dlpack(t).mark_layout_dynamic(leading_dim=ct.get_leading_dim(t)) - - -def reference_combined_attention(q, k_comp, v_comp, k_swa, v_swa, - attn_sink, scale, swa_len, is_causal=False): - m, hd = q.shape - n_comp = k_comp.shape[0] - n_swa = k_swa.shape[0] - k_comb = torch.cat([k_comp, k_swa], dim=0) - v_comb = torch.cat([v_comp, v_swa], dim=0) - scores = q.float() @ k_comb.float().T * scale - scores[:, n_comp:] += attn_sink - if swa_len < n_swa: - scores[:, n_comp + swa_len:] = float('-inf') - if is_causal: - for i in range(m): - for j in range(n_swa): - if j > i: - scores[i, n_comp + j] = float('-inf') - max_s = scores.max(dim=-1, keepdim=True).values - exp_s = (scores - max_s).exp() - sum_s = exp_s.sum(dim=-1, keepdim=True).clamp(min=1e-30) - return (exp_s / sum_s @ v_comb.float()).to(torch.bfloat16) - - -def python_kv_merge(segment_results): - """O = sum_i(exp(lse_i) * O_i_norm) / sum_i(exp(lse_i))""" - lse_stack = torch.stack([r[1] for r in segment_results], dim=0) - lse_max = lse_stack.max(dim=0).values - numerator = torch.zeros_like(segment_results[0][0]) - denominator = torch.zeros(lse_max.shape[0], dtype=torch.float32, device=lse_max.device) - for o_norm, lse in segment_results: - exp_lse = (lse - lse_max).exp() - numerator += exp_lse.unsqueeze(1) * o_norm.float() - denominator += exp_lse - return numerator / denominator.unsqueeze(1).clamp(min=1e-30) - - -def run_fmha(q, k, v, kernel_obj, compiled, stream, swa_len, sink_bias=None): - """Run FMHA with single KV tile, return (O_norm, LSE).""" - m = q.shape[0] - hd = v.shape[1] - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - lse = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - rs = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - - mQ = to_cute(q.unsqueeze(-1)) - mK = to_cute(k.unsqueeze(-1)) - mV = to_cute(v.unsqueeze(-1)) - mC = to_cute(c); mLSE = to_cute(lse); mRS = to_cute(rs) - - if sink_bias is not None: - mSB = to_cute(sink_bias) - compiled(mQ, mK, mV, mC, stream, mLSE, swa_len=swa_len, sink_bias=mSB, row_sums=mRS) - else: - compiled(mQ, mK, mV, mC, stream, mLSE, swa_len=swa_len, row_sums=mRS) - torch.cuda.synchronize() - - o_norm = c[:, :, 0].float() / rs[:, 0, 0].float().unsqueeze(1).clamp(min=1e-30) - return o_norm, lse[:, 0, 0].clone() - - -def test_d5c_multitile(): - print("=== D5c Multi-Tile: Sink Bias + Python KV Merge ===\n") - - hd = 64; m = 128; n_comp = 96; n_swa = 160; n_total = 256 - swa_len = 100; scale = 1.0 / math.sqrt(hd); seg_size = 128 - torch.manual_seed(42) - - q = torch.randn(m, hd, dtype=torch.bfloat16, device='cuda') - k_comp = torch.randn(n_comp, hd, dtype=torch.bfloat16, device='cuda') - v_comp = torch.randn(n_comp, hd, dtype=torch.bfloat16, device='cuda') - k_swa = torch.randn(n_swa, hd, dtype=torch.bfloat16, device='cuda') - v_swa = torch.randn(n_swa, hd, dtype=torch.bfloat16, device='cuda') - attn_sink_val = 0.5 - attn_sink = torch.tensor([attn_sink_val], dtype=torch.float32, device='cuda') - - k_combined = torch.cat([k_comp, k_swa], dim=0) - v_combined = torch.cat([v_comp, v_swa], dim=0) - - # Full reference - ref = reference_combined_attention(q, k_comp, v_comp, k_swa, v_swa, - attn_sink_val, scale, swa_len) - - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - # Per-segment reference and kernel verification - # Segment 0: positions 0-127 (n_comp=96, 32 SWA tokens, all valid, sink bias on SWA) - k0 = k_combined[0:128]; v0 = v_combined[0:128] - scores0 = q.float() @ k0.float().T * scale - scores0[:, n_comp:] += attn_sink_val - # swa_len=100 but only 32 SWA positions → no D3 masking - ref0 = (torch.softmax(scores0, dim=-1) @ v0.float()).to(torch.bfloat16) - lse0_ref = torch.logsumexp(scores0, dim=-1) - - k_seg0 = FmhaKernel(head_dim=hd, s_k=128, normalize=False, - apply_swa_mask=True, is_causal=False, n_comp=n_comp, apply_sink_bias=True) - c0 = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - l0 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - r0 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - comp0 = cute.compile(k_seg0, to_cute(q.unsqueeze(-1)), to_cute(k0.unsqueeze(-1)), - to_cute(v0.unsqueeze(-1)), to_cute(c0), stream, to_cute(l0), - swa_len=100, sink_bias=to_cute(attn_sink), row_sums=to_cute(r0)) - comp0(to_cute(q.unsqueeze(-1)), to_cute(k0.unsqueeze(-1)), - to_cute(v0.unsqueeze(-1)), to_cute(c0), stream, to_cute(l0), - swa_len=100, sink_bias=to_cute(attn_sink), row_sums=to_cute(r0)) - torch.cuda.synchronize() - ok0 = c0[:, :, 0].float() / r0[:, 0, 0].float().unsqueeze(1).clamp(min=1e-30) - cos0 = torch.nn.functional.cosine_similarity(ok0.flatten().unsqueeze(0), ref0.flatten().unsqueeze(0).float()).item() - lse0_kern = l0[:, 0, 0] - print(f'Seg0: cos {cos0:.6f} LSE_kern[0]={lse0_kern[0].item():.4f} LSE_ref[0]={lse0_ref[0].item():.4f}') - - # Segment 1: positions 128-255 (all SWA, n_comp=0, sink bias on all, D3 mask at position 68) - k1 = k_combined[128:256]; v1 = v_combined[128:256] - scores1 = q.float() @ k1.float().T * scale - scores1 += attn_sink_val # all SWA → sink bias on all - # Valid SWA: absolute 96-195. In this segment (128-255): positions 0-67 valid, 68-127 masked - swa_len_seg1 = 68 # n_comp + swa_len - k_start = 96 + 100 - 128 = 68 - scores1[:, swa_len_seg1:] = float('-inf') - ref1 = (torch.softmax(scores1, dim=-1) @ v1.float()).to(torch.bfloat16) - lse1_ref = torch.logsumexp(scores1, dim=-1) - - k_seg1 = FmhaKernel(head_dim=hd, s_k=128, normalize=False, - apply_swa_mask=True, is_causal=False, n_comp=0, apply_sink_bias=True) - c1 = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - l1 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - r1 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - comp1 = cute.compile(k_seg1, to_cute(q.unsqueeze(-1)), to_cute(k1.unsqueeze(-1)), - to_cute(v1.unsqueeze(-1)), to_cute(c1), stream, to_cute(l1), - swa_len=swa_len_seg1, sink_bias=to_cute(attn_sink), row_sums=to_cute(r1)) - comp1(to_cute(q.unsqueeze(-1)), to_cute(k1.unsqueeze(-1)), - to_cute(v1.unsqueeze(-1)), to_cute(c1), stream, to_cute(l1), - swa_len=swa_len_seg1, sink_bias=to_cute(attn_sink), row_sums=to_cute(r1)) - torch.cuda.synchronize() - ok1 = c1[:, :, 0].float() / r1[:, 0, 0].float().unsqueeze(1).clamp(min=1e-30) - cos1 = torch.nn.functional.cosine_similarity(ok1.flatten().unsqueeze(0), ref1.flatten().unsqueeze(0).float()).item() - lse1_kern = l1[:, 0, 0] - print(f'Seg1: cos {cos1:.6f} LSE_kern[0]={lse1_kern[0].item():.4f} LSE_ref[0]={lse1_ref[0].item():.4f}') - - # Merge - o_merged = python_kv_merge([(ok0, lse0_kern), (ok1, lse1_kern)]) - cos = torch.nn.functional.cosine_similarity( - o_merged.flatten().unsqueeze(0), ref.flatten().unsqueeze(0).float()).item() - status = "PASS" if cos >= 0.99 else "FAIL" - print(f'\nMerged: cos {cos:.6f} {status}') - if cos < 0.99: - print(f' kernel[0,:4]={o_merged[0,:4].tolist()}') - print(f' ref[0,:4]={ref[0,:4].tolist()}') - - -if __name__ == '__main__': - test_d5c_multitile() diff --git a/tests/unit/test_fmha_pv16.py b/tests/unit/test_fmha_pv16.py deleted file mode 100644 index 8406ffee..00000000 --- a/tests/unit/test_fmha_pv16.py +++ /dev/null @@ -1,52 +0,0 @@ -"""Test FMHA with pv_n_tile=16 (N=16 sub-tiles for PV GEMM).""" -import torch -import math -import sys -import os - -sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..')) - -from dsv4.kernels.attention.production import dsv4_attention_per_head - - -def test_fmha_pv16(hd): - sk = 128 - scale = 1.0 / math.sqrt(hd) - - torch.manual_seed(42) - q = torch.randn(1, 1, hd, dtype=torch.bfloat16, device='cuda') - k = torch.randn(sk, hd, dtype=torch.bfloat16, device='cuda') - v = torch.randn(hd, sk, dtype=torch.bfloat16, device='cuda') - - o = dsv4_attention_per_head(q, k, v, scale=scale, swa_len=sk) - - # Reference - q_ref = q[0, 0].float() - k_ref = k.float() - v_ref = v.float() - - s = (q_ref @ k_ref.T) * scale - p = torch.softmax(s, dim=-1) - o_ref = (p @ v_ref.T).to(torch.bfloat16) - - o_f = o[0, 0].float() - o_ref_f = o_ref.float() - - cs = torch.nn.functional.cosine_similarity(o_f.unsqueeze(0), o_ref_f.unsqueeze(0)).item() - print(f"HD={hd} pv_n_tile=16: cosine={cs:.8f}") - - if cs < 0.999: - print(f" FAILED") - print(f" o[0:4] = {o_f[0:4].tolist()}") - print(f" o_ref[0:4] = {o_ref_f[0:4].tolist()}") - return False - print(f" PASSED") - return True - - -if __name__ == '__main__': - all_pass = True - for hd in [16, 64, 128]: - if not test_fmha_pv16(hd): - all_pass = False - sys.exit(0 if all_pass else 1) diff --git a/tests/unit/test_fmha_v3_stage_d1.py b/tests/unit/test_fmha_v3_stage_d1.py deleted file mode 100644 index 2302271b..00000000 --- a/tests/unit/test_fmha_v3_stage_d1.py +++ /dev/null @@ -1,113 +0,0 @@ -""" -FMHA v3 Stage D1: Parameterized HEAD_DIM (64 → 512). - -The kernel outputs un-normalized O + LSE. -Test compares un-normalized O against FP32 reference. -External normalization (O_norm = O_unnorm / row_sum) uses LSE for the D5 merge. -""" -import torch, math -import cutlass.cute as cute -import cutlass.torch as ct -import cuda.bindings.driver as cuda -from dsv4.kernels.attention.fmha import FmhaKernel - - -def test_head_dim(hd, n_kv=128): - m = 128 - torch.manual_seed(42) - - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(n_kv, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(n_kv, hd, dtype=torch.bfloat16, device='cuda') - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - - # FP32 reference - qf = q[:, :, 0].float() - kf = k[:, :, 0].float() - scale = 1.0 / math.sqrt(hd) - attn_max = (qf @ kf.T * scale).max(dim=-1, keepdim=True)[0] - attn_exp = torch.exp(qf @ kf.T * scale - attn_max) - attn_sum = attn_exp.sum(dim=-1, keepdim=True) - ref_norm = (attn_exp / attn_sum) @ v.float() - ref_unnorm = attn_exp @ v.float() - ref_lse = (torch.log(attn_sum.squeeze(-1)) + attn_max.squeeze(-1))[0].item() - - lse_tensor = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - - # Use TMEM-P for all hd (Saves SMEM, and works correctly after qk_mma_tiler fix) - # SMEM-P is available for hd>64 but requires more SMEM budget - kernel = FmhaKernel(head_dim=hd, s_k=n_kv, use_smem_p=False) - pv_n_tile = kernel.pv_n_tile - n_pv_tiles = kernel.n_pv_tiles - - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - # Compile once - v_tile = v[:, 0:pv_n_tile].contiguous() - v_kernel = v_tile.unsqueeze(-1) - c_tile = torch.zeros(m, pv_n_tile, 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_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) - - print(f'hd={hd}, n={n_kv} (pv_n_tile={pv_n_tile}, n_pv_tiles={n_pv_tiles}): Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE) - - lse_val = None - for nt in range(n_pv_tiles): - v_start = nt * pv_n_tile - v_end = v_start + pv_n_tile - v_tile = v[:, v_start:v_end].contiguous() - v_kernel = v_tile.unsqueeze(-1) - c_tile = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - lse_tensor.zero_() - - 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_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) - - compiled(mQ, mK, mV, mC, stream, mLSE) - torch.cuda.synchronize() - - c[:, v_start:v_end, :] = c_tile - if nt == 0: - lse_val = lse_tensor[0, 0, 0].item() - - out_unnorm = c[:, :, 0].float() - out_norm = out_unnorm / attn_sum - cos_unnorm = torch.nn.functional.cosine_similarity( - out_unnorm.flatten().unsqueeze(0), ref_unnorm.flatten().unsqueeze(0) - ).item() - cos_norm = torch.nn.functional.cosine_similarity( - out_norm.flatten().unsqueeze(0), ref_norm.flatten().unsqueeze(0) - ).item() - lse_err = abs(lse_val - ref_lse) if lse_val is not None else float('inf') - - status = "PASS" if cos_unnorm >= 0.99 else "FAIL" - print(f'hd={hd}, n={n_kv}: cos_unnorm {cos_unnorm:.6f} cos_norm {cos_norm:.6f} lse_err {lse_err:.6f} {status}') - return cos_unnorm - - -def test(): - print("=== Stage D1: Parameterized HEAD_DIM ===") - print("(Kernel outputs un-normalized O + LSE; TMEM-P path)\n") - - cos64 = test_head_dim(64, 128) - cos128 = test_head_dim(128, 128) - cos256 = test_head_dim(256, 128) - cos512 = test_head_dim(512, 128) - - print("\n=== Summary ===") - print(f"hd=64: cos={cos64:.6f} {'PASS' if cos64 >= 0.99 else 'FAIL'}") - print(f"hd=128: cos={cos128:.6f} {'PASS' if cos128 >= 0.99 else 'FAIL'}") - print(f"hd=256: cos={cos256:.6f} {'PASS' if cos256 >= 0.99 else 'FAIL'}") - print(f"hd=512: cos={cos512:.6f} {'PASS' if cos512 >= 0.99 else 'FAIL'}") - - -if __name__ == '__main__': - test() diff --git a/tests/unit/test_fmha_v3_stage_d5b.py b/tests/unit/test_fmha_v3_stage_d5b.py deleted file mode 100644 index e8027be6..00000000 --- a/tests/unit/test_fmha_v3_stage_d5b.py +++ /dev/null @@ -1,204 +0,0 @@ -""" -FMHA v3 Stage D5b: SWA + Sink Merge (Python-level). - -Tests the full DSV4 attention pipeline: -1. Run FMHA with compressed KV (normalize=False) → o_unnorm, lse -2. Run FMHA with SWA KV (normalize=False) → o_unnorm, lse -3. Merge with sink weights in Python: - numerator = o_unnorm_comp + exp(attn_sink) * o_unnorm_swa - denominator = exp(lse_comp) + exp(attn_sink) * exp(lse_swa) - output = numerator / denominator - -Uses hd=64 TMEM-P path. The un-normalized merge formula is mathematically -equivalent to the normalized merge but avoids the exp(lse)*o_norm multiply -that can lose precision for large lse values. -""" -import torch, math -import cutlass.cute as cute -import cutlass.torch as ct -import cuda.bindings.driver as cuda -from dsv4.kernels.attention.fmha import FmhaKernel - - -def run_fmha_unnorm(q, k, v, kernel_obj, compiled_kernel, stream): - """Run FMHA with normalize=False, return un-normalized O and LSE.""" - m = 128 # M tile - hd = v.shape[1] - pv_n_tile = kernel_obj.pv_n_tile - n_pv_tiles = kernel_obj.n_pv_tiles - - c_unnorm = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - - for nt in range(n_pv_tiles): - v_start = nt * pv_n_tile - v_end = v_start + pv_n_tile - v_tile = v[:, v_start:v_end].contiguous().unsqueeze(-1) - c_tile = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - lse_tile = torch.zeros(m, 1, 1, dtype=torch.float32, 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_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_tile)) - mC = ct.from_dlpack(c_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tile)) - - compiled_kernel(mQ, mK, mV, mC, stream, mLSE) - torch.cuda.synchronize() - - c_unnorm[:, v_start:v_end, :] = c_tile - if nt == 0: - lse_val = lse_tile[0, 0, 0].item() - - o_unnorm = c_unnorm[:, :, 0] # (m, hd) BF16 - return o_unnorm, lse_val - - -def test(): - print("=== Stage D5b: SWA + Sink Merge (Python) ===\n") - - hd = 64 - m = 128 - n_kv = 128 # same length for both compressed and SWA KV - torch.manual_seed(42) - - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k_comp = torch.randn(n_kv, hd, 1, dtype=torch.bfloat16, device='cuda') - v_comp = torch.randn(n_kv, hd, dtype=torch.bfloat16, device='cuda') - k_swa = torch.randn(n_kv, hd, 1, dtype=torch.bfloat16, device='cuda') - v_swa = torch.randn(n_kv, hd, dtype=torch.bfloat16, device='cuda') - - # Per-head sink weight (learnable parameter, in LOG domain) - attn_sink = torch.tensor([0.5], dtype=torch.float32, device='cuda') - scale = 1.0 / math.sqrt(hd) - - # === FP32 Reference: normalized merge === - qf = q[:, :, 0].float() - kf_comp = k_comp[:, :, 0].float() - vf_comp = v_comp.float() - kf_swa = k_swa[:, :, 0].float() - vf_swa = v_swa.float() - - # Compressed KV attention - attn_comp = qf @ kf_comp.T * scale - o_norm_comp = torch.softmax(attn_comp, dim=-1) @ vf_comp - attn_comp_max = attn_comp.max(dim=-1, keepdim=True)[0] - lse_comp = torch.logsumexp(attn_comp, dim=-1, keepdim=True) # (m, 1) - - # SWA KV attention - attn_swa = qf @ kf_swa.T * scale - o_norm_swa = torch.softmax(attn_swa, dim=-1) @ vf_swa - lse_swa = torch.logsumexp(attn_swa, dim=-1, keepdim=True) - - # Un-normalized outputs: o_unnorm = exp(attn - max) @ V (NOT exp(attn) @ V) - attn_comp_exp = torch.exp(attn_comp - attn_comp_max) - attn_swa_exp = torch.exp(attn_swa - attn_swa.max(dim=-1, keepdim=True)[0]) - o_unnorm_comp_ref = attn_comp_exp @ vf_comp - o_unnorm_swa_ref = attn_swa_exp @ vf_swa - - # Normalized merge (reference, numerically stable) - lse_max = torch.max(lse_comp, lse_swa) - exp_lse_comp = torch.exp(lse_comp - lse_max) - exp_lse_swa = torch.exp(lse_swa - lse_max) - exp_sink = attn_sink.exp() - - numerator_norm = (exp_lse_comp * o_norm_comp + exp_sink * exp_lse_swa * o_norm_swa) - denominator_norm = (exp_lse_comp + exp_sink * exp_lse_swa).clamp(min=1e-30) - ref_output = numerator_norm / denominator_norm # (m, hd) - - # Un-normalized merge (same result, different form): - # numerator = o_unnorm_comp + exp(sink) * o_unnorm_swa - # denominator = exp(lse_comp) + exp(sink) * exp(lse_swa) - # But o_unnorm = o_norm * exp(lse), so: - # numerator = o_norm_comp * exp(lse_comp) + exp(sink) * o_norm_swa * exp(lse_swa) - # This is identical to the normalized formula. Let's verify: - numerator_unnorm = o_unnorm_comp_ref + exp_sink * o_unnorm_swa_ref - denominator_unnorm = (lse_comp.exp() + exp_sink * lse_swa.exp()).clamp(min=1e-30) - ref_output_unnorm = numerator_unnorm / denominator_unnorm - - unnorm_vs_norm_cos = torch.nn.functional.cosine_similarity( - ref_output.flatten().unsqueeze(0), - ref_output_unnorm.flatten().unsqueeze(0) - ).item() - print(f"Reference: normalized vs unnorm cos = {unnorm_vs_norm_cos:.6f}") - - # Use the numerically stable (log-sum-exp) version for reference - # Un-normalized merge with log-sum-exp stability: - # Multiply num and denom by exp(-lse_max) - numerator_unnorm_stable = (o_unnorm_comp_ref * exp_lse_comp + exp_sink * o_unnorm_swa_ref * exp_lse_swa) - denominator_unnorm_stable = (exp_lse_comp + exp_sink * exp_lse_swa).clamp(min=1e-30) - ref_output_stable = numerator_unnorm_stable / denominator_unnorm_stable - - stable_cos = torch.nn.functional.cosine_similarity( - ref_output.flatten().unsqueeze(0), - ref_output_stable.flatten().unsqueeze(0) - ).item() - print(f"Reference: stable unnorm merge cos = {stable_cos:.6f}") - - # === Kernel: Run FMHA twice (normalize=False) and merge === - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - kernel = FmhaKernel(head_dim=hd, s_k=n_kv, normalize=False) - - # Compile - print('Compiling kernel...', flush=True) - v_tile = v_comp[:, 0:kernel.pv_n_tile].contiguous().unsqueeze(-1) - c_tile = torch.zeros(m, kernel.pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - lse_tile = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) - mK = ct.from_dlpack(k_comp).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k_comp)) - mV = ct.from_dlpack(v_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_tile)) - mC = ct.from_dlpack(c_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tile)) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE) - - # Run compressed KV - print('Running compressed KV...', flush=True) - o_kern_comp, lse_kern_comp = run_fmha_unnorm(q, k_comp, v_comp, kernel, compiled, stream) - - # Run SWA KV - print('Running SWA KV...', flush=True) - o_kern_swa, lse_kern_swa = run_fmha_unnorm(q, k_swa, v_swa, kernel, compiled, stream) - - # Kernel-level merge with sink weights - lse_comp_t = torch.tensor(lse_kern_comp, dtype=torch.float32, device='cuda') - lse_swa_t = torch.tensor(lse_kern_swa, dtype=torch.float32, device='cuda') - exp_lse_kern_comp = lse_comp_t.exp() - exp_lse_kern_swa = lse_swa_t.exp() - exp_sink_kern = attn_sink[0].exp() - - # numerator = o_unnorm_comp + exp(sink) * o_unnorm_swa - # denominator = exp(lse_comp) + exp(sink) * exp(lse_swa) - kern_numerator = o_kern_comp.float() + exp_sink_kern * o_kern_swa.float() - kern_denominator = (exp_lse_kern_comp + exp_sink_kern * exp_lse_kern_swa).clamp(min=1e-30) - kern_output = kern_numerator / kern_denominator # (m, hd) - - # Compare with reference (use the stable version) - cos = torch.nn.functional.cosine_similarity( - kern_output.flatten().unsqueeze(0), - ref_output_stable.flatten().unsqueeze(0) - ).item() - max_abs = (kern_output - ref_output_stable).abs().max().item() - - status = "PASS" if cos >= 0.93 else "FAIL" - print(f'\nMerge result: cos {cos:.6f} max_abs {max_abs:.4f} {status}') - - # Check individual attention passes - cos_comp = torch.nn.functional.cosine_similarity( - o_kern_comp.flatten().unsqueeze(0).float(), - o_unnorm_comp_ref.flatten().unsqueeze(0) - ).item() - cos_swa = torch.nn.functional.cosine_similarity( - o_kern_swa.flatten().unsqueeze(0).float(), - o_unnorm_swa_ref.flatten().unsqueeze(0) - ).item() - print(f' Compressed KV unnorm cos: {cos_comp:.6f}') - print(f' SWA KV unnorm cos: {cos_swa:.6f}') - print(f' LSE comp: kernel={lse_kern_comp:.6f} ref={lse_comp[0,0].item():.6f} err={abs(lse_kern_comp-lse_comp[0,0].item()):.6f}') - print(f' LSE swa: kernel={lse_kern_swa:.6f} ref={lse_swa[0,0].item():.6f} err={abs(lse_kern_swa-lse_swa[0,0].item()):.6f}') - - if cos < 0.93: - print(f' kern[0,:4]={kern_output[0,:4].tolist()}') - print(f' ref[0,:4]={ref_output_stable[0,:4].tolist()}') - - -if __name__ == '__main__': - test() diff --git a/tests/unit/test_production.py b/tests/unit/test_production.py index ba94aa03..f2d73d01 100644 --- a/tests/unit/test_production.py +++ b/tests/unit/test_production.py @@ -13,7 +13,7 @@ Tests: import torch import math import pytest -from dsv4.kernels.attention.production import dsv4_attention, dsv4_attention_per_head +from dsv4.kernels.attention.production import dsv4_attention, dsv4_attention_per_head # noqa: E402 # --------------------------------------------------------------------------- diff --git a/tests/unit/test_smem_acc.py b/tests/unit/test_smem_acc.py deleted file mode 100644 index 3e8cc07c..00000000 --- a/tests/unit/test_smem_acc.py +++ /dev/null @@ -1,118 +0,0 @@ -""" -Test SMEM accumulator FMHA kernel: multi-KV-tile with in-kernel O accumulation. -No Python KV merge needed — the kernel handles acc_scale internally. -""" -import torch, math, sys -import cutlass.cute as cute -import cutlass.torch as ct -import cuda.bindings.driver as cuda -from dsv4.kernels.attention.fmha_smem_acc import FmhaKernel - - -def test_smem_acc(hd=64, s_k=256, use_smem_p=False, normalize=False): - m = 128 - n_kv_tiles = s_k // 128 - torch.manual_seed(42) - - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - - # FP32 reference - qf = q[:, :, 0].float() - kf = k[:, :, 0].float() - scale = 1.0 / math.sqrt(hd) - attn_max = (qf @ kf.T * scale).max(dim=-1, keepdim=True)[0] - attn_exp = torch.exp(qf @ kf.T * scale - attn_max) - attn_sum = attn_exp.sum(dim=-1, keepdim=True) - ref_norm = (attn_exp / attn_sum) @ v.float() - ref_unnorm = attn_exp @ v.float() - - lse_tensor = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - row_sums_tensor = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - - kernel = FmhaKernel(head_dim=hd, s_k=s_k, use_smem_p=use_smem_p, normalize=normalize) - pv_n_tile = kernel.pv_n_tile - n_pv_tiles = kernel.n_pv_tiles - - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - # Compile - v_tile = v[:, 0:pv_n_tile].contiguous() - v_kernel = v_tile.unsqueeze(-1) - c_tile = torch.zeros(m, pv_n_tile, 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_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) - mRS = ct.from_dlpack(row_sums_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(row_sums_tensor)) - - # Simple GMEM tensor (non-dynamic-layout) for SMEM accumulator TMA store - c_simple_tensor = c_tile.clone() - mCSimple = ct.from_dlpack(c_simple_tensor) # No mark_layout_dynamic! - - print(f' hd={hd}, s_k={s_k} ({n_kv_tiles} KV tiles, pv_n_tile={pv_n_tile}, n_pv_tiles={n_pv_tiles}): Compiling...', flush=True) - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, lse=mLSE, row_sums=mRS, c_simple=mCSimple) - - for nt in range(n_pv_tiles): - v_start = nt * pv_n_tile - v_end = v_start + pv_n_tile - v_tile = v[:, v_start:v_end].contiguous() - v_kernel = v_tile.unsqueeze(-1) - c_tile = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - lse_tensor.zero_() - row_sums_tensor.zero_() - - 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_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) - mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) - mRS = ct.from_dlpack(row_sums_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(row_sums_tensor)) - - mCSimple = ct.from_dlpack(c_tile) # No mark_layout_dynamic! - - compiled(mQ, mK, mV, mC, stream, lse=mLSE, row_sums=mRS, c_simple=mCSimple) - torch.cuda.synchronize() - - c[:, v_start:v_end, :] = c_tile - - out = c[:, :, 0].float() - - if normalize: - cos = torch.nn.functional.cosine_similarity( - out.flatten().unsqueeze(0), ref_norm.flatten().unsqueeze(0) - ).item() - ref = ref_norm - else: - cos = torch.nn.functional.cosine_similarity( - out.flatten().unsqueeze(0), ref_unnorm.flatten().unsqueeze(0) - ).item() - ref = ref_unnorm - - status = "PASS" if cos >= 0.99 else "FAIL" - print(f' hd={hd}, s_k={s_k} ({n_kv_tiles} tiles): cos {cos:.6f} {status}') - return cos - - -def test(): - print("=== SMEM Accumulator FMHA: In-Kernel Multi-KV-Tile O Accumulation ===\n") - - # Single KV tile (s_k=128): should work like fmha.py - print("--- Single KV tile (s_k=128) ---") - test_smem_acc(64, 128) - test_smem_acc(128, 128) - - # Multi KV tile: the SMEM accumulator approach should handle this correctly - print("\n--- Multi KV tile (s_k=256+) ---") - test_smem_acc(64, 256) - test_smem_acc(64, 384) - test_smem_acc(64, 512) - test_smem_acc(128, 256) - - -if __name__ == '__main__': - test() diff --git a/tests/unit/test_smem_budget.py b/tests/unit/test_smem_budget.py deleted file mode 100644 index 7acd0ad1..00000000 --- a/tests/unit/test_smem_budget.py +++ /dev/null @@ -1,71 +0,0 @@ -"""Diagnostic: Print exact SMEM budget for various HEAD_DIM values.""" -import math -from dsv4.kernels.attention.fmha import FmhaKernel - - -def smem_budget(hd, s_k=128, use_smem_p=None): - """Reproduce FmhaKernel._setup SMEM allocation sizes.""" - kv_stage = 1 if hd > 128 else 2 - q_stage = 1 - k_tile = min(hd, 256) - pv_n_tile = min(hd, 256) - pv_mma_tiler = (128, pv_n_tile, k_tile) - qk_mma_tiler = (128, 128, k_tile) - - # Rough SMEM sizes (actual CuTe layouts have padding/alignment) - # sQ: (M, K_tile) BF16 * q_stage - sQ = 128 * k_tile * 2 * q_stage - # sK: (K_tile, N) BF16 * kv_stage - sK = k_tile * 128 * 2 * kv_stage - # sV: (K_tile_V, N) BF16 * kv_stage, where K_tile_V = pv_n_tile - sV = pv_n_tile * 128 * 2 * kv_stage - # sC: (M, N) BF16 * num_c_stage - num_c_stage = 1 if hd > 256 else 2 - sC = 128 * pv_n_tile * 2 * num_c_stage - # sP: only if SMEM-P - sP = 0 - if use_smem_p is None: - use_smem_p = hd > 64 - if use_smem_p: - sP = 128 * pv_n_tile * 2 * 1 # 1 stage - - total = sQ + sK + sV + sC + sP - limit = 232 * 1024 - - print(f"hd={hd}: k_tile={k_tile}, kv_stage={kv_stage}, pv_n_tile={pv_n_tile}, " - f"num_c_stage={num_c_stage}, use_smem_p={use_smem_p}") - print(f" sQ={sQ//1024}KB, sK={sK//1024}KB, sV={sV//1024}KB, sC={sC//1024}KB, sP={sP//1024}KB") - print(f" Total={total//1024}KB vs {limit//1024}KB limit → {'✅ FITS' if total <= limit else '❌ OVER'}") - print() - - # Now compute with overlap: sQ/sV share same region - if not use_smem_p: - overlap_total = max(sQ, sV) + sK + sC - print(f" With sQ/sV overlap: {overlap_total//1024}KB → {'✅ FITS' if overlap_total <= limit else '❌ OVER'}") - else: - overlap_total = max(sQ, sV) + sK + sC + sP - print(f" With sQ/sV overlap (+sP): {overlap_total//1024}KB → {'✅ FITS' if overlap_total <= limit else '❌ OVER'}") - - # Another option: also overlap sK with sC (K consumed before C written) - # After QK GEMM, sK is done. Then PV starts, writes O to TMEM, then epilogue writes C. - # But sK is needed for ALL kv_tiles... and sC is written after all PV is done. - # Actually sK is needed per kv_tile, so it stays alive through the whole QK loop. - # Overlap sK/sC doesn't work. - - # Option: reduce pv_n_tile to 128 - pv_n_tile_small = 128 - sV_small = pv_n_tile_small * 128 * 2 * kv_stage - sC_small = 128 * pv_n_tile_small * 2 * num_c_stage - total_small = sQ + sK + sV_small + sC_small + sP - print(f" With pv_n_tile=128: sV={sV_small//1024}KB, sC={sC_small//1024}KB, " - f"Total={total_small//1024}KB → {'✅ FITS' if total_small <= limit else '❌ OVER'}") - overlap_small = max(sQ, sV_small) + sK + sC_small - print(f" With pv_n_tile=128 + sQ/sV overlap: {overlap_small//1024}KB → {'✅ FITS' if overlap_small <= limit else '❌ OVER'}") - print() - - -print("=== SMEM Budget Analysis ===\n") -smem_budget(64) -smem_budget(128) -smem_budget(256) -smem_budget(512) diff --git a/tests/unit/test_smem_p_coord.py b/tests/unit/test_smem_p_coord.py deleted file mode 100644 index 7f8fbc0d..00000000 --- a/tests/unit/test_smem_p_coord.py +++ /dev/null @@ -1,137 +0,0 @@ -""" -SMEM-P Coordinate Verification Test. - -Writes a known pattern to sP using the coordinate-indexed approach -(identical to FmhaKernel's SMEM-P path), then reads sP back -and verifies on the host. - -This test uses the FmhaKernel class to set up all layouts (MMA, SMEM, TMEM) -inside the JIT context, then writes a test pattern and reads it back. -""" -import torch, math -import cutlass, cutlass.cute as cute -import cutlass.utils as utils -from cutlass.cute.nvgpu import tcgen05 -from cutlass import Float32, BFloat16, Int32, const_expr -import cutlass.torch as ct -import cuda.bindings.driver as cuda -from dsv4.kernels.attention.fmha import FmhaKernel - - -def test_smem_p_coords(): - head_dim = 256 - s_k = 128 - m = 128 - pv_n_tile = min(head_dim, 256) - - # Use FmhaKernel to do the actual test - # We modify the kernel to write a test pattern instead of P values - kernel = FmhaKernel(head_dim=head_dim, s_k=s_k, use_smem_p=True, normalize=False) - - q = torch.randn(m, head_dim, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(s_k, head_dim, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, head_dim, dtype=torch.bfloat16, device='cuda') - c = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - lse = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - v_tile = v[:, 0:pv_n_tile].contiguous().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_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_tile)) - mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) - mLSE = ct.from_dlpack(lse).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse)) - - print("Compiling FmhaKernel (hd=256, SMEM-P, normalize=False)...") - try: - compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE) - except Exception as e: - print(f"COMPILE FAILED: {e}") - import traceback - traceback.print_exc() - return - print("Running...") - try: - compiled(mQ, mK, mV, mC, stream, mLSE) - except Exception as e: - print(f"RUN FAILED: {e}") - import traceback - traceback.print_exc() - return - torch.cuda.synchronize() - - # The kernel writes P to sP using the coordinate-indexed approach - # then reads it back via PV MMA. The output should be close to - # the reference attention output. - out = c[:, :, 0].float() - - # FP32 reference - qf = q[:, :, 0].float() - kf = k[:, :, 0].float() - scale = 1.0 / math.sqrt(head_dim) - attn = qf @ kf.T * scale - attn = torch.softmax(attn, dim=-1) - ref = attn @ v[:, 0:pv_n_tile].float() - - cos = torch.nn.functional.cosine_similarity( - out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0) - ).item() - print(f"hd=256, n=128: cos {cos:.6f} {'PASS' if cos >= 0.97 else 'FAIL'}") - - if cos < 0.97: - # Print first few output vs reference values - print(f" out[0,:4]={out[0,:4].tolist()}") - print(f" ref[0,:4]={ref[0,:4].tolist()}") - print(f" out[1,:4]={out[1,:4].tolist()}") - print(f" ref[1,:4]={ref[1,:4].tolist()}") - - # Check if output is zero (sP not written) or non-zero but wrong - out_norm = out.norm().item() - ref_norm = ref.norm().item() - print(f" out norm: {out_norm:.4f}, ref norm: {ref_norm:.4f}") - - # Check if output is proportional to ref (scaling issue) - if out_norm > 0 and ref_norm > 0: - scale_ratio = out_norm / ref_norm - scaled_out = out / scale_ratio - scaled_cos = torch.nn.functional.cosine_similarity( - scaled_out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0) - ).item() - print(f" Scaled cos (out/scale_ratio): {scaled_cos:.6f}") - - # Also test hd=64 TMEM-P as regression - print("\n--- Regression: hd=64 TMEM-P ---") - kernel64 = FmhaKernel(head_dim=64, s_k=s_k, use_smem_p=False, normalize=False) - q64 = torch.randn(m, 64, 1, dtype=torch.bfloat16, device='cuda') - k64 = torch.randn(s_k, 64, 1, dtype=torch.bfloat16, device='cuda') - v64 = torch.randn(s_k, 64, dtype=torch.bfloat16, device='cuda') - c64 = torch.zeros(m, 64, 1, dtype=torch.bfloat16, device='cuda') - lse64 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - - mQ64 = ct.from_dlpack(q64).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q64)) - mK64 = ct.from_dlpack(k64).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k64)) - v64_tile = v64.unsqueeze(-1) - mV64 = ct.from_dlpack(v64_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v64_tile)) - mC64 = ct.from_dlpack(c64).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c64)) - mLSE64 = ct.from_dlpack(lse64).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse64)) - - compiled64 = cute.compile(kernel64, mQ64, mK64, mV64, mC64, stream, mLSE64) - compiled64(mQ64, mK64, mV64, mC64, stream, mLSE64) - torch.cuda.synchronize() - - out64 = c64[:, :, 0].float() - qf64 = q64[:, :, 0].float() - kf64 = k64[:, :, 0].float() - scale64 = 1.0 / math.sqrt(64) - attn64 = qf64 @ kf64.T * scale64 - attn64 = torch.softmax(attn64, dim=-1) - ref64 = attn64 @ v64.float() - cos64 = torch.nn.functional.cosine_similarity( - out64.flatten().unsqueeze(0), ref64.flatten().unsqueeze(0) - ).item() - print(f"hd=64, n=128: cos {cos64:.6f} {'PASS' if cos64 >= 0.97 else 'FAIL'}") - - -if __name__ == '__main__': - test_smem_p_coords() diff --git a/tests/unit/test_smem_p_diag.py b/tests/unit/test_smem_p_diag.py deleted file mode 100644 index 7fd12126..00000000 --- a/tests/unit/test_smem_p_diag.py +++ /dev/null @@ -1,516 +0,0 @@ -""" -Diagnostic test for SMEM-P: verify the TV layout mapping from TMEM-load to sP. - -This test does NOT compute attention. It: -1. Creates the TMEM-load copy and extracts its TV layout -2. Creates sP with PV A-operand layout -3. Builds atom_layout_tv mapping (tid,vid) → sP address -4. Validates the mapping is correct (all 16384 elements covered, no overlaps) -5. Tests the make_cotiled_copy approach -""" -import torch, math -import cutlass, cutlass.cute as cute -import cutlass.utils as utils -import cutlass.utils.blackwell_helpers as bh -from cutlass.cute.nvgpu import tcgen05, cpasync -from cutlass import Float32, BFloat16, Int32 -import cutlass.torch as ct -import cuda.bindings.driver as cuda - - -def diag_smem_p_tv(): - """Print TV layout info from TMEM-load copy and sP to build the R→S copy.""" - head_dim = 256 # First SMEM-P head dim (not 64 which uses TMEM-P) - s_k = 128 - m = 128 - pv_n_tile = min(head_dim, 256) - scale_softmax = 1.0 / math.sqrt(head_dim) - scale_softmax_log2 = scale_softmax * math.log2(math.e) - - # Create tensors (shapes don't matter much, we just need the layouts) - q = torch.randn(m, head_dim, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(s_k, head_dim, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, head_dim, dtype=torch.bfloat16, device='cuda') - c = torch.zeros(m, head_dim, 1, dtype=torch.bfloat16, device='cuda') - - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - # Create CuTe tensors - 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)) - - # V layout: (pv_n_tile, s_k, 1) for FMHA - v_fmha = v[:, 0:pv_n_tile].contiguous() - v_kernel = v_fmha.unsqueeze(-1) - mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel)) - mC = ct.from_dlpack(c[:, 0:pv_n_tile, :]).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c[:, 0:pv_n_tile, :])) - - # MMA setup - a_major = cute.LayoutEnum.from_tensor(mQ).mma_major_mode() - b_major = cute.LayoutEnum.from_tensor(mK).mma_major_mode() - v_major = cute.LayoutEnum.from_tensor(mV).mma_major_mode() - - qk_mma = utils.sm100.make_trivial_tiled_mma( - BFloat16, BFloat16, a_major, b_major, Float32, - tcgen05.CtaGroup.ONE, (128, 128), tcgen05.OperandSource.SMEM - ) - pv_a_major = a_major # SMEM-P: A from SMEM - pv_source = tcgen05.OperandSource.SMEM - pv_mma = utils.sm100.make_trivial_tiled_mma( - BFloat16, BFloat16, pv_a_major, v_major, Float32, - tcgen05.CtaGroup.ONE, (128, pv_n_tile), pv_source - ) - - # Compute layouts (mirrors FmhaKernel._setup) - 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, pv_n_tile, pv_ik * (128 // pv_ik)) - - # P SMEM layout (PV A-operand, 1 stage) - p_smem_s = utils.sm100.make_smem_layout_a(pv_mma, pv_mma_tiler, BFloat16, 1) - - # P TMEM layout (for reference) - p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, pv_mma_tiler, BFloat16, 1) - - # QK C-fragment - qk_thr = qk_mma.get_slice(0) - qk_as = qk_thr.partition_shape_C(qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - tStS0 = cute.make_tensor(tStS.iterator, tStS.layout) # offset 0 for diagnostics - - # PV C-fragment - pv_thr = pv_mma.get_slice(0) - pv_as = pv_thr.partition_shape_C(pv_mma_tiler[:2]) - tOtO = pv_thr.make_fragment_C(pv_as) - - # ====== TMEM-load copy (softmax reads S from TMEM) ====== - 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, tStS0) - - print("=== TMEM-Load Copy Info ===") - print(f" layout_tv_tiled: {tiled_tmem_load.layout_tv_tiled}") - print(f" layout_dst_tv_tiled: {tiled_tmem_load.layout_dst_tv_tiled}") - print(f" layout_src_tv_tiled: {tiled_tmem_load.layout_src_tv_tiled}") - print(f" TV shape: {tiled_tmem_load.layout_tv_tiled.shape}") - print(f" TV stride: {tiled_tmem_load.layout_tv_tiled.stride}") - - # Softmax thread partition - sfw_idx = 0 # thread 0 in softmax warps - thr_load = tiled_tmem_load.get_slice(sfw_idx) - tTMEM_LOADtS = thr_load.partition_S(tStS0) - - # Coordinate identity tensor - cS = cute.make_identity_tensor((128, 128)) - tScS = qk_thr.partition_C(cS) - tTMEM_LOADcS = thr_load.partition_D(tScS) - - print(f"\n=== Per-Thread Partition (thread 0) ===") - print(f" tTMEM_LOADtS shape: {cute.shape(tTMEM_LOADtS)}") - print(f" tTMEM_LOADtS layout: {tTMEM_LOADtS.layout}") - print(f" tTMEM_LOADcS shape: {cute.shape(tTMEM_LOADcS)}") - print(f" tTMEM_LOADcS layout: {tTMEM_LOADcS.layout}") - - # ====== sP layout (PV A-operand SMEM) ====== - # Remove stage dim - sP_stage = p_smem_s.outer - if len(cute.shape(sP_stage)) > 3: - # Try removing the stage dimension - sP_stage = sP_stage # might need slicing - - print(f"\n=== sP Layout (PV A-operand SMEM) ===") - print(f" p_smem_s.outer shape: {cute.shape(p_smem_s.outer)}") - print(f" p_smem_s.outer layout: {p_smem_s.outer}") - print(f" p_smem_s.inner (swizzle): {p_smem_s.inner}") - - # Flatten sP to 1D for address computation - sP_flat = cute.coalesce(p_smem_s.outer) - print(f" sP_flat shape: {cute.shape(sP_flat)}") - print(f" sP_flat layout: {sP_flat}") - - # ====== Build atom_layout_tv ====== - # The TV layout of tiled_tmem_load maps (tid, vid) → tStS0 coordinate - # We need to compose with (m, k) → sP address mapping - # - # Approach: - # 1. Get the TV layout: (tid, vid) → flat S/P address in tStS0 - # 2. Convert tStS0 flat address to (m, k) coordinate - # 3. Convert (m, k) to sP coordinate - # 4. Convert sP coordinate to sP flat address - # - # The TV layout from tiled_tmem_load is in terms of tStS0's data layout. - # tStS0 has layout ((128,128),1,1):((65536,1),0,0) - # So the flat address IS (m * 65536 + k), which is just m*65536 + k - # But tStS0's actual coordinate mapping is (m, k) where m is row, k is column - - # Let's look at what layout_dst_tv_tiled gives us - dst_tv = tiled_tmem_load.layout_dst_tv_tiled - print(f"\n=== Destination TV Layout ===") - print(f" shape: {dst_tv.shape}") - print(f" stride: {dst_tv.stride}") - print(f" size: {cute.size(dst_tv)}") - - # The destination TV layout maps (tid, vid) to the destination coordinate space. - # For the TMEM-load, destination is tStS0. - # tStS0 layout: ((128,128),1,1):((65536,1),0,0) - # So (tid, vid) → (addr_m * 65536 + addr_k * 1, 0, 0) - # The first component encodes (m, k) in the 128x128 S matrix - - # We need to understand how (tid, vid) maps to (m, k) - # For a 128-thread partition (4 softmax warps × 32 threads/warp): - # Total values: 128 * 128 = 16384 - - num_softmax_threads = 128 # 4 warps - # Each thread has tTMEM_LOADcS size = 32 * 4 = 128 values (from shape ((32,1),4,1,1)) - - # The TV layout should have shape (num_threads, values_per_thread) - # From the TMEM-load TV layout: - tv_shape = dst_tv.shape - tv_stride = dst_tv.stride - print(f"\n TV shape breakdown: tid_dim={tv_shape[0] if tv_shape else '?'}, vid_dim={tv_shape[1] if len(tv_shape) > 1 else '?'}") - - # Now let's try to build the sP address mapping - # sP outer layout (from diagnostics): ((128,16),1,(4,2),1):(((64,1),0,(16,8192)),0) - # With swizzle S<3,4,3> - - # For make_cotiled_copy, we need: - # atom_layout_tv: (tid, vid) → sP flat address - # data_layout: sP flat layout (coalesced) - - # Step 1: Extract (m, k) from dst_tv - # dst_tv maps (tid, vid) → tStS0 address - # tStS0 address = m * 65536 + k (from layout ((128,128),1,1):((65536,1),0,0)) - # So m = addr // 65536, k = addr % 65536 - # But since tStS0 has stride (65536, 1) for the first group ((128, 128)): - # addr = m * 65536 + k - # This gives us the (m, k) mapping implicitly. - - # Step 2: Map (m, k) to sP address - # sP coordinate: ((m, k%16), 0, ((k//16)%4, k//64), 0) - # sP layout applies swizzle and stride - # sP flat address = sP.layout(((m, k%16), 0, ((k//16)%4, k//64), 0)) - - # For make_cotiled_copy, we need: - # atom_layout_tv: (tid, vid) → sP flat address - # We can compute this by composing the TV layout with a mapping from - # tStS0 address to sP address. - - # Let me think about this differently. - # We know dst_tv maps (tid, vid) → tStS0 address (a flat integer) - # tStS0 address encodes (m, k) as m*65536 + k - # We need to convert this to an sP address - - # Actually, let's use the identity tensor approach. - # cS = make_identity_tensor((128, 128)) maps (m, k) → (m, k) - # tScS = qk_thr.partition_C(cS) maps thread's C-fragment coordinates to (m, k) - # tTMEM_LOADcS = thr_load.partition_D(tScS) maps TMEM-load partition coordinates to (m, k) - - # The key: tTMEM_LOADcS already gives us (m, k) per (thread_coord, vid_coord) - # We just need to evaluate ALL (thread, vid) pairs to get the full TV → (m, k) mapping - # Then compose with (m, k) → sP address - - # Let's print some coordinate values to verify - print(f"\n=== Coordinate Verification ===") - # For thread 0: tTMEM_LOADcS shape is ((32,1),4,1,1) - # Values should be (m, k) pairs covering some subset of [0,128) x [0,128) - # Let's iterate and print a few - - # We need to do this at trace time inside @cute.kernel - # For now, print the layout properties - print(f" tScS shape: {cute.shape(tScS)}") - print(f" tScS layout: {tScS.layout}") - print(f" tTMEM_LOADcS shape: {cute.shape(tTMEM_LOADcS)}") - print(f" tTMEM_LOADcS layout: {tTMEM_LOADcS.layout}") - - # Try to understand the TV mapping from the layout properties - # The tTMEM_LOADcS layout: ((32,1),4,1,1):((1@1,0),32@1,0,0) - # This means: for index ((j0,0), j1, 0, 0): - # j0 ranges 0..31, j1 ranges 0..3 - # coordinate = (j0 * 1, j1 * 32) = (j0, j1 * 32) - # But wait, the @1 in stride means there's a tmem-specific indexing - # Actually, the stride is ((1@1, 0), 32@1, 0, 0) - # @1 might mean "stride 1 in the second dimension of a 2D coordinate" - - # Let me look at it from the TMEM load's perspective: - # Ld32x32bOp with Repetition(32) loads a 32x32 tile of FP32 from TMEM - # Each warp has 32 threads. 4 warps = 128 threads. - # Each thread loads 32 values per invocation, repeated 32 times? - # No - Repetition(32) means 32 repetitions of the base operation - # Ld32x32bOp base: 32 threads each load 32 values from a 32x32 TMEM tile - # Repetition(32): each thread loads 32 * 32 = 1024 values total? - # That's too many. With 128 threads, that's 128 * 1024 = 131072 >> 16384 - - # Hmm, let me reconsider. The partition_S/D shapes tell the story: - # tTMEM_LOADtS shape: (((32,32),1),4,1,1) - source (TMEM data) - # tTMEM_LOADcS shape: ((32,1),4,1,1) - destination (coordinate) - # - # The (32,1) in cS vs (32,32) in tS: - # The coordinate is 2D (m,k) so (32,1) means 32 coordinates - # The data is FP32 so (32,32) means 32*32 = 1024 values? No that's too many - # - # Actually, looking at the shapes: - # tTMEM_LOADtS: (((32,32),1),4,1,1) - this is the data register layout - # Inner (32,32) = 1024 values per fragment, 4 fragments = 4096 per thread - # 128 threads * 4096 = 524288 ≠ 16384 - # - # This doesn't add up. The shape must represent the register layout differently. - # Let me just count: 32*1*4*1*1 = 128 elements per thread (coordinate) - # And 32*32*1*4*1*1 = 4096 elements per thread (data) - # 128 * 128 = 16384 coordinates ✓ (each coord is 2D = 32768 scalars) - # But 128 * 4096 = 524288 data elements ≠ 16384 - # - # Something's off. The tTMEM_LOADtS data shape includes the FP32 register layout, - # not element count. The actual element count per thread matches the coordinate count - # because each (m,k) coordinate corresponds to one P value. - - print(f"\n=== Element Count Verification ===") - cs_size = cute.size(tTMEM_LOADcS) # should be 128 (coordinates per thread) - ts_size = cute.size(tTMEM_LOADtS) # data elements per thread - print(f" tTMEM_LOADcS size: {cs_size} (coordinates per thread)") - print(f" tTMEM_LOADtS size: {ts_size} (data elements per thread)") - print(f" 128 threads * {cs_size} coords = {128 * cs_size} (should be 16384)") - - # ====== Try make_cotiled_copy approach ====== - # We need atom_layout_tv: (tid, vid) → sP flat address - # The TMEM-load TV layout gives us (tid, vid) → tStS0 address - # tStS0 address encodes (m, k) - # We need to compose with a (m, k) → sP address mapping - - # Step 1: Get the TV layout that maps to the S/P matrix coordinates - # dst_tv = tiled_tmem_load.layout_dst_tv_tiled - # This maps (tid, vid) → coordinate in tStS0 - # But tStS0 has a specific layout ((128,128),1,1):((65536,1),0,0) - # So the TV output is in terms of tStS0's flat address space - - # Step 2: Build a layout that maps tStS0 flat address → sP flat address - # tStS0 address = m * 65536 + k (from stride) - # We need to map this to the sP address - # sP address = sP.layout(m, k%16, (k//16)%4, k//64) (after swizzle) - - # This requires iterating over all (m, k) and computing sP addresses - # which we can do at Python time since both layouts are static - - # Let's compute the mapping table - sP_outer = p_smem_s.outer - sP_swizzle = p_smem_s.inner - - # Coalesce sP to get a flat 1D layout for address computation - sP_coalesced = cute.coalesce(sP_outer) - print(f"\n=== sP Coalesced Layout ===") - print(f" shape: {cute.shape(sP_coalesced)}") - print(f" layout: {sP_coalesced}") - - # Compute (m, k) → sP address mapping for all 16384 elements - # For verification: build the address table - sP_shape = cute.shape(sP_outer) - print(f"\n=== Full sP shape: {sP_shape} ===") - - # Let's try a different approach: use composition directly - # We know that: - # 1. The TMEM-load TV layout maps (tid, vid) → tStS0 coordinate - # 2. tStS0's layout has shape (128, 128) in the first group - # 3. We need (128, 128) → sP address - - # Build a layout that maps the 128x128 P matrix to sP addresses - # For each (m, k) in [0,128) x [0,128): - # sP index = ((m, k%16), 0, ((k//16)%4, k//64), 0) - # sP address = sP_outer(sP index) (then swizzle applied by CuTe) - - # We can build this as a CuTe layout: - # shape = (128, 128), mapping to sP addresses - # For k in [0, 128): - # k0 = k % 16 - # k1 = (k // 16) % 4 - # k2 = k // 64 - # sP_addr = sP_outer((m, k0), 0, (k1, k2), 0) - - # But CuTe layouts are affine (linear). The (k%16, (k//16)%4, k//64) decomposition - # is NOT affine unless 128 = 16 * 4 * 2 exactly (which it is!) - # k = k0 + 16*k1 + 64*k2 where k0∈[0,16), k1∈[0,4), k2∈[0,2) - # So k0 = k % 16, k1 = (k//16) % 4, k2 = k//64 - # And the sP shape ((128, 16), 1, (4, 2), 1) has: - # mode 0: (128, 16) → stride (64, 1) - # mode 2: (4, 2) → stride (16, 8192) - # So sP_addr(m, k0, k1, k2) = m*64 + k0*1 + k1*16 + k2*8192 - - # Wait, but sP has swizzle S<3,4,3> which XORs some bits - # The swizzle is applied by CuTe when we use the tensor directly - # For make_cotiled_copy, we need the unswizzled (logical) layout - # because CuTe handles the swizzle at access time - - # Actually, let me re-read make_cotiled_copy's requirement: - # atom_layout_tv: (tid, vid) -> data addr - # data_layout: data coord -> data addr - # - # The "data addr" here is the LOGICAL address (before swizzle) - # because CuTe applies swizzle at tensor access time - # So we need the sP outer layout WITHOUT swizzle for data_layout - # But p_smem_s.outer IS the logical layout (swizzle is in p_smem_s.inner) - - # So for make_cotiled_copy: - # data_layout = sP_outer (logical, no swizzle) - # atom_layout_tv: (tid, vid) -> sP logical address - - # The sP logical layout with shape ((128,16),1,(4,2),1) and strides - # can be coalesced to a 1D layout mapping (m, k) → address - - # Let me compute the sP logical layout as a (128, 128) → address mapping - # sP_outer((m, k0), 0, (k1, k2), 0) = m*64 + k0*1 + k1*16 + k2*8192 - - # If we view this as a function of (m, k) where k = k0 + 16*k1 + 64*k2: - # sP_addr(m, k) = m*64 + (k%16)*1 + ((k//16)%4)*16 + (k//64)*8192 - # = m*64 + (k%16) + 16*((k//16)%4) + 8192*(k//64) - - # This is NOT a simple affine function of k alone (due to the modular decomposition) - # But it IS affine if we treat k as decomposed into (k0, k1, k2) - # k = k0 + 16*k1 + 64*k2 - # sP_addr = 64*m + 1*k0 + 16*k1 + 8192*k2 - - # So the sP logical layout, viewed as a (128, 128) → address map, - # has stride (64, ?) where the k-stride varies based on k's position - # This is the standard column-major with subtiling - - # For make_cotiled_copy, we need to express this as a CuTe layout - # The sP outer layout in its native form IS the correct data_layout - - print("\n=== Plan for make_cotiled_copy ===") - print("1. data_layout = coalesce(sP_outer) — flat 1D sP address space") - print("2. atom_layout_tv = compose(tiled_tmem_load.layout_dst_tv_tiled, tStS_to_sP_mapping)") - print("3. tStS_to_sP_mapping: tStS address → sP address") - print(" tStS layout: (128,128) → address, stride=(65536,1)") - print(" sP layout: (128,16,4,2) → address (after coalescing), strides vary") - print("") - print("The problem: tStS has stride 65536 for m, while sP has stride 64 for m") - print("These are fundamentally different address spaces.") - print("We need to: (tid,vid) → (m,k) via identity, then (m,k) → sP_addr via sP layout") - - # Actually, the cleaner approach: - # 1. Build identity tensor for (128, 128) P matrix - # 2. The TMEM-load partition + coordinate partition gives us (tid, vid) → (m, k) - # 3. The sP layout gives us (m, k) → sP address - # 4. Compose these to get (tid, vid) → sP address - - # The identity tensor cS = make_identity_tensor((128, 128)) - # maps (m, k) → (m, k) with layout (128, 128):((1, 128)) — row-major - # Wait, identity tensor layout is just (128, 128) with stride (1, 128)? Let me check. - - # Actually, make_identity_tensor creates a tensor where each element holds its coordinate - # The layout is just an enumerate layout: (128, 128):(1, 128) in row-major - # So cS(m, k) = m + 128*k (flat index in row-major order) - - # The tScS = qk_thr.partition_C(cS) partitions the identity by QK C-fragment threads - # The tTMEM_LOADcS = thr_load.partition_D(tScS) further partitions by TMEM-load threads - - # So tTMEM_LOADcS maps (thread_local_index) → (m, k) coordinate - # The shape is ((32,1),4,1,1) with layout ((1@1,0),32@1,0,0) - # For index ((j0,0), j1, 0, 0): - # value = (j0, j1*32) — thread 0's coordinates - - # But we need the FULL (tid, vid) → (m, k) mapping for ALL 128 threads - # The tTMEM_LOADcS is for a SINGLE thread (thread 0) - - # The TV layout (layout_dst_tv_tiled) gives us the FULL mapping: - # (tid, vid) → destination coordinate (in tStS0's address space) - # But the destination is tStS0 which has layout ((128,128),1,1):((65536,1),0,0) - # So the TV output is in tStS0's flat address: m*65536 + k - - # To convert to sP address, we need to: - # 1. Extract (m, k) from tStS0 flat address: m = addr // 65536, k = addr % 65536 - # But this isn't simple because CuTe layouts aren't easily decomposable - # 2. Map (m, k) to sP address: use sP_outer((m, k%16), 0, ((k//16)%4, k//64), 0) - - # Let's try using composition directly. - # We have: - # dst_tv: (tid, vid) → tStS0_addr - # tStS0 layout: ((128,128),1,1):((65536,1),0,0) - # We need: (tid, vid) → sP_addr - - # Approach: build a "reindex" layout that maps tStS0_addr → sP_addr - # For each possible tStS0_addr, compute the corresponding sP_addr - - # But tStS0_addr = m*65536 + k, and we need to decompose back to (m, k) - # With m ∈ [0, 128), k ∈ [0, 128), and stride 65536 for m and 1 for k - # We can build the inverse: given addr, m = addr // 65536, k = addr % 65536 - # Then sP_addr = 64*m + (k%16) + 16*((k//16)%4) + 8192*(k//64) - - # This is a 16384-element lookup table. We can build it at Python time. - - # Let me build the lookup table - print(f"\n=== Building (tid, vid) → sP_addr mapping ===") - - # Get the TV layout - tv = dst_tv - tv_shape = tv.shape - tv_stride = tv.stride - - # We need to evaluate the TV layout for all (tid, vid) pairs - # and then map the resulting tStS0 address to sP address - - # First, let's understand the TV layout structure - # It should have shape (num_threads, values_per_thread) - # num_threads = 128 (4 softmax warps) - # values_per_thread = 128 (32*4 from the coordinate shape) - - print(f" TV shape: {tv_shape}") - print(f" TV stride: {tv_stride}") - print(f" Total elements: {cute.size(tv)} (should be 16384)") - - # Evaluate the TV layout to get all (tid, vid) → tStS0_addr values - # We can do this by iterating over (tid, vid) and computing the layout value - - # CuTe layout evaluation: layout(idx) = sum(idx_i * stride_i) - # For a 2D layout (tid, vid) → addr: - # addr = tid * stride[0] + vid * stride[1] (simplified) - - # But the TV layout might have nested shapes/strides - # Let me just try to evaluate it - - # Actually, for make_cotiled_copy, I don't need to manually compute the mapping. - # I can use CuTe's composition operation: - # atom_layout_tv = composition(sP_flat_layout, dst_tv) - # This would compose dst_tv: (tid, vid) → tStS0_addr - # with a mapping tStS0_addr → sP_addr - - # But I need to build the tStS0_addr → sP_addr mapping first - # Let me build it as a CuTe layout - - # tStS0 has layout ((128,128),1,1):((65536,1),0,0) - # Coalescing: (128,128):(65536,1) — so tStS0_addr = m*65536 + k - # The inverse is: m = addr // 65536, k = addr % 65536 - - # sP has layout ((128,16),1,(4,2),1):((64,1),0,(16,8192),0) (before swizzle) - # Coalescing: ((128,16,4,2),1,1,1):((64,1,16,8192),0,0,0) → flat - # sP_addr = 64*m + 1*k0 + 16*k1 + 8192*k2 - # = 64*m + (k%16) + 16*((k//16)%4) + 8192*(k//64) - - # To build tStS0_addr → sP_addr as a CuTe layout: - # We need a layout with shape matching tStS0's domain (16384 elements) - # But CuTe layouts are affine, and the mapping is NOT affine - # (due to the modular decomposition of k) - - # So we CAN'T use composition directly. We need to build a custom layout. - - # Alternative: build the sP layout as a (128, 128) → addr layout directly - # sP_addr(m, k) = 64*m + (k%16) + 16*((k//16)%4) + 8192*(k//64) - # This can be expressed as: - # sP_addr = 64*m + f(k) where f(k) = (k%16) + 16*((k//16)%4) + 8192*(k//64) - - # f(k) IS representable as a CuTe layout because 16*4*2 = 128 - # k can be decomposed as (k0, k1, k2) where k0=k%16, k1=(k//16)%4, k2=k//64 - # f(k) = k0 + 16*k1 + 8192*k2 = (k0, k1, k2) → (1, 16, 8192) → addr - - # So the full sP layout in (m, k) coordinates: - # (m, k) → (m, k0, k1, k2) → (64*m + k0 + 16*k1 + 8192*k2) - # = (m, k0, k1, k2) → (64, 1, 16, 8192) → sP_addr - - # And the tStS0 layout: - # (m, k) → (65536*m + k) → tStS0_addr - - # The mapping tStS0_addr → (m, k) is the INVERSE of the tStS0 layout: - # (65536, 1) → (m, k) — this is what left_inverse gives - - # Then (m, k) → sP_addr via the sP layout - - # So the full chain \ No newline at end of file diff --git a/tests/unit/test_smem_p_write.py b/tests/unit/test_smem_p_write.py deleted file mode 100644 index d1982548..00000000 --- a/tests/unit/test_smem_p_write.py +++ /dev/null @@ -1,199 +0,0 @@ -""" -SMEM-P Write Diagnostic: Verify the coordinate-indexed write to sP. - -This test: -1. Creates an S matrix (identity for simplicity) -2. Writes S values to sP using the coordinate-indexed approach -3. Reads sP back via the PV MMA's A-operand fragment -4. Verifies the round-trip - -If the round-trip is correct, the coordinate extraction and sP indexing are right. -If not, we need to debug the coordinate mapping. -""" -import torch, math -import cutlass, cutlass.cute as cute -import cutlass.utils as utils -from cutlass.cute.nvgpu import tcgen05, cpasync -from cutlass import Float32, BFloat16, Int32, const_expr -from cutlass.utils import LayoutEnum -import cutlass.torch as ct -import cuda.bindings.driver as cuda - - -@cute.jit -def smem_p_write_test(s_data, sP_out, qk_mma, pv_mma, qk_mma_tiler, pv_mma_tiler, p_smem_s): - """Write S values to sP and read them back.""" - tidx, _, _ = cute.arch.thread_idx() - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - - # Only use 4 softmax warps + 1 MMA warp for this test - if warp_idx >= 5: - return - - # SMEM allocation - @cute.struct - class SS: - dummy: cute.struct.MemRange[cutlass.Int64, 2] - smem = utils.SmemAllocator() - st = smem.allocate(SS) - tmem_bar = pipeline.NamedBarrier(barrier_id=2, num_threads=32 * 5) # 5 warps - - # Allocate sP in SMEM - sP = smem.allocate_tensor(element_type=BFloat16, layout=p_smem_s.outer, byte_alignment=128, swizzle=p_smem_s.inner) - sP_nostage = sP[(None, None, None, 0)] - - # Allocate TMEM for S - tmem = utils.TmemAllocator(st.dummy.ptr, barrier_for_retrieve=tmem_bar, allocator_warp_id=0, is_two_cta=False) - if warp_idx < 4: - tmem.allocate(128) # enough for 128×128 FP32 S - tmem.wait_for_alloc() - tmem_ptr = tmem.retrieve_ptr(Float32) - - # QK C-fragment (S in TMEM) - qk_thr = qk_mma.get_slice(0) - qk_as = qk_thr.partition_shape_C(qk_mma_tiler[:2]) - tStS = qk_thr.make_fragment_C(qk_as) - tStS0 = cute.make_tensor(tStS.iterator, tStS.layout) - - # TMEM-load copy - 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, tStS0) - sfw_idx = tidx % 128 # 4 softmax warps - thr_load = tiled_tmem_load.get_slice(sfw_idx) - tTMEM_LOADtS = thr_load.partition_S(tStS0) - - # Coordinate identity tensor - cS = cute.make_identity_tensor((128, 128)) - tScS = qk_thr.partition_C(cS) - tTMEM_LOADcS = thr_load.partition_D(tScS) - - # Softmax warps: read S from TMEM, write to sP - if warp_idx < 4: - # First: copy input data to TMEM (S data) - # Use the TMEM store to write s_data to TMEM - 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, tStS0) - thr_store = tiled_tmem_store.get_slice(sfw_idx) - tTMEM_STOREtS = thr_store.partition_D(tStS0) - - # Copy s_data to register, then to TMEM - tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, Float32) - - # Load from global → register → TMEM - gS = cute.local_tile(s_data, (128, 128), (0, 0)) - tCgS = qk_thr.partition_C(gS) - - # Simple: load from global to register using universal copy - cute.copy(tiled_tmem_store, tTMEM_LOADrS, tTMEM_STOREtS) - cute.arch.fence_view_async_tmem_store() - - # Actually, we need to first put the data into TMEM. - # The simplest approach: just write identity values to TMEM directly. - # For testing, write the coordinate (m + 128*k) as the S value. - # Then verify sP has the same values. - - # Fill TMEM with test values: S[m, k] = m + 128*k (as BF16) - # Use the register bridge pattern - rS_words = cute.make_rmem_tensor(tTMEM_STOREtS.shape, Float32) - rS_bf16 = cute.make_tensor(cute.recast_ptr(rS_words.iterator, dtype=BFloat16), 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)) - rS_bf16_frg = cute.logical_divide(rS_bf16, cute.make_layout(frg_tile)) - - for j in range(frg_cnt): - for k in range(cute.size(tTMEM_LOADrS_frg, mode=[0])): - # Get (m, k) coordinate - coord = tTMEM_LOADcS[(k, 0), j, 0, 0] - m_val = coord[0] - k_val = coord[1] - # S = m + 128*k (unique value per position) - val = Float32(1) * m_val + Float32(128) * k_val - tTMEM_LOADrS_frg[k, j] = val - s_vec = tTMEM_LOADrS_frg[None, j].load() - rS_bf16_frg[None, j].store(s_vec.to(BFloat16)) - - # Store to TMEM - cute.copy(tiled_tmem_store, rS_words, tTMEM_STOREtS) - cute.arch.fence_view_async_tmem_store() - - # Now read S from TMEM and write to sP - cute.copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS) - cute.arch.fence_view_async_tmem_load() - - # Write to sP using coordinate-indexed store - rP_bf16_test = cute.make_tensor(cute.recast_ptr(rS_words.iterator, dtype=BFloat16), tTMEM_LOADrS.layout) - # Copy the BF16 values from S load to P buffer - for j in range(frg_cnt): - s_vec = tTMEM_LOADrS_frg[None, j].load() - rP_bf16_test_frg = cute.logical_divide(rP_bf16_test, cute.make_layout(frg_tile)) - rP_bf16_test_frg[None, j].store(s_vec.to(BFloat16)) - - for j0 in range(32): - for j1 in range(4): - coord = tTMEM_LOADcS[(j0, 0), j1, 0, 0] - m_coord = coord[0] - k_coord = coord[1] - k0 = k_coord % 16 - k1 = (k_coord // 16) % 4 - k2 = k_coord // 64 - sP_nostage[(m_coord, k0), 0, (k1, k2)] = rP_bf16_test[(j0, 0), j1, 0, 0] - - cute.arch.fence_proxy("async.shared", space="cta") - - # Signal MMA warp that sP is ready - sync_bar = pipeline.NamedBarrier(barrier_id=3, num_threads=32 + 32) - sync_bar.arrive() - - # MMA warp: read from sP and write to output - if warp_idx == 4: - sync_bar.arrive_and_wait() - - # Read sP using PV MMA's A-operand fragment - pv_thr = pv_mma.get_slice(0) - tCrP = pv_mma.make_fragment_A(sP) - - # Read sP values and write to output - gOut = cute.local_tile(sP_out, (128, 128), (0, 0)) - tCgOut = pv_thr.partition_C(gOut) - - # Simple: read from sP using the fragment and store to output - # We need to iterate over the fragment's K dimension - for kb in cutlass.range(cute.size(tCrP, mode=[2]), unroll_full=True): - for i in cutlass.range(cute.size(tCrP, mode=[0]), vectorize=True): - pass # Can't easily extract individual values from SMEM fragment - - # Alternative: read sP directly using the sP tensor (not fragment) - # This tests if the sP writes are correct - for m in range(128): - for k0 in range(16): - for k1 in range(4): - for k2 in range(2): - val = sP_nostage[(m, k0), 0, (k1, k2)] - # Expected: m + 128*(k0 + 16*k1 + 64*k2) = m + 128*k - # But we can't write to global memory from here easily - - if warp_idx < 4: - tmem.relinquish_alloc_permit() - tmem.free(tmem_ptr) - - -def main(): - head_dim = 256 - s_k = 128 - - # This is too complex for a first test. Let me do something simpler: - # Just verify that the coordinate mapping (j0, j1) -> (m, k) is correct - # by printing coordinates from the identity tensor. - - # Actually, we can't print from inside @cute.kernel easily. - # Let me try a different approach: create a simple test that - # writes known values to sP using the coordinate approach, - # then reads them back and checks correctness. - - # Simplest possible test: write to sP in host code and read in kernel - pass - -if __name__ == '__main__': - main() diff --git a/tests/unit/test_tmem_budget.py b/tests/unit/test_tmem_budget.py deleted file mode 100644 index b380cbc1..00000000 --- a/tests/unit/test_tmem_budget.py +++ /dev/null @@ -1,100 +0,0 @@ -"""TMEM column budget probe for FMHA at various head_dims. - -Uses @cute.jit to construct MMA objects inside a compiled context. -Only prints the critical budget numbers. -""" -import torch, math -import cuda.bindings.driver as cuda -import 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 -from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -import cutlass.torch as ct - - -class BudgetProbe: - def __init__(self, head_dim): - self.hd = head_dim - - @cute.jit - def __call__(self, mQ, mK, stream): - a_major = LayoutEnum.from_tensor(mQ).mma_major_mode() - b_major = LayoutEnum.from_tensor(mK).mma_major_mode() - - # QK MMA: always (128, 128) - qk_mma = utils.sm100.make_trivial_tiled_mma( - BFloat16, BFloat16, a_major, b_major, - Float32, tcgen05.CtaGroup.ONE, (128, 128), - tcgen05.OperandSource.SMEM, - ) - qk_thr = qk_mma.get_slice(0) - qk_as = qk_thr.partition_shape_C((128, 128)) - tStS = qk_thr.make_fragment_C(qk_as) - s_cols = find_tmem_tensor_col_offset(tStS) - cute.printf("hd=%d s_cols=%d", Int32(self.hd), Int32(s_cols)) - - # PV MMA: (128, hd), TMEM-P (P from TMEM, K-major) - # tcgen05 MMA max N = 256. At hd > 256, skip TMEM-P and SMEM-P full-hd probes. - pv_n_tile = self.hd if self.hd <= 256 else 256 - pv_a_major_tmem = cute.nvgpu.OperandMajorMode.K - pv_mma_tmem = utils.sm100.make_trivial_tiled_mma( - BFloat16, BFloat16, pv_a_major_tmem, b_major, - Float32, tcgen05.CtaGroup.ONE, (128, pv_n_tile), - tcgen05.OperandSource.TMEM, - ) - pv_thr_tmem = pv_mma_tmem.get_slice(0) - pv_as_tmem = pv_thr_tmem.partition_shape_C((128, pv_n_tile)) - tOtO_tmem = pv_thr_tmem.make_fragment_C(pv_as_tmem) - o_cols_tmem = find_tmem_tensor_col_offset(tOtO_tmem) - cute.printf("hd=%d pv_n_tile=%d o_cols_tmem=%d", Int32(self.hd), Int32(pv_n_tile), Int32(o_cols_tmem)) - - # PV MMA: (128, pv_n_tile), SMEM-P (P from SMEM, same a_major as Q) - pv_mma_smem = utils.sm100.make_trivial_tiled_mma( - BFloat16, BFloat16, a_major, b_major, - Float32, tcgen05.CtaGroup.ONE, (128, pv_n_tile), - tcgen05.OperandSource.SMEM, - ) - pv_thr_smem = pv_mma_smem.get_slice(0) - pv_as_smem = pv_thr_smem.partition_shape_C((128, pv_n_tile)) - tOtO_smem = pv_thr_smem.make_fragment_C(pv_as_smem) - o_cols_smem = find_tmem_tensor_col_offset(tOtO_smem) - cute.printf("hd=%d pv_n_tile=%d o_cols_smem=%d", Int32(self.hd), Int32(pv_n_tile), Int32(o_cols_smem)) - - # P columns in TMEM (TMEM-P path only) - p_cols = 128 * 16 // 32 # pv_mma_tiler[2] * bf16_width / fp32_width - cute.printf("hd=%d p_cols=%d", Int32(self.hd), Int32(p_cols)) - - # TMEM-P total - tmem_p0 = 32 - p_end = tmem_p0 + p_cols - o_after = s_cols if s_cols > p_end else p_end - tmem_o0 = ((o_after + 31) // 32) * 32 - total_tmem_p = tmem_o0 + o_cols_tmem - cute.printf("hd=%d TMEM-P_total=%d", Int32(self.hd), Int32(total_tmem_p)) - - # SMEM-P total (O reuses S space, starting at col 0) - cute.printf("hd=%d SMEM-P_total=%d", Int32(self.hd), Int32(o_cols_smem)) - - -def probe_hd(hd): - print(f"\n=== HEAD_DIM={hd} ===", flush=True) - 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') - - 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)) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - probe = BudgetProbe(head_dim=hd) - print(f' Compiling hd={hd}...', flush=True) - compiled = cute.compile(probe, mQ, mK, stream) - compiled(mQ, mK, stream) - torch.cuda.synchronize() - print(f' Done.', flush=True) - - -if __name__ == '__main__': - for hd in [64, 128, 256, 512]: - probe_hd(hd) diff --git a/tests/unit/test_tmem_roundtrip_minimal.py b/tests/unit/test_tmem_roundtrip_minimal.py deleted file mode 100644 index b0c2a47f..00000000 --- a/tests/unit/test_tmem_roundtrip_minimal.py +++ /dev/null @@ -1,154 +0,0 @@ -""" -D1.5 Phase 2: NO-OP TMEM round-trip test inside FMHA context. - -Strategy: Run FMHA with s_k=128 (single KV tile, no rescale). -Then add a correction_rescale with scale=1.0 (NO-OP) after PV. -If output is bitwise identical to without rescale → round-trip works. -If output differs → round-trip corrupts data. - -This tests the EXACT CUTLASS correction_rescale pattern: - - Both load and store atoms use Repetition(corr_tile_size) - - Both copies built from the SAME tOtO_i tensor (composition-tiled) - - Register buffer sized from load partition_D - -Variants: - V1: Repetition(16) + composition — CUTLASS exact pattern - V2: Repetition(32) + composition -""" -import torch, math -import cutlass.cute as cute -import cutlass.torch as ct -import cuda.bindings.driver as cuda -from dsv4.kernels.attention.fmha import FmhaKernel - - -def reference_attention(q, k, v, scale): - """FP32 reference: returns un-normalized O.""" - qf = q.float() - kf = k.float() - attn = qf @ kf.T * scale - attn_max = attn.max(dim=-1, keepdim=True)[0] - attn_exp = torch.exp(attn - attn_max) - ref_unnorm = attn_exp @ v.float() - return ref_unnorm - - -def test(): - hd = 64 - s_k = 128 - m = 128 - scale = 1.0 / math.sqrt(hd) - torch.manual_seed(42) - - q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') - k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda') - v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') - c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - - # FP32 reference - ref_unnorm = reference_attention(q[:, :, 0], k[:, :, 0], v, scale) - - lse_tensor = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - row_sums_tensor = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - # Test 1: Baseline s_k=128, no rescale, TMEM-P - kernel1 = FmhaKernel(head_dim=hd, s_k=s_k, use_smem_p=False, normalize=False) - pv_n_tile = kernel1.pv_n_tile - v_tile = v[:, 0:pv_n_tile].contiguous() - v_kernel = v_tile.unsqueeze(-1) - c_tile1 = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - lse1 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - rs1 = torch.zeros(m, 1, 1, dtype=torch.float32, 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)) - mC1 = ct.from_dlpack(c_tile1).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile1)) - mLSE1 = ct.from_dlpack(lse1).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse1)) - mRS1 = ct.from_dlpack(rs1).mark_layout_dynamic(leading_dim=ct.get_leading_dim(rs1)) - - print(f'Test 1: s_k=128 baseline (no rescale)', flush=True) - compiled1 = cute.compile(kernel1, mQ, mK, mV, mC1, stream, mLSE1, row_sums=mRS1) - compiled1(mQ, mK, mV, mC1, stream, mLSE1, row_sums=mRS1) - torch.cuda.synchronize() - - out1 = c_tile1[:, :, 0].float() - cos1 = torch.nn.functional.cosine_similarity( - out1.flatten().unsqueeze(0), ref_unnorm.flatten().unsqueeze(0) - ).item() - print(f' cos_unnorm={cos1:.6f} {"PASS" if cos1 >= 0.999 else "FAIL"}') - - # Test 2: s_k=256 with the CUTLASS correction_rescale pattern - # This is the REAL test — does multi-KV-tile O rescale work? - s_k2 = 256 - k2 = torch.randn(s_k2, hd, 1, dtype=torch.bfloat16, device='cuda') - v2 = torch.randn(s_k2, hd, dtype=torch.bfloat16, device='cuda') - c_tile2 = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - lse2 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - rs2 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - - ref_unnorm2 = reference_attention(q[:, :, 0], k2[:, :, 0], v2, scale) - - # Use the EXISTING Python KV merge as oracle - # Run per-segment and merge - kernel_s128 = FmhaKernel(head_dim=hd, s_k=128, use_smem_p=False, normalize=False) - - # Segment 0 - c_seg0 = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - lse_seg0 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - rs_seg0 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - mK0 = ct.from_dlpack(k2[:128]).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k2[:128])) - v2_t0 = v2[:128, 0:pv_n_tile].contiguous().unsqueeze(-1) - mV0 = ct.from_dlpack(v2_t0).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v2_t0)) - mC_s0 = ct.from_dlpack(c_seg0).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_seg0)) - mLSE_s0 = ct.from_dlpack(lse_seg0).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_seg0)) - mRS_s0 = ct.from_dlpack(rs_seg0).mark_layout_dynamic(leading_dim=ct.get_leading_dim(rs_seg0)) - compiled_s0 = cute.compile(kernel_s128, mQ, mK0, mV0, mC_s0, stream, mLSE_s0, row_sums=mRS_s0) - compiled_s0(mQ, mK0, mV0, mC_s0, stream, mLSE_s0, row_sums=mRS_s0) - - # Segment 1 - c_seg1 = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - lse_seg1 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - rs_seg1 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - mK1 = ct.from_dlpack(k2[128:]).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k2[128:])) - v2_t1 = v2[128:, 0:pv_n_tile].contiguous().unsqueeze(-1) - mV1 = ct.from_dlpack(v2_t1).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v2_t1)) - mC_s1 = ct.from_dlpack(c_seg1).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_seg1)) - mLSE_s1 = ct.from_dlpack(lse_seg1).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_seg1)) - mRS_s1 = ct.from_dlpack(rs_seg1).mark_layout_dynamic(leading_dim=ct.get_leading_dim(rs_seg1)) - compiled_s1 = cute.compile(kernel_s128, mQ, mK1, mV1, mC_s1, stream, mLSE_s1, row_sums=mRS_s1) - compiled_s1(mQ, mK1, mV1, mC_s1, stream, mLSE_s1, row_sums=mRS_s1) - - torch.cuda.synchronize() - - # Python KV merge (proven correct, cos 0.999998) - o0 = c_seg0[:, :, 0].float() - o1 = c_seg1[:, :, 0].float() - rs0 = rs_seg0[:, 0, 0].float() - rs1_val = rs_seg1[:, 0, 0].float() - lse0 = lse_seg0[:, 0, 0].float() - lse1_val = lse_seg1[:, 0, 0].float() - - # D5 merge: O = sum(exp(lse_i) * O_i) / sum(exp(lse_i)) - # where O_i is NORMALIZED (O_i = O_unnorm_i / row_sum_i) - o0_norm = o0 / rs0.unsqueeze(1).clamp(min=1e-30) - o1_norm = o1 / rs1_val.unsqueeze(1).clamp(min=1e-30) - w0 = torch.exp(lse0).unsqueeze(1) - w1 = torch.exp(lse1_val).unsqueeze(1) - oracle = (w0 * o0_norm + w1 * o1_norm) / (w0 + w1) - - cos_oracle = torch.nn.functional.cosine_similarity( - oracle.flatten().unsqueeze(0), ref_unnorm2.flatten().unsqueeze(0) - ).item() - print(f'\nOracle (Python KV merge): cos={cos_oracle:.6f}', flush=True) - - # Now test: s_k=256 with in-kernel rescale (if we add it to FmhaKernel) - # This test will FAIL until we implement the fix, but serves as the target - print(f'\nTest 2: s_k=256 in-kernel rescale (NOT YET IMPLEMENTED)', flush=True) - print(f' This test requires adding correction_rescale to FmhaKernel') - print(f' See Phase 4 of MAY_24_2026_PLAN_NEW.md') - - -if __name__ == '__main__': - test() diff --git a/tests/unit/test_ultra_minimal.py b/tests/unit/test_ultra_minimal.py deleted file mode 100644 index 60e8f890..00000000 --- a/tests/unit/test_ultra_minimal.py +++ /dev/null @@ -1,31 +0,0 @@ -"""Ultra-minimal test: Float32 comparison + Int32 assignment in CuTeDSL.""" -import torch -import cutlass -import cutlass.cute as cute -import cutlass.torch as cutlass_torch - - -@cute.kernel -def ultra_minimal_kernel( - input_f32: cute.Tensor, - output_i32: cute.Tensor, -): - tidx, _, _ = cute.arch.thread_idx() - if tidx == cutlass.Int32(0): - x = cutlass.Float32(3.7) # no load, just a constant - r = cutlass.Int32(0) - if x > cutlass.Float32(2.0): - r = cutlass.Int32(1) - cute.arch.store(output_i32.iterator, r) - - -if __name__ == "__main__": - out = torch.zeros(1, dtype=torch.int32, device='cuda') - dummy = torch.zeros(1, dtype=torch.float32, device='cuda') - dc = cutlass_torch.from_dlpack(dummy).mark_layout_dynamic(leading_dim=0) - oc = cutlass_torch.from_dlpack(out).mark_layout_dynamic(leading_dim=0) - print("Compiling...") - compiled = cute.compile(ultra_minimal_kernel, dc, oc) - print("Running...") - compiled(dc, oc) - print(f'Result: {out.item()} (expected 1)') diff --git a/tests/working_softmax_maybe.py b/tests/working_softmax_maybe.py deleted file mode 100644 index f322bd20..00000000 --- a/tests/working_softmax_maybe.py +++ /dev/null @@ -1,497 +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. - A no-op TMEM-load-then-TMEM-store visibly corrupts data. Use the paired - atoms from `utils.sm100.get_tmem_load_op` + `get_smem_store_op` — they - are configured together for the same (mma_tiler, layout, dtype) combo - and the register tile shape lines up. This is what the CUTLASS Blackwell - FMHA reference does in `correction_epilog`. - -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 cutlass.range loop variable. - -2. Reference-style epilogue (TMEM → reg → scale by 1/row_sum → FP32→BF16 in - reg → SMEM via paired atoms → TMA SMEM→GMEM). One pass, no TMEM - round-trip, no `epilogue_tma_store` helper. Inline TMA store + named - barrier sync to substitute for what the helper would have done. - -3. Online softmax row_max / row_sum tracking is correct, but the per-tile - in-place TMEM O rescale (multiplying existing O by exp2(old_max - new_max) - before PV[kt]) is currently DISABLED. Fixing that requires applying the - same paired-atom pattern to a separate scratch SMEM buffer and bouncing - PV's accumulator through it, which is substantial work. For now, the - kernel is correct when row_max growth across tiles is mild. Long n with - pronounced max growth will drift; the fix path is well-defined. - -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() - # 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 ===== - # NOTE: using kt from cutlass.range works for n=128 (single tile). - # Multi-tile (n>128) loads from tile 0 only — the JIT constant-folds kt. - # TODO: fix multi-tile TMA indexing (kv_coord pattern from 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() - 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) - - row_max = -Float32.inf - row_sum = Float32(0.0) - scale_log2 = Float32(self.scale_softmax_log2) - - # Per-tile softmax loop. - # Online softmax row_max/row_sum tracking is maintained, but the - # in-place TMEM O rescale (which would multiply existing O by - # exp2(old_max - new_max) before PV[kt]) is DISABLED — this is the - # correctness compromise for hand-paired TMEM atoms not working. - # The fix path is to integrate the rescale into the same paired - # tmem_load/smem_store epilogue pattern we use below for normalize. - # For now: kernel is correct when row_max growth across tiles is - # mild (typical for short n with random data); for very long n - # the missing rescale shows as accuracy drift. - 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, fused with scale). - 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) - - # row_sum rescale (correct even without O rescale — row_sum - # is a register variable, not in TMEM). - # 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, - # 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() - - si_handle.release() - softmax_done_bar.arrive() - - # === Reference-style scaled epilogue (no TMEM round-trip) === - # - # Pattern (mirrors CUTLASS Blackwell FMHA reference's - # correction_epilog): for each column sub-tile, - # 1. TMEM -> registers via PAIRED tmem_load atom - # 2. scale in registers (1/row_sum) - # 3. FP32 -> BF16 conversion in registers - # 4. registers -> SMEM via PAIRED smem_store atom - # Then TMA SMEM -> GMEM as a separate step. - # - # Critical: the load and store atoms MUST be a matched pair. - # Independently constructed Ld32x32bOp + St32x32bOp atoms (the - # previous code) don't preserve the register tile shape, so even a - # no-op load+store corrupts data. Using utils.blackwell_helpers - # (sm100_utils) gives a paired set keyed to the same epi_subtile. - - # Wait for MMA's PV[N-1] to commit before reading O. - final_o_bar.arrive_and_wait() - - # === O normalization via TMEM load → scale → TMEM store === - # Matches CUTLASS reference's correction_rescale pattern exactly. - - 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_atom = cute.make_copy_atom( - tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), - self.acc_dtype, - ) - tmem_store_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_atom, tOtO_i) - tiled_tmem_store_o = tcgen05.make_tmem_copy(tmem_store_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) - - # 2D register tensor: (frg_shape, n_corr_tiles) - tTMrO = cute.make_rmem_tensor( - (tTMEM_LOADcO.shape, 128 // corr_tile_size), self.acc_dtype - ) - - inv_row_sum = Float32(1.0) / row_sum - - for i in range(HEAD_DIM // corr_tile_size): - 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_o, 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(tiled_tmem_store_o, tTMrO_i, 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) - - # 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