Skip to content

decode_kernel_profiling.md

Latest commit

 

History

History
135 lines (102 loc) · 5.32 KB

decode_kernel_profiling.md

File metadata and controls

135 lines (102 loc) · 5.32 KB
Raw

Triton decode kernel + MBU analysis

Goal

Replace the pure-PyTorch inner SSM recurrence in autoregressive decode (seqlen = 1) with a fused Triton kernel, verify bit-exact parity against the naive scan, and measure how close we get to memory-bandwidth saturation on an A10G (peak ≈ 600 GB/s).

Kernel design

File: kernels/scan_decode.py

Per decode step the math per (batch, d_inner) element is:

delta_A = exp(dt * A)           # (N,)
delta_B = dt * B                # (N,)
h_new   = delta_A * h_old + delta_B * x
y       = sum_n C[n] * h_new[n]
y      += D * x                 # optional
y      *= silu(z)               # optional

Launch shape:

  • grid = (d_inner, batch)
  • One program owns one output scalar y[b, d] and its N-wide state row h[b, d, :].
  • BLOCK_N = next_pow2(d_state) with masking to cover non-pow2 state sizes.

Dtype contract (matches mamba_minimal.scan_naive):

  • Inputs may be fp16 / bf16 / fp32.
  • Accumulation always in fp32 (loads are cast).
  • State buffer is fp32 end-to-end (HF convention).
  • y is written in fp32 and cast to x.dtype in the wrapper.

State is written to a freshly allocated new_state buffer each call. One allocation per decode step is fine for tok/s regimes we care about and keeps the API purely functional.

Correctness

tests/test_scan_decode_triton.py — 23 / 23 pass.

  • fp32: atol=rtol=1e-5 across five shapes × {with/without D} × {with/without z}.
  • fp16 / bf16 at Mamba-130m shape: atol=rtol=5e-3 / 8e-3.
  • Parity vs selective_scan_naive at L=1 (the reference-parity oracle, which is bit-exact vs mamba_ssm.selective_scan_ref).

End-to-end: generate_native(model, ...) with use_triton=True is token-exact vs the pure-PyTorch path and vs HF Mamba-130m.

Microbench — kernel-in-isolation

benchmarks/decode_kernel.py, 1000 iters, CUDA events, A10G.

Bytes-moved model per step (per program):

  • Vector loads/stores (A row, B, C, h_old, h_new) = 5 * N fp32
  • ~6 scalar loads/stores ≈ 24 B
  • Total ≈ B * D * (20 * N + 24) bytes
Shape Dtype Triton (µs) PyTorch (µs) Speedup Triton GB/s % of 600 GB/s
B=1, D=1536, N=16 fp32 60.5 153.6 2.5x 8.7 1.5%
B=4, D=1536, N=16 fp32 59.6 138.5 2.3x 35.5 5.9%
B=8, D=1536, N=16 fp32 58.4 152.3 2.6x 72.4 12.1%
B=1, D=1536, N=16 fp16 75.3 201.5 2.7x 7.0 1.2%
B=1, D=1536, N=16 bf16 76.3 209.8 2.7x 6.9 1.2%
B=1, D=3072, N=16 fp32 61.5 145.9 2.4x 17.2 2.9%
B=1, D=1536, N=64 fp32 58.3 140.0 2.4x 34.4 5.7%

Takeaways:

  1. Consistent 2.3–2.7× over pure-PyTorch equivalent. PyTorch here is the same math with torch.exp, broadcast multiplies, and a sum — several kernel launches and intermediate allocations.
  2. Floor around ~60 µs per call across B, D, N. That's the launch / dispatch latency, not bandwidth. Doubling D doesn't meaningfully increase wall time until the kernel is compute/memory bound.
  3. MBU grows with batch. B=1 → 1.5%, B=8 → 12.1% of the 600 GB/s peak. At Mamba-130m decode shapes one step simply isn't enough work to saturate an A10G: B*D*20*N = 1*1536*20*16 ≈ 490 kB — that's a sub-microsecond transfer at peak bandwidth; everything else is overhead.

End-to-end (parallel-scan benchmark with --use-triton)

Running generate_native on Mamba-130m, 8-token prompt, 128 new tokens:

  • mamba_minimal.parallel (pure PyTorch decode): 73.1 tok/s
  • mamba_minimal.parallel+triton (Triton decode): 72.1 tok/s

The Triton kernel does not move end-to-end tok/s because decode wall time is dominated by the four nn.Linear calls per step (in_proj, x_proj, dt_proj, out_proj) plus Python-level per-step overhead. The SSM recurrence — the only part we replaced — is a small fraction of the step. This is exactly the headroom a production decode path would claim next by CUDA-graphing the step or fusing all four Linears + SSM into a single decode kernel. That's out of scope here.

nsight-compute note

ncu is not installed in this environment. MBU numbers above come from our own bytes-moved accounting (batch * d_inner * (5*N*4 + 6*4) bytes per step, divided by measured time). For a hardware-counter cross-check on a box that has ncu:

ncu --set full --kernel-name _selective_scan_decode_kernel \
    --launch-count 1 python benchmarks/decode_kernel.py

Expected counters to inspect: dram__bytes.sum (DRAM traffic), sm__throughput.avg.pct_of_peak_sustained_elapsed (SM utilization), l1tex__t_bytes.sum (L1 traffic vs DRAM to spot cache reuse).

Prefill path

No new prefill kernel here. The existing Triton fused chunked scan at kernels/scan_fused.py covers prefill; when unavailable the block falls back to the pure-PyTorch Blelloch scan (selective_scan_parallel).

Verdict

  • Correctness: bit-exact at fp32, tight at fp16 / bf16, parity vs the naive reference and HF generation preserved.
  • Kernel speedup in isolation: 2.3–2.7× over PyTorch equivalent.
  • MBU: 1.5% at B=1, 12% at B=8 — bound by launch latency, not bandwidth, at Mamba-130m decode shapes on A10G.
  • End-to-end tok/s: flat. The decode bottleneck on this model is the Linear projections and Python overhead, not the SSM recurrence. Real end-to-end wins need CUDA graphs or a fully fused decode step.