…ate decode

BEGIN_PUBLIC
[Kernels][GPU] Coalesce, vectorize and parallelize selective_scan_update decode

Rewrites the Mamba/Mamba2 per-token decode kernel selective_scan_update_gpu
from a register-spilling, uncoalesced scalar kernel into two coalesced,
vectorized, state-parallel layouts selected at compile time by d_state, launched
over a 2D (dim-tile, batch) grid so (b, d) come from block indices with no
per-thread division:

* d_state > WARP_SIZE (Mamba2/3, d_state in {64,128,256}): warp-per-row -- one
  warp owns a (batch, dim) row, each lane owns d_state/WARP_SIZE CONTIGUOUS
  state elements loaded as one aligned 128-bit float4, sums them in-register,
  and the output sum over the state dim finishes with a single warp.sum (no
  shared memory, no barrier). A block packs 4 rows, matching the
  state-spaces/mamba Triton BLOCK_SIZE_M tiling.
* d_state <= WARP_SIZE (Mamba1/2, d_state in {2..32}): warp-cooperative -- each
  lane owns min(4, d_state) contiguous elements (an aligned float4),
  d_state/that lanes cooperate on a row and finish with a narrow
  warp.lane_group_sum (no shared memory, no barriers).

Key efficiency points, each verified on GB10 with ncu:
* Consecutive lanes on consecutive state elements -> coalesced loads.
* 16-byte-aligned width-4 loads emit one LD.E.128 instead of 4 scalar loads,
  cutting L1 load sectors ~3.6x (now identical to Triton's).
* 2D grid removes the per-thread divmod by the runtime dim, dropping executed
  instructions below Triton's.
* All four global loads (A, B, state, C) issue before any is consumed --
  the kernel is long-scoreboard (global-load-latency) bound at ~99% occupancy
  (ncu), so 4 loads in flight share one latency stall instead of serializing.
* O(1) registers/thread (28-35, zero spill) vs the old kernel's four
  SIMD[MAX_D_STATE=256] vectors (255-register cap -> ~3.6 MB spill, ~8%
  occupancy, ~78% memory-bound on scattered transactions, batch*dim/128 blocks).

Measured on GB10 (FP32, dim=1536, nsys eager kernel-trace median, MAX and the
Triton selective_state_update collected back-to-back / clock-matched): the
d_state=128 path reaches parity with Triton at the Mamba2 batched decode
profile (B=128: ~833 us on nsys kernel-trace both sides; ~860-896 us vs ~836 us
on amortized wall-clock -- parity within the measurement spread, from 2.9x
behind on the original kernel), and every profile is many-x faster than the
prior MAX kernel.
Triton retains a ~1.1-1.5x edge at small d_state / batch=1, where the kernel is
memory-latency-bound at ~99% occupancy; MAX matches or beats Triton there on
every measurable resource (instructions, L1/L2 sectors, occupancy, registers),
so the residual is instruction-scheduling codegen, not algorithm or traffic.

Adds real-dim Mamba1 (d_state=16) and Mamba2 (d_state=128) decode regression
tests and a microbenchmark.

Stacks on the selective_scan ops refactor / d_state generalization.
END_PUBLIC

Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
Signed-off-by: Evan <[email protected]>

BEGIN_PUBLIC
[Kernels][GPU][NFC] Factor shared selective_scan_update decode helpers

The two compile-time decode layouts in selective_scan_update_gpu (warp-per-row
for d_state > WARP_SIZE, warp-cooperative for d_state <= WARP_SIZE) carried
near-identical per-row scalar prologue, vectorized recurrence, and output
epilogue, duplicated across the `comptime if`/`else` arms.

Factor the shared work into three `@always_inline` helpers -- `_decode_row_prologue`
(dt + bias + softplus, x load), `_decode_row_recur` (the hoisted A/B/state/C
loads, state recurrence, state store, and in-register partial sum), and
`_decode_row_finalize` (D skip term, silu(z) gating, output store). Each layout
arm now contains only its distinctive logic: index setup, vector width, and the
cross-lane reduction (`warp.sum` vs `lane_group_sum`).

Pure refactor. The helpers are `@always_inline`, so the generated code -- and
the load-hoisting order the kernel depends on -- is unchanged. Verified on GB10
(FP32, dim=1536): 14/14 GPU goldens pass, and the bench is within run-to-run
noise at every profile (d=128 B=128 ~823 us, parity with Triton preserved;
d=16/128 B=1 and d=16 B=128 all unchanged).
END_PUBLIC

Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
Signed-off-by: Evan <[email protected]>

… tidy prologue

BEGIN_PUBLIC
[Kernels][GPU][NFC] Share selective_scan_update decode launch config; tidy prologue

Readability cleanup of the selective_scan_update decode path, no behavior
change:

* Extract the decode launch grid/block math into two helpers next to the
  kernel constants -- `selective_scan_update_decode_grid_dim_x(dim, d_state)`
  and `selective_scan_update_decode_block_dim(d_state)`. The op dispatch, the
  microbenchmark, and the GPU test now all derive grid/block from these instead
  of each re-deriving the warp-per-row vs warp-cooperative tiling inline. This
  also lets selective_scan_ops.mojo drop its duplicated `_WARP_COOP_BLOCK` /
  `_DECODE_WARPS_PER_BLOCK` / `_DECODE_ROWS_PER_THREAD` "must match" copies and
  the now-unused WARP_SIZE import.
* `_decode_row_prologue` now returns `(dt_val, x_val)`; the trivial
  `dt_x = dt_val * x_val` moves into `_decode_row_recur` where it is consumed.
  Call sites use named tuple unpacking instead of positional `pro[0/1/2]`.
* Drop a duplicate `from std.math import ceildiv` in the test and the unused
  `ceildiv` import in the bench.

The launch helpers compute exactly the values the inline code did, and the
prologue/recur change is `@always_inline` (relocated multiply, same inlined
code). Verified on GB10 (FP32, dim=1536): 14/14 GPU goldens pass and the bench
is within run-to-run noise at every profile (d=128 B=128 ~825 us parity with
Triton preserved; d=16 B=128 104-124 us noise band unchanged).
END_PUBLIC

Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
Signed-off-by: Evan <[email protected]>

…n-decode-opt

# Conflicts:
#	max/kernels/src/state_space/selective_scan.mojo
#	max/kernels/src/state_space/varlen_selective_scan.mojo
#	max/kernels/src/state_space/varlen_selective_scan_ops.mojo