feat(sp4): Task A13.2 — retrofit moe_expert_util_ema with Pearls A+D
Replaces the kernel's hardcoded `alpha` with the shared `pearls_ad_update`
host-side helper. 9 ISV slots retrofit: 8 per-expert util + 1 gate entropy.
- Kernel `moe_expert_util_ema_update` signature: drops `(isv,
isv_util_base, isv_entropy_index, alpha)` for `(scratch_buf,
scratch_idx_util_base=44, scratch_idx_entropy=52)`.
- Single-block single-thread; computes per-expert col_mean in a single
forward pass (collapses prior dual-pass), writes each to scratch slots
[44..52) and Shannon entropy to slot 52.
- 9 ISV slots wired with Pearls A+D: ISV[118..126) (per-expert util,
Wiener offsets 132..156) + ISV[126] (gate entropy, Wiener offset 156).
- Launcher `gpu_moe_head.rs::launch_expert_util_ema`: drops alpha + isv
args; takes scratch_dev_ptr + 2 scratch_idx args.
- Wrapper `GpuDqnTrainer::launch_moe_expert_util_ema(_ema_alpha_unused)`:
sync + Pearls A+D loop over 9 slots via `apply_pearls(scratch_idx,
isv_idx)` closure.
- `debug_assert_eq!(MOE_NUM_EXPERTS, 8)` guards against K changes
silently breaking the contiguous scratch layout.
Behavior: stationary signals converge to the same value at adaptive rate.
The 9 slots stay semantically identical (per-expert col_mean, gate
entropy); HEALTH_DIAG mirror + the adaptive λ controller
`moe_lambda_eff_update` (which reads ISV[MOE_GATE_ENTROPY_EMA_INDEX])
continue to consume them unchanged.
Tests: `sp4_moe_expert_util_ema_writes_step_obs_via_pearl_a_then_converges_pearl_d`
drives kernel with B=128 K=8 perfect-uniform gate → analytical col_mean=1/8,
entropy=ln(8)≈2.0794. Asserts slot values ∈ ±1e-6/1e-5, non-target slots
remain 0; verifies Pearl A bootstrap + Pearl D convergence.
Per `feedback_no_atomicadd.md`,
`feedback_no_htod_htoh_only_mapped_pinned.md`. Build: `cargo check -p ml
--lib --tests --offline` clean (11 pre-existing warnings, no new warnings).
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -10151,24 +10151,86 @@ impl GpuDqnTrainer {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Phase 3 T3.5: launch `moe_expert_util_ema_update` ISV producer.
|
||||
/// Phase 3 T3.5; SP4 Layer A Task A13.2 retrofit (2026-05-01): launch
|
||||
/// `moe_expert_util_ema_update` step-observation producer, sync, then
|
||||
/// apply Pearls A+D host-side for all 9 ISV slots.
|
||||
///
|
||||
/// Single-block single-thread cold-path kernel. Reads `moe_gate_softmax_buf [B, K]`
|
||||
/// and EMA-updates ISV[MOE_EXPERT_UTIL_EMA_BASE..+8] (per-expert mean gate)
|
||||
/// and ISV[MOE_GATE_ENTROPY_EMA_INDEX] (gate entropy).
|
||||
pub fn launch_moe_expert_util_ema(&self, ema_alpha: f32) -> Result<(), MLError> {
|
||||
/// Single-block single-thread cold-path kernel. Reads
|
||||
/// `moe_gate_softmax_buf [B, K]` and writes per-expert col_mean to
|
||||
/// `producer_step_scratch_buf[44..52)` (one slot per expert; K=8) plus
|
||||
/// gate Shannon entropy to slot 52. Host then applies Pearls A+D for
|
||||
/// each slot in turn:
|
||||
/// - ISV[MOE_EXPERT_UTIL_EMA_BASE+k=118+k] for k in 0..K
|
||||
/// ← Wiener offset (44+k)*3
|
||||
/// - ISV[MOE_GATE_ENTROPY_EMA_INDEX=126]
|
||||
/// ← Wiener offset 52*3=156
|
||||
///
|
||||
/// **α dropped per SP4 — Pearls A+D adapt α from per-slot signal-vs-noise
|
||||
/// variance.** The `_ema_alpha_unused` argument is preserved so callers
|
||||
/// (training_loop.rs) compile unchanged.
|
||||
pub fn launch_moe_expert_util_ema(&self, _ema_alpha_unused: f32) -> Result<(), MLError> {
|
||||
use crate::cuda_pipeline::sp4_wiener_ema::{pearls_ad_update, WienerState};
|
||||
|
||||
debug_assert!(self.isv_signals_dev_ptr != 0,
|
||||
"launch_moe_expert_util_ema: isv_signals_dev_ptr must be non-zero");
|
||||
|
||||
const SCRATCH_IDX_UTIL_BASE: usize = 44;
|
||||
const SCRATCH_IDX_ENTROPY: usize = 52;
|
||||
debug_assert_eq!(MOE_NUM_EXPERTS, 8,
|
||||
"MOE_NUM_EXPERTS must be 8 for the contiguous 8-slot scratch layout");
|
||||
|
||||
let gate_ptr = self.moe_gate_softmax_buf.raw_ptr();
|
||||
let scratch_dev = self.producer_step_scratch_buf.dev_ptr;
|
||||
self.moe_head.launch_expert_util_ema(
|
||||
gate_ptr,
|
||||
self.isv_signals_dev_ptr,
|
||||
scratch_dev,
|
||||
self.config.batch_size,
|
||||
MOE_NUM_EXPERTS,
|
||||
MOE_EXPERT_UTIL_EMA_BASE,
|
||||
MOE_GATE_ENTROPY_EMA_INDEX,
|
||||
ema_alpha,
|
||||
)
|
||||
SCRATCH_IDX_UTIL_BASE,
|
||||
SCRATCH_IDX_ENTROPY,
|
||||
)?;
|
||||
self.stream.synchronize()
|
||||
.map_err(|e| MLError::ModelError(format!("moe_expert_util_ema sync: {e}")))?;
|
||||
|
||||
// ── Pearls A+D host-side for all 9 slots (8 experts + entropy) ──
|
||||
let apply_pearls = |scratch_idx: usize, isv_idx: usize| {
|
||||
let step_obs = unsafe {
|
||||
std::ptr::read_volatile(self.producer_step_scratch_buf.host_ptr.add(scratch_idx))
|
||||
};
|
||||
// Skip degenerate-zero observations — gate softmax always sums
|
||||
// to 1, so col_means are non-zero in steady state. A 0 here
|
||||
// means the kernel hasn't run (cold start before forward
|
||||
// populates moe_gate_softmax_buf) — Wiener state untouched.
|
||||
if step_obs == 0.0 { return; }
|
||||
|
||||
let wiener_offset = scratch_idx * 3;
|
||||
let prev_x_mean = unsafe {
|
||||
std::ptr::read_volatile(self.isv_signals_pinned.add(isv_idx))
|
||||
};
|
||||
let mut state = unsafe {
|
||||
let p = self.wiener_state_buf.host_ptr;
|
||||
WienerState {
|
||||
sample_var: std::ptr::read_volatile(p.add(wiener_offset)),
|
||||
diff_var: std::ptr::read_volatile(p.add(wiener_offset + 1)),
|
||||
x_lag: std::ptr::read_volatile(p.add(wiener_offset + 2)),
|
||||
}
|
||||
};
|
||||
let new_x_mean = pearls_ad_update(prev_x_mean, &mut state, step_obs);
|
||||
unsafe {
|
||||
std::ptr::write_volatile(self.isv_signals_pinned.add(isv_idx), new_x_mean);
|
||||
let wp = self.wiener_state_buf.host_ptr;
|
||||
std::ptr::write_volatile(wp.add(wiener_offset), state.sample_var);
|
||||
std::ptr::write_volatile(wp.add(wiener_offset + 1), state.diff_var);
|
||||
std::ptr::write_volatile(wp.add(wiener_offset + 2), state.x_lag);
|
||||
}
|
||||
};
|
||||
|
||||
for k in 0..MOE_NUM_EXPERTS {
|
||||
apply_pearls(SCRATCH_IDX_UTIL_BASE + k, MOE_EXPERT_UTIL_EMA_BASE + k);
|
||||
}
|
||||
apply_pearls(SCRATCH_IDX_ENTROPY, MOE_GATE_ENTROPY_EMA_INDEX);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Adaptive load-balance λ controller (2026-04-27).
|
||||
|
||||
@@ -265,31 +265,33 @@ impl GpuMoeHead {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Launch `moe_expert_util_ema_update` ISV producer on device pointers.
|
||||
/// Launch `moe_expert_util_ema_update` step-observation producer on
|
||||
/// device pointers. SP4 Layer A Task A13.2 retrofit (2026-05-01): writes
|
||||
/// per-expert col_mean to `scratch_buf[scratch_idx_util_base..+K)` and
|
||||
/// gate Shannon entropy to `scratch_buf[scratch_idx_entropy]`. The host
|
||||
/// applies Pearls A+D via `pearls_ad_update`.
|
||||
pub(crate) fn launch_expert_util_ema(
|
||||
&self,
|
||||
gate_ptr: u64,
|
||||
isv_dev_ptr: u64,
|
||||
scratch_dev_ptr: u64,
|
||||
b: usize,
|
||||
k: usize,
|
||||
isv_util_base: usize,
|
||||
isv_entropy_index: usize,
|
||||
alpha: f32,
|
||||
scratch_idx_util_base: usize,
|
||||
scratch_idx_entropy: usize,
|
||||
) -> Result<(), MLError> {
|
||||
let b_i32 = b as i32;
|
||||
let k_i32 = k as i32;
|
||||
let util_base_i32 = isv_util_base as i32;
|
||||
let entropy_idx_i32 = isv_entropy_index as i32;
|
||||
let util_base_i32 = scratch_idx_util_base as i32;
|
||||
let entropy_idx_i32 = scratch_idx_entropy as i32;
|
||||
unsafe {
|
||||
self.stream
|
||||
.launch_builder(&self.moe_expert_util_ema_update)
|
||||
.arg(&gate_ptr)
|
||||
.arg(&isv_dev_ptr)
|
||||
.arg(&scratch_dev_ptr)
|
||||
.arg(&b_i32)
|
||||
.arg(&k_i32)
|
||||
.arg(&util_base_i32)
|
||||
.arg(&entropy_idx_i32)
|
||||
.arg(&alpha)
|
||||
.launch(LaunchConfig {
|
||||
grid_dim: (1, 1, 1),
|
||||
block_dim: (1, 1, 1),
|
||||
|
||||
@@ -207,35 +207,33 @@ extern "C" __global__ void moe_softmax_backward(
|
||||
|
||||
extern "C" __global__ void moe_expert_util_ema_update(
|
||||
const float* __restrict__ gate, /* [B, K] */
|
||||
float* __restrict__ isv, /* [ISV_TOTAL_DIM] */
|
||||
float* __restrict__ scratch_buf, /* producer_step_scratch_buf */
|
||||
int B,
|
||||
int K,
|
||||
int isv_util_base,
|
||||
int isv_entropy_index,
|
||||
float alpha
|
||||
int scratch_idx_util_base, /* slot 44 — first of K util scratch slots */
|
||||
int scratch_idx_entropy /* slot 52 — gate-entropy scratch slot */
|
||||
) {
|
||||
/* Single-block, single-thread cold-path cadence kernel (matches
|
||||
* h_s2_rms_ema_update / aux_heads_loss_ema_update shape). */
|
||||
/* SP4 Layer A Task A13.2 retrofit (2026-05-01): write step
|
||||
* observations only; the host launcher applies Pearls A+D host-side
|
||||
* via `pearls_ad_update`. Writes one scratch slot per expert util
|
||||
* (`scratch_idx_util_base..+K`) plus one for gate entropy
|
||||
* (`scratch_idx_entropy`). **α dropped per SP4 — Pearls A+D adapt α
|
||||
* from per-slot signal-vs-noise variance.** Single-block,
|
||||
* single-thread cold-path cadence (matches retrofit family shape). */
|
||||
if (blockIdx.x != 0 || threadIdx.x != 0) return;
|
||||
|
||||
const float one_minus_alpha = 1.0f - alpha;
|
||||
|
||||
for (int k = 0; k < K; ++k) {
|
||||
float sum = 0.0f;
|
||||
for (int b = 0; b < B; ++b) sum += gate[b * K + k];
|
||||
float col_mean = sum / (float)B;
|
||||
float prev = isv[isv_util_base + k];
|
||||
isv[isv_util_base + k] = one_minus_alpha * prev + alpha * col_mean;
|
||||
}
|
||||
|
||||
/* Single forward pass: compute col_mean[k] for each expert, write to
|
||||
* scratch util slot, accumulate Shannon entropy. Avoids the 2× pass
|
||||
* over `gate` the prior implementation did (one for util, one for
|
||||
* entropy reread). */
|
||||
float entropy = 0.0f;
|
||||
for (int k = 0; k < K; ++k) {
|
||||
float sum = 0.0f;
|
||||
for (int b = 0; b < B; ++b) sum += gate[b * K + k];
|
||||
float col_mean = sum / (float)B;
|
||||
const float col_mean = sum / (float)B;
|
||||
scratch_buf[scratch_idx_util_base + k] = col_mean;
|
||||
if (col_mean > 1e-9f) entropy += -col_mean * logf(col_mean);
|
||||
}
|
||||
float prev_e = isv[isv_entropy_index];
|
||||
isv[isv_entropy_index] = one_minus_alpha * prev_e + alpha * entropy;
|
||||
scratch_buf[scratch_idx_entropy] = entropy;
|
||||
__threadfence_system();
|
||||
}
|
||||
|
||||
@@ -1696,3 +1696,117 @@ fn sp4_aux_label_scale_ema_writes_step_obs_via_pearl_a_then_converges_pearl_d()
|
||||
}
|
||||
assert!(((x_mean - expected_mean_abs) / expected_mean_abs).abs() < 0.01);
|
||||
}
|
||||
|
||||
// ── SP4 Task A13.2: moe_expert_util_ema retrofit ─────────────────────────────
|
||||
|
||||
const SP4_MOE_KERNELS_CUBIN: &[u8] =
|
||||
include_bytes!(concat!(env!("OUT_DIR"), "/moe_kernels.cubin"));
|
||||
|
||||
fn load_moe_expert_util_ema_kernel(stream: &Arc<CudaStream>) -> CudaFunction {
|
||||
let module = stream
|
||||
.context()
|
||||
.load_cubin(SP4_MOE_KERNELS_CUBIN.to_vec())
|
||||
.expect("load moe_kernels cubin");
|
||||
module
|
||||
.load_function("moe_expert_util_ema_update")
|
||||
.expect("load moe_expert_util_ema_update function")
|
||||
}
|
||||
|
||||
/// SP4 Task A13.2: `moe_expert_util_ema_update` writes K col_means
|
||||
/// to scratch[44..52) plus gate Shannon entropy to scratch[52]. Pearls A+D
|
||||
/// bootstrap + convergence verified host-side via `pearls_ad_update`.
|
||||
///
|
||||
/// Test gate: B=128 samples, K=8 experts, perfect uniform softmax →
|
||||
/// each col_mean = 1/8, entropy = ln(8) ≈ 2.0794.
|
||||
#[test]
|
||||
#[ignore = "requires GPU"]
|
||||
fn sp4_moe_expert_util_ema_writes_step_obs_via_pearl_a_then_converges_pearl_d() {
|
||||
use ml::cuda_pipeline::sp4_wiener_ema::{pearls_ad_update, WienerState};
|
||||
|
||||
const B: usize = 128;
|
||||
const K: usize = 8;
|
||||
const SCRATCH_IDX_UTIL_BASE: usize = 44;
|
||||
const SCRATCH_IDX_ENTROPY: usize = 52;
|
||||
const SP4_PRODUCER_COUNT: usize = 69;
|
||||
|
||||
// Uniform gate: each row sums to 1, all entries equal 1/K.
|
||||
let uniform: f32 = 1.0 / K as f32;
|
||||
let gate: Vec<f32> = vec![uniform; B * K];
|
||||
|
||||
let stream = make_test_stream();
|
||||
let kernel = load_moe_expert_util_ema_kernel(&stream);
|
||||
|
||||
let gate_buf = unsafe { MappedF32Buffer::new(B * K) }.expect("alloc gate");
|
||||
gate_buf.write_from_slice(&gate);
|
||||
let scratch_buf = unsafe { MappedF32Buffer::new(SP4_PRODUCER_COUNT) }
|
||||
.expect("alloc scratch");
|
||||
|
||||
let gate_dev = gate_buf.dev_ptr;
|
||||
let scratch_dev = scratch_buf.dev_ptr;
|
||||
let b_i32 = B as i32;
|
||||
let k_i32 = K as i32;
|
||||
let util_base_i32 = SCRATCH_IDX_UTIL_BASE as i32;
|
||||
let entropy_idx_i32 = SCRATCH_IDX_ENTROPY as i32;
|
||||
|
||||
unsafe {
|
||||
stream.launch_builder(&kernel)
|
||||
.arg(&gate_dev)
|
||||
.arg(&scratch_dev)
|
||||
.arg(&b_i32)
|
||||
.arg(&k_i32)
|
||||
.arg(&util_base_i32)
|
||||
.arg(&entropy_idx_i32)
|
||||
.launch(LaunchConfig {
|
||||
grid_dim: (1, 1, 1),
|
||||
block_dim: (1, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
})
|
||||
.expect("launch moe_expert_util_ema_update");
|
||||
}
|
||||
stream.synchronize().expect("sync");
|
||||
|
||||
let host = scratch_buf.read_all();
|
||||
|
||||
// Verify only target slots written.
|
||||
for (i, &v) in host.iter().enumerate() {
|
||||
let in_util = i >= SCRATCH_IDX_UTIL_BASE && i < SCRATCH_IDX_UTIL_BASE + K;
|
||||
let is_entropy = i == SCRATCH_IDX_ENTROPY;
|
||||
if !in_util && !is_entropy {
|
||||
assert_eq!(v, 0.0, "wrote outside target slots: scratch[{i}]={v}");
|
||||
}
|
||||
}
|
||||
|
||||
// Each col_mean should equal 1/8.
|
||||
for k in 0..K {
|
||||
let col_mean = host[SCRATCH_IDX_UTIL_BASE + k];
|
||||
assert!(
|
||||
(col_mean - uniform).abs() < 1e-6,
|
||||
"expert {k} col_mean {col_mean} vs expected {uniform}",
|
||||
);
|
||||
}
|
||||
|
||||
// Entropy = ln(K).
|
||||
let entropy = host[SCRATCH_IDX_ENTROPY];
|
||||
let expected_entropy: f32 = (K as f32).ln();
|
||||
assert!(
|
||||
(entropy - expected_entropy).abs() < 1e-5,
|
||||
"entropy {entropy} vs expected {expected_entropy}",
|
||||
);
|
||||
|
||||
// Pearl A bootstrap on the entropy slot (representative — same logic
|
||||
// for util slots).
|
||||
let mut state = WienerState::ZERO;
|
||||
let new_x_mean = pearls_ad_update(0.0, &mut state, entropy);
|
||||
assert_eq!(new_x_mean, entropy, "Pearl A returns step_obs directly");
|
||||
assert_eq!(state.x_lag, entropy);
|
||||
|
||||
// Pearl D convergence — 1000 stationary observations.
|
||||
let mut x_mean = new_x_mean;
|
||||
for _ in 0..1000 {
|
||||
x_mean = pearls_ad_update(x_mean, &mut state, entropy);
|
||||
}
|
||||
assert!(
|
||||
((x_mean - expected_entropy) / expected_entropy).abs() < 0.01,
|
||||
"Pearl D: stationary entropy converged to {x_mean} (expected {expected_entropy})",
|
||||
);
|
||||
}
|
||||
|
||||
@@ -2328,3 +2328,5 @@ SP4 Layer A Tasks A14+A15 — Pearl C engagement counter wired into all 5 Adam k
|
||||
SP4 Layer A Task A13.0 — h_s2_rms_ema retrofit (Pearls A+D) + buffer growth (2026-05-01): retrofit the existing per-step EMA producer in `crates/ml/src/cuda_pipeline/h_s2_rms_ema_kernel.cu` to consume the shared `pearls_ad_update` (Task A3) host-side helper instead of the hardcoded `ema_alpha`. Kernel signature changed to `(const float* h_s2, int B, int SH2, float* scratch_buf, int scratch_idx)` — RMS = sqrt(sum_sq / (B*SH2)) is reduced via the existing 256-thread shmem tree (no atomicAdd) and written to `producer_step_scratch_buf[scratch_idx=40]` with `__threadfence_system()`. Launcher `GpuDqnTrainer::launch_h_s2_rms_ema(_ema_alpha_unused: f32)` now syncs the stream after the kernel launch and applies Pearls A+D host-side via zero-copy mapped-pinned reads/writes of `isv_signals_pinned[H_S2_RMS_EMA_INDEX=96]` + `wiener_state_buf[120..123)` (scratch slot 40 × 3 = wiener offset 120). Degenerate-zero short-circuit before mutating ISV/Wiener state. **Buffer growth in same commit**: `SP4_PRODUCER_COUNT 47 → 69` (40 SP4 + 29 Task A13 retrofit producers); `wiener_state_buf 141 → 207` floats; producer_step_scratch_buf grows 47 → 69 entries. Stable layout doc-comments updated in `producer_step_scratch_buf` field comment, `launch_sp4_target_q_p99` slot-table comment, `reset_sp4_wiener_state` doc, `state_reset_registry.rs::sp4_wiener_state` description, and 5 `wiener_offset + 2 < 141` safety comments in retrofit launchers (now `< 207`). Test cubin reference `SP4_PRODUCER_COUNT: usize = 47` in `tests/sp4_producer_unit_tests.rs` updated to 69 across all 6 occurrences. **Unit test** `sp4_h_s2_rms_ema_writes_step_rms_via_pearl_a_then_converges_pearl_d` (`#[ignore]`-gated for GPU): drives the production kernel kernel-direct on a constant-5.0 stationary signal of B=4, SH2=64 (256 floats), asserts step_rms ∈ ±1e-4 of analytical RMS=5.0, asserts all non-target scratch slots remain 0, then exercises Pearl A bootstrap (`prev=0, state=ZERO` → returns step_rms directly + seeds x_lag) and Pearl D convergence (1000 stationary observations → x_mean within 1% of 5.0). Behaviour: cold-path producer with no consumer-facing change beyond the EMA mechanism (slot 96 is read by `mag_concat_qdir`'s adaptive-scale path); slot stays semantically identical (RMS of `save_h_s2`), only the EMA blending logic changes from hardcoded α to Wiener-optimal α*. Per `feedback_no_atomicadd.md`, `feedback_no_htod_htoh_only_mapped_pinned.md`, `feedback_no_partial_refactor.md`. `cargo check -p ml --lib --tests --offline` clean (11 pre-existing warnings, no new warnings); `cargo test -p ml --lib sp4_wiener_ema --offline` 6/6 passing.
|
||||
|
||||
SP4 Layer A Task A13.1 — aux_heads_loss_ema + aux_label_scale_ema retrofit (Pearls A+D) (2026-05-01): both kernels in `crates/ml/src/cuda_pipeline/aux_heads_loss_ema_kernel.cu` retrofit in the same commit per `feedback_no_partial_refactor.md` (single shared cubin, single producer family). Kernel `aux_heads_loss_ema_update` signature: `(next_bar_loss_scalar, regime_loss_scalar, scratch_buf, scratch_idx_next_bar=41, scratch_idx_regime=42)` — the two scalar reductions (already produced by `aux_*_loss_reduce` upstream) forward through to scratch slots [41..43) with `__threadfence_system()`. Kernel `aux_label_scale_ema_update` signature: `(label, B, scratch_buf, scratch_idx=43)` — single-block 256-thread shmem reduce for `mean(|label|)` writes to scratch slot 43 with `__threadfence_system()`. Three new ISV slots wired with Pearls A+D: ISV[AUX_NEXT_BAR_MSE_EMA_INDEX=113] (Wiener offset 123), ISV[AUX_REGIME_CE_EMA_INDEX=114] (Wiener offset 126), ISV[AUX_LABEL_SCALE_EMA_INDEX=117] (Wiener offset 129). Launchers updated in `gpu_aux_heads.rs::AuxHeadsForwardOps::launch_loss_ema` (drops `isv_dev_ptr`/`ema_alpha`/two `isv_*_index` args; takes scratch_dev_ptr + 2 scratch_idx args) and `launch_label_scale_ema` (same pattern, single scratch_idx arg). Wrapper `GpuDqnTrainer::launch_aux_heads_loss_ema(_ema_alpha_unused: f32)` now syncs the stream after the kernel launch and applies Pearls A+D for both slots in a `for` loop over `[(SCRATCH_IDX_NEXT_BAR=41, ISV[113]), (SCRATCH_IDX_REGIME=42, ISV[114])]`. The `aux_label_scale_ema` launch within `aux_heads_forward` (Step 2b — happens MID-STEP, before `next_bar_loss_reduce` and backward consumers read ISV[117]) gains an inline sync + Pearls A+D update so consumers see the up-to-date scale this step. Degenerate-zero short-circuit at all three slots prevents Wiener-state advancement when the upstream reduction yields 0 (e.g., before first aux head forward populates the scalars). **Unit tests** (collocated, `#[ignore]`-gated for GPU): (1) `sp4_aux_heads_loss_ema_writes_step_obs_via_pearl_a_then_converges_pearl_d` — drives kernel with stationary `nb_loss=0.5`, `rg_loss=1.2`, asserts step obs ∈ ±1e-6 of inputs at scratch[41]/[42], non-target slots remain 0, Pearl A returns inputs directly + seeds x_lag, 1000 stationary observations converge to within 1% via Pearl D. (2) `sp4_aux_label_scale_ema_writes_step_obs_via_pearl_a_then_converges_pearl_d` — drives kernel with mixed-sign B=256 labels (alternating ±3.0, mean_abs=3.0), asserts step obs ∈ ±1e-4 of analytical mean at scratch[43], non-target slots remain 0, Pearl A bootstrap + Pearl D convergence behave identically. Behaviour: cold-path producer + within-step producer (label scale) — slots stay semantically identical (next-bar MSE, regime CE, label-scale mean_abs); only the EMA blending logic changes. Per `feedback_no_atomicadd.md`, `feedback_no_htod_htoh_only_mapped_pinned.md`, `feedback_no_partial_refactor.md` — both kernels of the same `.cu` file retrofit in this commit; `cargo check -p ml --lib --tests --offline` clean (11 pre-existing warnings, no new warnings).
|
||||
|
||||
SP4 Layer A Task A13.2 — moe_expert_util_ema retrofit (Pearls A+D) (2026-05-01): retrofit the existing per-step MoE expert-utilisation + gate-entropy producer in `crates/ml/src/cuda_pipeline/moe_kernels.cu::moe_expert_util_ema_update`. Kernel signature changed: drops `(isv, isv_util_base, isv_entropy_index, alpha)` for `(scratch_buf, scratch_idx_util_base=44, scratch_idx_entropy=52)`. Single-block single-thread; computes per-expert col_mean of `gate [B, K]` in a single forward pass, writing each to `scratch_buf[scratch_idx_util_base + k]` for k in 0..K, and accumulating Shannon entropy (`-Σ col_mean·ln(col_mean)` with 1e-9 floor) for slot `scratch_idx_entropy`. The dual-pass formulation in the prior implementation (one pass for util, one re-read for entropy) collapses into a single pass — same output, no algorithmic change beyond the scratch-write routing. **9 new ISV slots wired with Pearls A+D**: `ISV[MOE_EXPERT_UTIL_EMA_BASE+k=118+k]` for k in 0..8 (Wiener offsets `(44+k)*3` = 132..156), plus `ISV[MOE_GATE_ENTROPY_EMA_INDEX=126]` (Wiener offset 156..159). Launcher `gpu_moe_head.rs::launch_expert_util_ema` drops `(isv_dev_ptr, isv_util_base, isv_entropy_index, alpha)` for `(scratch_dev_ptr, scratch_idx_util_base, scratch_idx_entropy)`. Wrapper `GpuDqnTrainer::launch_moe_expert_util_ema(_ema_alpha_unused: f32)` syncs the stream after launch and applies Pearls A+D in a per-slot loop via `apply_pearls(scratch_idx, isv_idx)` closure: 8 expert util pairs + 1 entropy pair. Degenerate-zero short-circuit at every slot prevents Wiener-state advancement when the kernel hasn't run (cold start before forward populates `moe_gate_softmax_buf`). `debug_assert_eq!(MOE_NUM_EXPERTS, 8)` guards against future K changes silently breaking the contiguous scratch layout. **Unit test** `sp4_moe_expert_util_ema_writes_step_obs_via_pearl_a_then_converges_pearl_d` (`#[ignore]`-gated for GPU): drives kernel with B=128 K=8 perfect-uniform gate (all entries 1/8) → analytical col_mean=1/8 per expert, entropy=ln(8)≈2.0794. Asserts each util slot ∈ ±1e-6, entropy slot ∈ ±1e-5, all non-target scratch slots remain 0; exercises Pearl A bootstrap (entropy slot returns directly + seeds x_lag) + Pearl D convergence (1000 stationary observations within 1% of analytical). Behaviour: cold-path producer with no consumer-facing change beyond the EMA mechanism — the 9 slots stay semantically identical (per-expert col_mean, gate entropy) and continue to be read by HEALTH_DIAG mirror + the adaptive λ controller `moe_lambda_eff_update` (which reads ISV[MOE_GATE_ENTROPY_EMA_INDEX]). Per `feedback_no_atomicadd.md`, `feedback_no_htod_htoh_only_mapped_pinned.md`. `cargo check -p ml --lib --tests --offline` clean (11 pre-existing warnings, no new warnings).
|
||||
|
||||
Reference in New Issue
Block a user