forked from ccf-ai-infra/TileOPs-Metax
2f6dc23c47
## Summary ### da_cumsum Accepts raw `dt` and applies the full pipeline as compile-time-conditional steps: - Add per-head `dt_bias` (`has_dt_bias=True`) - Softplus with overflow bypass at `dt > 20` (`dt_softplus=True`) - Clamp to `[dt_min, dt_max]` - Inclusive prefix sum of `dA = dt_out * A` Returns two outputs: `dt_out` (processed dt) and `dA_cumsum`. The kernel signature is fixed regardless of flags; unused inputs are dummy-zeroed at the op boundary **only when `has_dt_bias=False`** — calling `forward` with `dt_bias=None` when `has_dt_bias=True` now raises `ValueError` immediately instead of silently computing un-biased results. ### ssd_chunk_scan Tensor layouts updated to match the official `_chunk_scan_fwd`: - `x`, `C`, `out` changed from chunk-fused `[B,C,L,H,P]` to seqlen-fused `[B,S,H,P]` - `cb` changed from head-owned `[B,C,H,L,L]` to group-owned `[B,C,G,L,L]` - `prev_states` axis order changed from `[B,C,H,N,P]` to `[B,C,H,P,N]` (P before N, official convention) - `dt` layout changed from `[B,C,L,H]` to `[B,H,C,L]` - `n_groups` added as a constructor parameter `dA_l` shared-memory load moved to just before it is consumed (after the history path), eliminating a redundant `sync_threads` stall. ### ssd_state_passing Fixed output convention to match Mamba-2 spec: `out[:,c]` now holds the state *before* chunk `c`, so `out[:,0] = initial_states` and `out[:,c+1] = s_c` for `c` in `[0, C-2]`. Reference implementations in tests and benchmarks updated to match. ### All Mamba kernels Added `@functools.lru_cache` to all five kernel factory functions (`da_cumsum`, `ssd_chunk_scan`, `ssd_chunk_state`, `ssd_decode`, `ssd_state_passing`) to prevent redundant TileLang recompilation on repeated calls with identical static parameters. ## Test plan - [x] Run `tests/ops/test_mamba.py` — all existing and new test cases pass (includes new smoke test `test_da_cumsum_fwd_missing_bias_raises`) - [x] Run `benchmarks/ops/bench_mamba.py` — benchmark executes cleanly with no regressions --------- Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> |
||
|---|---|---|
| .. | ||
| attention | ||
| __init__.py | ||
| activation.py | ||
| ada_layer_norm.py | ||
| ada_layer_norm_zero.py | ||
| argreduce.py | ||
| base.py | ||
| batch_norm.py | ||
| binary_arith.py | ||
| deltanet.py | ||
| engram.py | ||
| fft.py | ||
| fp8_lighting_indexer.py | ||
| fp8_quant.py | ||
| fused_add_layer_norm.py | ||
| fused_add_rms_norm.py | ||
| gated_deltanet.py | ||
| gemm.py | ||
| gla.py | ||
| group_norm.py | ||
| grouped_gemm.py | ||
| instance_norm.py | ||
| layer_norm.py | ||
| logical_reduce.py | ||
| mamba.py | ||
| mhc.py | ||
| moe.py | ||
| nsa_utils.py | ||
| reduce.py | ||
| rms_norm.py | ||
| softmax.py | ||
| topk_selector.py | ||
| vector_norm.py | ||
| workload_base.py | ||