Two ops for masked-diffusion token decoding (OmniVoice's per-step select): masked_diffusion_scores fuses the classifier-free-guidance log-softmax over a (rows, C*V) logit block, the MASK-id exclusion, the argmax or top-fraction Gumbel sample, and the per-position confidence minus a per-codebook layer penalty plus position-temperature Gumbel noise, writing a (C, T) prediction grid and a score grid that is -inf wherever the token grid is already unmasked; masked_diffusion_commit writes the k selected flat positions from the prediction into the token grid and stamps the step index into an unmask-step grid; the uniforms come from a splitmix64 counter hash over (seed, domain, index) in detail/hash_rng.h so CPU and CUDA draw bit-identical noise and the CUDA kernel avoids FMA contraction so scores match CPU to ~3e-6 with identical selections; top_k_rows loses its CUDA limit of k <= 6144 (the per-thread shared-memory staging) and Metal's k <= 256 through a rank-selection kernel above k = 256 and a partial_sort on CPU above k = 64, all keeping the descending-value lower-index-first order, verified on one row of 32768 with ties at k from 1 to the full row; tests/test_masked_diffusion.cpp checks the small cases against an in-test double-precision reference and CPU against CUDA at C=8 T=400 V=1025

Author: wlejon

CMakeLists.txt

@@ -185,6 +185,7 @@ add_library(brotensor_cpu STATIC
     src/cpu/resample1d.cpp
     src/cpu/log_exp_round.cpp
     src/cpu/sample_logits.cpp
+    src/cpu/masked_diffusion.cpp
     src/cpu/noise.cpp
     src/cpu/l2_norm.cpp
     src/cpu/gated_delta_rule.cpp
@@ -284,6 +285,7 @@ if(BROTENSOR_WITH_CUDA)
         src/cuda/codec_quant.cu
         src/cuda/resample1d.cu
         src/cuda/sample_logits.cu
+        src/cuda/masked_diffusion.cu
         src/cuda/noise.cu
         src/cuda/l2_norm.cu
         src/cuda/gated_delta_rule.cu
@@ -382,6 +384,7 @@ if(BROTENSOR_WITH_METAL)
         src/metal/resample1d.mm
         src/metal/log_exp_round.mm
         src/metal/sample_logits.mm
+        src/metal/masked_diffusion.mm
         src/metal/noise.mm
         src/metal/l2_norm.mm
         src/metal/l2_normalize.mm
@@ -464,6 +467,7 @@ if(BROTENSOR_WITH_METAL)
         src/metal/resample1d.mm
         src/metal/log_exp_round.mm
         src/metal/sample_logits.mm
+        src/metal/masked_diffusion.mm
         src/metal/noise.mm
         src/metal/l2_norm.mm
         src/metal/l2_normalize.mm

include/brotensor/detail/hash_rng.h

@@ -0,0 +1,67 @@
+#pragma once
+
+// brotensor detail/hash_rng.h — counter-based hash RNG shared by the CPU and
+// CUDA backends (and by tests that need to reproduce an op's noise).
+//
+// A stateless, order-independent uniform generator: the uniform for element
+// `index` of stream (`seed`, `domain`) is a pure function of those three
+// integers, so a CPU loop and a CUDA grid produce bit-identical noise for
+// bit-identical inputs. The mixer is splitmix64 (Steele/Lea/Flood) — a
+// bijective 64-bit finaliser with full avalanche, which is exactly what a
+// counter-based generator needs; the input is `seed ^ (domain + index)` so a
+// caller varies `seed` per step and each op keeps a distinct `domain` for
+// each of its noise streams (a position stream and a vocabulary stream must
+// not collide at index 0).
+//
+// The uniform takes the high 24 bits of the hash to a float in [0, 1) — the
+// same mapping the Philox ops use (rand_uniform, sample_logits), so the value
+// set is identical: multiples of 2^-24, never 1.0.
+//
+// Every function is `inline` and, under nvcc, `__host__ __device__`, so the
+// header compiles into both a .cpp and a .cu translation unit with identical
+// integer arithmetic. Metal cannot include this header (MSL is compiled from
+// source strings); src/metal/masked_diffusion.mm carries a verbatim copy.
+
+#include <cstdint>
+
+#if defined(__CUDACC__)
+#define BROTENSOR_HASH_RNG_HD __host__ __device__
+#else
+#define BROTENSOR_HASH_RNG_HD
+#endif
+
+namespace brotensor::detail {
+
+// Domain constants for the ops that draw from this generator. Arbitrary
+// large odd constants, far apart, so `domain + index` never overlaps across
+// streams for any index a tensor can hold.
+inline constexpr std::uint64_t kHashDomainMaskedDiffusionPosition = 0x5851F42D4C957F2DULL;
+inline constexpr std::uint64_t kHashDomainMaskedDiffusionClass    = 0x14057B7EF767814FULL;
+
+// splitmix64 finaliser: one round of the SplitMix64 output function applied
+// to `x` (including the golden-ratio increment, so x == 0 does not map to 0).
+BROTENSOR_HASH_RNG_HD inline std::uint64_t splitmix64(std::uint64_t x) {
+    x += 0x9E3779B97F4A7C15ULL;
+    x = (x ^ (x >> 30)) * 0xBF58476D1CE4E5B9ULL;
+    x = (x ^ (x >> 27)) * 0x94D049BB133111EBULL;
+    return x ^ (x >> 31);
+}
+
+// Raw 64-bit hash for element `index` of stream (seed, domain).
+BROTENSOR_HASH_RNG_HD inline std::uint64_t hash_u64(std::uint64_t seed,
+                                                    std::uint64_t domain,
+                                                    std::uint64_t index) {
+    return splitmix64(seed ^ (domain + index));
+}
+
+// Uniform in [0, 1): top 24 bits of the hash / 2^24. Exact in FP32.
+BROTENSOR_HASH_RNG_HD inline float hash_uniform(std::uint64_t seed,
+                                                std::uint64_t domain,
+                                                std::uint64_t index) {
+    return static_cast<float>(hash_u64(seed, domain, index) >> 40) *
+           (1.0f / 16777216.0f);
+}
+
+}  // namespace brotensor::detail
+
+#undef BROTENSOR_HASH_RNG_HD

include/brotensor/detail/op_table.h

@@ -624,6 +624,15 @@
     X(sample_logits_into,                      void,  (const ::brotensor::Tensor& logits, float temperature, int top_k, float top_p,          \
                                                        uint64_t key, ::brotensor::Tensor& counter, ::brotensor::Tensor& scratch,              \
                                                        ::brotensor::Tensor& indices))                                                         \
+    /* ─── Masked-diffusion token selection (OmniVoice codebook grids) ─── */                                                                   \
+    X(masked_diffusion_scores,                 void,  (const ::brotensor::Tensor& logits, const ::brotensor::Tensor& tokens,                  \
+                                                       int T, int C, int V, int mask_id,                                                      \
+                                                       float guidance_scale, float layer_penalty,                                             \
+                                                       float position_temperature, float class_temperature,                                   \
+                                                       float class_top_frac, uint64_t seed,                                                   \
+                                                       ::brotensor::Tensor& pred, ::brotensor::Tensor& scores))                               \
+    X(masked_diffusion_commit,                 void,  (const ::brotensor::Tensor& pred, const ::brotensor::Tensor& idx, int k, int step,      \
+                                                       ::brotensor::Tensor& tokens, ::brotensor::Tensor& unmask_step))                        \
     /* ─── L2 norm + Gated Delta Rule (linear-attention text path) ─── */                                                                      \
     X(l2_norm_forward,                         void,  (const ::brotensor::Tensor& X, int head_dim, int num_heads, float eps,                   \
                                                        ::brotensor::Tensor& Y))                                                                \

include/brotensor/ops/sampling.h

@@ -118,4 +118,82 @@ void randn_truncated(float lo, float hi,
                      uint64_t key, uint64_t counter,
                      Tensor& Y);
 
+
+// ─── Masked-diffusion token selection (OmniVoice-style codebook grids) ─────
+//
+// One step of a masked-diffusion language model over a (C codebooks, T frames)
+// token grid with vocabulary V, where id `mask_id` marks a still-masked cell.
+// Two ops split the step: `masked_diffusion_scores` fuses classifier-free
+// guidance, log-softmax, per-cell prediction and the confidence score every
+// cell competes with; the host then picks the k best cells (top_k_rows over
+// the scores viewed as one (1, C*T) row) and `masked_diffusion_commit` writes
+// the chosen predictions back. FP32-only, implemented on CPU, CUDA and Metal.
+//
+// Math per cell (c, t), with c_logits / u_logits its conditional and
+// unconditional logit rows (upstream: OmniVoice._predict_tokens_with_scoring):
+//   if guidance_scale != 0:
+//       c = log_softmax(c_logits); u = log_softmax(u_logits)
+//       log_probs = log_softmax(c + guidance_scale * (c - u))
+//   else:
+//       log_probs = log_softmax(c_logits)
+//   log_probs[mask_id] = -inf
+//   if class_temperature > 0:
+//       k = ceil(class_top_frac * V), clamped to [1, V]
+//       filtered = log_probs with all but its k largest entries set to -inf
+//                  (ties at the k-th value keep the lower vocabulary index)
+//       pred = argmax(filtered / class_temperature + gumbel(u_class[v]))
+//   else:
+//       pred = argmax(log_probs)                       (ties: lowest index)
+//   confidence = max(log_probs)                        (UNfiltered, always)
+//   score = confidence - c * layer_penalty
+//   if position_temperature > 0:
+//       score = score / position_temperature + gumbel(u_pos)
+//   score = -inf where tokens[c, t] != mask_id         (already decided)
+//   gumbel(u) = -log(-log(u + 1e-10) + 1e-10), u ~ U[0, 1)
+//
+// Noise: the uniforms are a counter-based hash (detail/hash_rng.h,
+// splitmix64) of (seed, cell index c*T + t) for u_pos and of (seed, cell
+// index * V + v) for u_class, under two distinct domain constants — so the
+// CPU and CUDA backends draw bit-identical noise for the same `seed`, and the
+// caller gets fresh noise by varying `seed` per step. Nothing is consumed or
+// advanced; the same (inputs, seed) always yields the same result.
+//
+// Numerics: log-sum-exp accumulates in FP64 on CPU and CUDA (Metal, which has
+// no FP64, accumulates in FP32); everything else is FP32 with the same
+// operation order on every backend, so CPU and CUDA agree to the last ulp
+// except where their libm exp/log differ, and `pred` can only diverge on an
+// exact FP32 near-tie.
+//
+//   logits: (R, C*V) FP32; column c*V + v. R == 2*T when guidance_scale != 0
+//           (rows [0, T) conditional, rows [T, 2T) unconditional), R == T
+//           otherwise.
+//   tokens: (C, T) INT32 current grid; cell (c, t) is at c*T + t and is masked
+//           iff tokens == mask_id.
+//   pred:   (C, T) INT32 output, resized + dtype-set. The predicted id for
+//           every cell (never mask_id unless the whole row is -inf).
+//   scores: (C, T) FP32 output, resized + dtype-set. -inf where tokens !=
+//           mask_id.
+// Throws ("brotensor: masked_diffusion_scores: <reason>") for a non-FP32
+// logits / non-INT32 tokens, T/C/V < 1, mask_id outside [0, V), or a shape
+// that does not match (T, C, V, guidance_scale).
+void masked_diffusion_scores(const Tensor& logits, const Tensor& tokens,
+                             int T, int C, int V, int mask_id,
+                             float guidance_scale, float layer_penalty,
+                             float position_temperature, float class_temperature,
+                             float class_top_frac, std::uint64_t seed,
+                             Tensor& pred, Tensor& scores);
+
+
+// Commit the k selected cells: for i in [0, k): p = idx[i]; tokens[p] =
+// pred[p]; unmask_step[p] = step. `idx` holds flat cell indices c*T + t into
+// the (C, T) grids — the layout top_k_rows returns over `scores` viewed as
+// (1, C*T). tokens / unmask_step / pred are (C, T) INT32 (all pre-sized,
+// nothing is resized); idx is INT32 with at least k elements, of which only
+// the first k are read. k == 0 is a no-op. An index outside [0, C*T) is
+// ignored on every backend (a device kernel cannot throw), so the host should
+// only ever pass what top_k_rows produced. Throws for a dtype / shape / k
+// mismatch.
+void masked_diffusion_commit(const Tensor& pred, const Tensor& idx, int k, int step,
+                             Tensor& tokens, Tensor& unmask_step);
+
 }  // namespace brotensor