Gino Lu 7103eacc99 refactor(sparse_attn): caller-owned workspace + dtype-aware sizing ...

Replace process-lifetime lazy hipMalloc K-stats workspace with a caller-owned
buffer; expose sparge_blockmap_get_workspace_size() / compute_workspace_layout()
host helpers. Split the combined sparge_blockmap_fwd into stage launchers
(sparge_kstats_fwd_oneshot + sparge_blockmap_only_fwd_oneshot) so the chained
launch is timed end-to-end.

Make pooled_k storage dtype follow KDataType (fp16/bf16) instead of fp32 to halve
workspace footprint and match dense-FMHA precision. Tighten per-head superparam
pointers to required (non-null) and assert N_k <= 256 in jenga MakeKargs to
document the 256-bool LDS staging cap. Drop the obsolete VSA extra-LDS staging.

Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>

2026-05-17 02:34:23 -04:00

..

kernel

refactor(sparse_attn): caller-owned workspace + dtype-aware sizing

2026-05-17 02:34:23 -04:00

pipeline

refactor(sparse_attn): caller-owned workspace + dtype-aware sizing

2026-05-17 02:34:23 -04:00