[robertream]

Summary

Makes the cuTile kernel library element-generic, fixes a compiler bug that blocked some of those kernels, adds a sampling kernel, and lets the float kernels store fp8. Behavioral GPU coverage goes from 4 to 21 host-reference tests. One squashed commit.

1. Compiler fix: callee generic bound to a same-named caller var

When type inference computes the speculative expected type for a call argument, an unresolved callee generic in the parameter type was compiled against the caller's generic vars. If the caller bound a var of the same name but a different kind, the callee's param silently rebound to it.

A generic-element kernel with const D calling load_view_tko (view: &Partition<E, D>, where the op's own D is [i32; N]) hit exactly this: the kernel's scalar D shadowed the view op's array const D, and the resolver panicked with "Unexpected const arg D for variadic type Partition". Only the combination of a generic element and a colliding const name tripped it.

Fix: skip the speculative hint when the parameter type still mentions an unresolved callee generic, and let normal argument inference resolve the call. Regression coverage (generic_view_const_arg.rs) drives the raw load and store view ops over f16/bf16/f32; verified test-first (fails without the fix, passes with it).

2. Element-generic kernels

Adds element-generic siblings of the kernels, keeping the concrete _f16 forms (additive). ElementType for pure data-movement kernels, FloatElement where the math needs floats; data loads as the element and converts to f32 for the compute, converting back on store.

• Data movement (ElementType): embedding_batch, gather_row, kv_cache_update_seq{,_dynpos}, and the experimental copy_stacked_to_contiguous / copy_contiguous_to_stacked.
• Float math (FloatElement): add_2d, silu_mul_2d, rms_norm, add_rms_norm, qk_norm, rope_seq{,_dynpos}, qk_rope_dynpos, argmax_blocks, lm_head_argmax_blocks, flash_attn_causal_seq{,_dynpos}, and the experimental add_rms_norm_decode_raw, qk_norm_rope_kv_{prefill,decode}_raw, group_gemm_nt_desc.

Kernels that were unsuffixed yet f16 (fmha_prefill_causal, fmha_prefill_gqa, fmha_causal, fmha_decode_gqa_split, splitk_reduce_merge, fmha_prefill_gqa_lpt) are generalized in place over the element type; f16 callers pass f16 as E. Four load through the raw load_view_tko with const D, so they depend on the fix above. argmax_reduce_blocks_to_u32 stays concrete — it reduces already-f32 maxima + u32 indices, with no element to parameterize.

3. Sampling kernel

sampling::sample<E: FloatElement, const BLOCK> — temperature + top-k + top-p (nucleus) categorical sampling over a logit row, one CTA per batch row. Both filters are sort-free value-threshold bisections (k >= len / p >= 1.0 disable each, keeping the body branch-free over tiles); the draw is inverse-CDF via scan_sum against a per-row host uniform (no device RNG, so decode stays deterministic). Fills the gap between greedy argmax and real generation.

4. Element-type traits

The generic kernels are parameterized over two element-type traits, picked by whether the kernel does arithmetic:

• ElementType — the pre-existing base trait (const ZERO), already implemented across the scalar, integer, fp8, and packed sub-byte types. The pure data-movement kernels bound on this: they only load, copy, and store, so any element works. This PR adds no change to the trait — it just uses it as a kernel bound.
• FloatElement: ElementType — new in this PR. A marker trait bounding the kernels that compute in f32: they load the IO element, convert_tile it to f32, do the math, and narrow back on store. The bound exists because convert_tile's float conversion only bridges float↔float, so a float-math kernel must reject non-float elements at the type level. Implemented for f16, bf16, f32, f64, and the fp8 storage types f8e4m3fn / f8e5m2.

fp8 support is exactly those two FloatElement impls: the float kernels accept f8e4m3fn / f8e5m2 as storage (memory layout), computing in f32 — a bandwidth/footprint win for the memory-bound elementwise and norm kernels, not fp8 tensor-core compute and not a GEMM path.

Validation

• CPU: every generic form smoke-compiles; the resolver regression suite (generic_view_const_arg) and the compiler/library unit tests are green.
• GPU (H100), compiler fix: the cutile device suite (104 tests) passes through the fixed compiler.
• GPU (H100), kernel numerics vs host references — 21 tests over the safe Tensor-API generic kernels:
  • Attention: fmha_prefill_causal (f16 + bf16), fmha_causal, flash_attn_causal_seq, flash_attn_causal_seq_dynpos (causal masking), fmha_prefill_gqa (GQA head-packing), and fmha_decode_gqa_split + splitk_reduce_merge (split-K decode, end-to-end vs full softmax).
  • Norms: rms_norm, add_rms_norm, qk_norm vs an f32 RMSNorm reference.
  • Rope: rope_seq, rope_seq_dynpos, qk_rope_dynpos vs a GPT-NeoX rotation reference.
  • Argmax: argmax_blocks, lm_head_argmax_blocks (exact index + value).
  • KV cache: kv_cache_update_seq, kv_cache_update_seq_dynpos (exact scatter).
  • Pointwise / movement: add_2d, gather_row, embedding_batch (exact), silu_mul_2d over bf16 and f8e4m3fn.
  • sample — four behavioral checks (top-k=1 ⇒ argmax incl. the u=0 boundary, inverse-CDF bin selection, top-p collapse, masked-tail exclusion).

The raw-pointer (*mut E) and pointer-table (Tensor<i64>) external-ABI kernels (add_rms_norm_decode_raw, qk_norm_rope_kv_{prefill,decode}_raw, fmha_prefill_gqa_lpt, group_gemm_nt_desc, the kvbm copies) keep their <E> generalization but are left to their safe Tensor-API siblings for behavioral coverage.

---

I used an AI assistant (Claude) to write and check these changes. I reviewed the diff and validated the kernels on an H100 (21 numeric tests plus the 104-test device suite).

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.