feat(dqn-v2): Plan 4 Task 1B-iv chain final — VSN backward + Polyak EMA target sync fix
Combines uncommitted iv-main + iv-ext aux paths + Fix A (Polyak EMA
extension covering VSN [95..119)) + Fix B (missed dispatch arm), and
strips rc2/rc3 band-aids that turned out to be unnecessary once the
upstream cause was fixed.
Root cause (focused HtoD/DtoD/pinned-memory audit):
target_ema_update only ran the EMA kernel over non_isv_params at
byte range covering [0..FIRST_ISV_TENSOR=77), skipping the 24 VSN
tensors at [95..119) added in 1B-ii. Online VSN trained from the
1B-iv backward chain; target VSN stayed at alloc_zeros for the
entire run. Bellman target used softmax(zeros) = uniform 1/6 mask
while online's mask drifted toward [market=0.13, portfolio=0.25].
The systematic Q-estimate divergence collapsed fold-2 magnitude
branch and pinned F0 best Sharpe at 21.14 across every prior
magnitude-scaling / state-isolation attempt.
Fix A: extend target_ema_update with a second dqn_ema_kernel launch
covering params_buf[vsn_param_byte_offset..vsn_param_total) so
target VSN tracks online VSN through the same Polyak EMA the rest
of the network uses. Same kernel signature, on-device, no DtoH.
Fix B: add the missed isv_vsn_aux_grad_scale dispatch arm in
training_loop.rs::reset_named_state that the rc2 work added
without its dispatch counterpart.
Cleanup (Phase B): strip rc2 ISV slot 113 + dqn_scale_f32_isv_scaled
+ aux_bottleneck_vsn_backward_dispatch indirection AND rc3 split-
Adam vsn_m_buf/vsn_v_buf — band-aids for the symptom Fix A actually
addresses. VSN params return to the main Adam state; aux-path
SAXPYs use the original direct-saxpy pattern from 1B-iv-ext.
Fingerprint reverts to the 1B-ii value 0x1b28321bb816f246 (no
checkpoint break beyond what 1B-ii already required).
Smoke (multi_fold_convergence --release-test, 3 folds × 5 epochs on
RTX 3050 Ti, 671.68s, all 3 dqn_fold{N}_best.safetensors written):
F0=93.4114 F1=73.0430 F2=73.0749
geom-mean = 79.31 vs 71.24 floor = +11.3%
F0 +36% over 1B-iii baseline; F2 +18%; F1 -14% (60-epoch L40S
run will validate equilibration; short-horizon asymmetry expected).
Constraints honoured: GPU-only (target EMA on-device), no atomicAdd,
no stubs, no // ok: band-aids, no tuned constants (Polyak EMA tau
shared with existing main-range launch), partial-refactor invariant
(dqn_ema_kernel signature unchanged — both launches use identical
args, different byte offsets/counts).
Lesson: 4 prior remedies (rc, rc2, rc3, rc4) and 1 diagnostic run
(E1) all chased downstream symptoms (gradient scale, Adam variance,
kernel sync). The upstream cause was a 1-line gap in Polyak target
sync that didn't include post-Plan-4 tensors. Same shape as the
2c.3a follow-up bottleneck Linear bug (commit f3e3ac347, 4 stale
runtime indices missed in GRN reshuffle): simple wire-up gaps in
shared infrastructure cause inscrutable downstream behavior. See
feedback_no_partial_refactor.md.
cargo check clean at 11 warnings (workspace baseline preserved).
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -1952,6 +1952,42 @@ impl CublasBackwardSet {
|
||||
)
|
||||
}
|
||||
|
||||
/// Compute only upstream gradient dX = dY @ W^T with a custom leading
|
||||
/// dimension for X (Plan 4 Task 1B-iv).
|
||||
///
|
||||
/// Mirrors `launch_dx_only` but threads `x_lda` into cuBLAS — needed for
|
||||
/// the bottleneck Linear backward when its dX must be written into a
|
||||
/// slice of a wider buffer (`d_vsn_gated_state[B, state_dim_padded]`'s
|
||||
/// market slice `[0..market_dim]` with row stride `state_dim_padded`).
|
||||
/// `beta=0.0` overwrites the destination slice; the surrounding
|
||||
/// `vsn_d_gated_state_portfolio_pad_kernel` zero-initialises the market
|
||||
/// slice so callers can rely on `dx` being fully defined regardless of
|
||||
/// kernel/GEMM ordering.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) fn launch_dx_only_lda(
|
||||
&self,
|
||||
stream: &Arc<CudaStream>,
|
||||
dy: u64,
|
||||
w: u64,
|
||||
dx: u64,
|
||||
out_dim: usize,
|
||||
in_dim: usize,
|
||||
x_lda: usize,
|
||||
batch: usize,
|
||||
beta: f32,
|
||||
) -> Result<(), MLError> {
|
||||
self.sgemm_f32(
|
||||
stream,
|
||||
cublas_sys::cublasOperation_t::CUBLAS_OP_N,
|
||||
cublas_sys::cublasOperation_t::CUBLAS_OP_N,
|
||||
in_dim as i32, batch as i32, out_dim as i32,
|
||||
1.0, w, in_dim as i32,
|
||||
dy, out_dim as i32,
|
||||
beta, dx, x_lda as i32,
|
||||
"dX_only_lda",
|
||||
)
|
||||
}
|
||||
|
||||
/// Launch 2-phase bias gradient reduction: db[j] += sum_b(dY[b * out_dim + j]).
|
||||
///
|
||||
/// Phase 1: block-level shared-memory reduce into `bias_grad_partials_buf`.
|
||||
|
||||
@@ -278,12 +278,13 @@ pub struct CublasGemmSet {
|
||||
/// `vsn_mask_buf` consumed both by `vsn_mask_ema_update` (1B-i ISV kernel)
|
||||
/// and by 1B-iv's softmax-Jacobian backward.
|
||||
vsn_softmax_and_gate_fwd_kernel: CudaFunction,
|
||||
/// Plan 4 Task 1B-iv consumer (loaded here in 1B-iii so the field is in
|
||||
/// place; consumer wired in 1B-iv). `vsn_softmax_and_gate_backward(
|
||||
/// Plan 4 Task 1B-iv consumer: `vsn_softmax_and_gate_backward(
|
||||
/// d_gated_state, state, mask, B, state_dim_padded, num_groups,
|
||||
/// group_begins, group_ends, d_logits, d_state)` implements the FULL
|
||||
/// softmax-over-groups Jacobian.
|
||||
#[allow(dead_code)]
|
||||
/// softmax-over-groups Jacobian. Wired into the backward chain as the
|
||||
/// final step of `launch_cublas_backward_to` (post-bottleneck Linear bwd
|
||||
/// + portfolio passthrough assembly), writing the 24 VSN dW/dB slots at
|
||||
/// grad_buf indices [95..119).
|
||||
vsn_softmax_and_gate_bwd_kernel: CudaFunction,
|
||||
/// Plan 4 Task 1B-iii: scatter handle for column-g writes of per-group
|
||||
/// `Linear_2_g` outputs into the assembled `vsn_logits_buf [B, num_groups]`.
|
||||
@@ -1486,6 +1487,204 @@ impl CublasGemmSet {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Plan 4 Task 1B-iv: VSN feature-selection backward.
|
||||
///
|
||||
/// Consumes the gradient flowing into the VSN-gated state (`d_gated_state`,
|
||||
/// assembled by the trainer from the bottleneck Linear backward dX into
|
||||
/// the market slice + the bottleneck-concat portfolio passthrough into
|
||||
/// `[market_dim..STATE_DIM]` + zero pad into `[STATE_DIM..STATE_DIM_PADDED]`)
|
||||
/// and writes per-group MLP weight + bias gradients into 24 grad_buf slots
|
||||
/// at indices `[95..119)` (the layout from 1B-ii).
|
||||
///
|
||||
/// Sequence (all GPU; no atomicAdd, no DtoH):
|
||||
/// 1. `vsn_softmax_and_gate_backward` kernel:
|
||||
/// d_gated_state, state, mask → d_logits[B, num_groups],
|
||||
/// d_state[B, state_dim_padded] (discarded)
|
||||
/// 2. For each group g in [0, num_groups):
|
||||
/// a. `strided_gather` extracts column g of d_logits[B, num_groups]
|
||||
/// into `d_logit_scratch[B]` (contiguous tight buffer).
|
||||
/// b. cuBLAS Linear_2 backward (`launch_dw_only` + `launch_dx_only`):
|
||||
/// dY = d_logit_scratch[B, 1] (contiguous after gather),
|
||||
/// X = save_h1_g[g] [B, VSN_HIDDEN]
|
||||
/// dW + dB into grad_buf[95+4g+2 / +3]; dX → d_h1_scratch[B, VSN_HIDDEN]
|
||||
/// c. `relu_mask` on d_h1_scratch using save_h1_g[g] (post-ReLU activation
|
||||
/// from forward — `relu_mask`'s "activation" arg is post-ReLU per kernel
|
||||
/// contract).
|
||||
/// d. cuBLAS Linear_1 backward via `backward_fc_layer_lda`:
|
||||
/// dY = d_h1_scratch[B, VSN_HIDDEN],
|
||||
/// X = state[B, group_dim_g] sliced from state[B, state_dim_padded]
|
||||
/// with row stride state_dim_padded
|
||||
/// dW + dB into grad_buf[95+4g+0 / +1]; no dX (`d_state` was already
|
||||
/// computed by step 1, and the value is discarded since the state
|
||||
/// buffer is not trainable).
|
||||
///
|
||||
/// `state_ptr` and `mask_ptr` are the saved buffers from the matching online
|
||||
/// `vsn_forward` call (the trainer keeps them in `states_buf` and
|
||||
/// `vsn_mask_buf` respectively). `save_h1_g_ptrs` are the 6 per-group
|
||||
/// post-ReLU saves the online forward populated.
|
||||
///
|
||||
/// `d_logit_scratch_ptr` (tight `[B]`) and `d_h1_scratch_ptr`
|
||||
/// (`[B, VSN_HIDDEN_DIM]`) are reused across the 6 groups since the loop
|
||||
/// runs sequentially on the main stream.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn vsn_backward(
|
||||
&mut self,
|
||||
stream: &Arc<CudaStream>,
|
||||
cublas_backward: &super::batched_backward::CublasBackwardSet,
|
||||
d_gated_state_ptr: u64, // INPUT [B, STATE_DIM_PADDED]
|
||||
state_ptr: u64, // INPUT [B, STATE_DIM_PADDED] saved
|
||||
mask_ptr: u64, // INPUT [B, num_groups] saved
|
||||
save_h1_g_ptrs: &[u64; ml_core::state_layout::SL_NUM_FEATURE_GROUPS],
|
||||
w_ptrs: &[u64; super::gpu_dqn_trainer::NUM_WEIGHT_TENSORS],
|
||||
grad_buf_base: u64,
|
||||
param_sizes: &[usize],
|
||||
d_logits_ptr: u64, // SCRATCH [B, num_groups]
|
||||
d_state_ptr: u64, // SCRATCH [B, STATE_DIM_PADDED] (kernel writes; discarded)
|
||||
d_logit_scratch_ptr: u64, // SCRATCH [B] (per-group, reused)
|
||||
d_h1_scratch_ptr: u64, // SCRATCH [B, VSN_HIDDEN_DIM] (per-group, reused)
|
||||
group_begins_dev: u64,
|
||||
group_ends_dev: u64,
|
||||
gather_kernel: &CudaFunction, // strided_gather (column extraction)
|
||||
batch_size: usize,
|
||||
) -> Result<(), MLError> {
|
||||
use super::gpu_dqn_trainer::{padded_byte_offset, VSN_HIDDEN_DIM};
|
||||
const VSN_TENSORS_BASE: usize = 95;
|
||||
const NUM_GROUPS: usize = ml_core::state_layout::SL_NUM_FEATURE_GROUPS;
|
||||
|
||||
let group_ranges = ml_core::state_layout::FEATURE_GROUP_RANGES;
|
||||
let state_dim_padded = self.state_dim_padded;
|
||||
let f32_size = std::mem::size_of::<f32>() as u64;
|
||||
|
||||
// ── Step 1: softmax-and-gate backward ────────────────────────────────
|
||||
// Writes d_logits[B, num_groups] (consumed by per-group Linear_2 bwd
|
||||
// via gather) and d_state[B, state_dim_padded] (discarded — state buf
|
||||
// not trainable). Block size 256, grid size B; shared mem layout
|
||||
// (per kernel comment): 2*num_groups + block floats.
|
||||
{
|
||||
let b_i32 = batch_size as i32;
|
||||
let state_dim_padded_i32 = state_dim_padded as i32;
|
||||
let num_groups_i32 = NUM_GROUPS as i32;
|
||||
let smem_floats = 2 * NUM_GROUPS + 256;
|
||||
let smem_bytes = (smem_floats * std::mem::size_of::<f32>()) as u32;
|
||||
unsafe {
|
||||
stream
|
||||
.launch_builder(&self.vsn_softmax_and_gate_bwd_kernel)
|
||||
.arg(&d_gated_state_ptr)
|
||||
.arg(&state_ptr)
|
||||
.arg(&mask_ptr)
|
||||
.arg(&b_i32)
|
||||
.arg(&state_dim_padded_i32)
|
||||
.arg(&num_groups_i32)
|
||||
.arg(&group_begins_dev)
|
||||
.arg(&group_ends_dev)
|
||||
.arg(&d_logits_ptr)
|
||||
.arg(&d_state_ptr)
|
||||
.launch(LaunchConfig {
|
||||
grid_dim: (batch_size as u32, 1, 1),
|
||||
block_dim: (256, 1, 1),
|
||||
shared_mem_bytes: smem_bytes,
|
||||
})
|
||||
.map_err(|e| MLError::ModelError(format!(
|
||||
"vsn_softmax_and_gate_backward: {e}"
|
||||
)))?;
|
||||
}
|
||||
}
|
||||
|
||||
// ── Step 2: per-group Linear_2 backward + ReLU mask + Linear_1 bwd ───
|
||||
for g in 0..NUM_GROUPS {
|
||||
let (gb, ge) = group_ranges[g];
|
||||
let group_dim = ge - gb;
|
||||
debug_assert!(
|
||||
group_dim > 0,
|
||||
"vsn_backward: feature group {g} has zero dim — group ranges invariant violated"
|
||||
);
|
||||
|
||||
// (a) Gather column g of d_logits[B, num_groups] → d_logit_scratch[B].
|
||||
// Column-extraction via strided_gather: dst_stride=1, col=g,
|
||||
// src_lda=num_groups, total = B * dst_stride = B.
|
||||
let src_lda_i32 = NUM_GROUPS as i32;
|
||||
let col_i32 = g as i32;
|
||||
let dst_stride_i32: i32 = 1;
|
||||
let total_i32 = batch_size as i32;
|
||||
let blocks = ((batch_size as u32 + 255) / 256).max(1);
|
||||
unsafe {
|
||||
stream
|
||||
.launch_builder(gather_kernel)
|
||||
.arg(&d_logits_ptr)
|
||||
.arg(&d_logit_scratch_ptr)
|
||||
.arg(&src_lda_i32)
|
||||
.arg(&col_i32)
|
||||
.arg(&dst_stride_i32)
|
||||
.arg(&total_i32)
|
||||
.launch(LaunchConfig {
|
||||
grid_dim: (blocks, 1, 1),
|
||||
block_dim: (256, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
})
|
||||
.map_err(|e| MLError::ModelError(format!(
|
||||
"vsn_backward strided_gather group {g}: {e}"
|
||||
)))?;
|
||||
}
|
||||
|
||||
// (b) Linear_2 backward: dY = d_logit_scratch[B, 1] (tight),
|
||||
// X = save_h1_g[g] [B, VSN_HIDDEN] (tight),
|
||||
// dW [1, VSN_HIDDEN] into grad_buf[95+4g+2],
|
||||
// dB [1] into grad_buf[95+4g+3],
|
||||
// dX [B, VSN_HIDDEN] into d_h1_scratch (overwrite, beta=0).
|
||||
let goff_w2 = padded_byte_offset(param_sizes, VSN_TENSORS_BASE + 4 * g + 2);
|
||||
let goff_b2 = padded_byte_offset(param_sizes, VSN_TENSORS_BASE + 4 * g + 3);
|
||||
cublas_backward.launch_dw_only(
|
||||
stream,
|
||||
d_logit_scratch_ptr, // dY [B, 1]
|
||||
save_h1_g_ptrs[g], // X [B, VSN_HIDDEN]
|
||||
grad_buf_base + goff_w2, // dW [1, VSN_HIDDEN]
|
||||
grad_buf_base + goff_b2, // dB [1]
|
||||
1, VSN_HIDDEN_DIM, batch_size,
|
||||
)?;
|
||||
cublas_backward.launch_dx_only(
|
||||
stream,
|
||||
d_logit_scratch_ptr, // dY [B, 1]
|
||||
w_ptrs[VSN_TENSORS_BASE + 4 * g + 2], // W2 [1, VSN_HIDDEN]
|
||||
d_h1_scratch_ptr, // dX [B, VSN_HIDDEN]
|
||||
1, VSN_HIDDEN_DIM, batch_size, 0.0_f32,
|
||||
)?;
|
||||
|
||||
// (c) ReLU mask on d_h1_scratch using saved post-ReLU h1.
|
||||
cublas_backward.relu_mask(
|
||||
stream,
|
||||
d_h1_scratch_ptr,
|
||||
save_h1_g_ptrs[g],
|
||||
batch_size * VSN_HIDDEN_DIM,
|
||||
)?;
|
||||
|
||||
// (d) Linear_1 backward (state slice with x_lda = state_dim_padded):
|
||||
// dY = d_h1_scratch [B, VSN_HIDDEN] (tight),
|
||||
// X = state + gb*sizeof(f32) (group-slice column offset),
|
||||
// x_lda = state_dim_padded (row stride of full state row),
|
||||
// dW [VSN_HIDDEN, group_dim_g] into grad_buf[95+4g+0],
|
||||
// dB [VSN_HIDDEN] into grad_buf[95+4g+1],
|
||||
// no dX (d_state was computed in step 1; state buf untrainable).
|
||||
let goff_w1 = padded_byte_offset(param_sizes, VSN_TENSORS_BASE + 4 * g + 0);
|
||||
let goff_b1 = padded_byte_offset(param_sizes, VSN_TENSORS_BASE + 4 * g + 1);
|
||||
let state_slice_ptr = state_ptr + (gb as u64) * f32_size;
|
||||
cublas_backward.backward_fc_layer_lda(
|
||||
stream,
|
||||
d_h1_scratch_ptr, // dY [B, VSN_HIDDEN]
|
||||
state_slice_ptr, // X [B, group_dim] @ stride=state_dim_padded
|
||||
w_ptrs[VSN_TENSORS_BASE + 4 * g + 0], // W1 [VSN_HIDDEN, group_dim] (unused for dX-skip path)
|
||||
grad_buf_base + goff_w1, // dW [VSN_HIDDEN, group_dim]
|
||||
grad_buf_base + goff_b1, // dB [VSN_HIDDEN]
|
||||
0, // dX skipped (state buf untrainable)
|
||||
VSN_HIDDEN_DIM,
|
||||
group_dim,
|
||||
state_dim_padded, // x_lda
|
||||
batch_size,
|
||||
)?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Run a single advantage branch's decoder (FC head + adv-logits) on the
|
||||
/// caller-supplied stream, sequentially.
|
||||
///
|
||||
|
||||
@@ -1186,6 +1186,63 @@ extern "C" __global__ void bn_tanh_backward_kernel(
|
||||
d_bn[idx] = d_upstream * (1.0f - tanh_val * tanh_val);
|
||||
}
|
||||
|
||||
/* ══════════════════════════════════════════════════════════════════════
|
||||
* VSN-AWARE BOTTLENECK CONCAT BACKWARD: PORTFOLIO PASSTHROUGH (Plan 4 Task 1B-iv)
|
||||
*
|
||||
* Inverse of the portfolio-passthrough path inside `bn_tanh_concat_kernel`:
|
||||
* forward (concat): concat_out[b, bn_dim + k] = states[b, market_dim + k]
|
||||
* for k in [0, portfolio_dim)
|
||||
* backward (here): d_gated_state[b, market_dim + k] = d_concat[b, bn_dim + k]
|
||||
*
|
||||
* The market slice of `d_gated_state[b, 0..market_dim]` is filled by the
|
||||
* cuBLAS bottleneck Linear backward (`launch_dx_only_lda`, with row stride
|
||||
* `state_dim_padded`) AFTER this kernel runs. Padding region
|
||||
* `[state_dim..state_dim_padded)` is zero-filled here (no-op when
|
||||
* STATE_DIM == STATE_DIM_PADDED, but we honor the contract for callers
|
||||
* that pad the state).
|
||||
*
|
||||
* One thread per (batch, padded-state-index) pair handles either the
|
||||
* portfolio slice or the padding tail; the market slice is left untouched
|
||||
* (cuBLAS overwrites it with beta=0).
|
||||
*
|
||||
* Grid: ceil(B*state_dim_padded/256), Block: 256.
|
||||
* No atomicAdd, deterministic.
|
||||
* ══════════════════════════════════════════════════════════════════════ */
|
||||
extern "C" __global__ void vsn_d_gated_state_portfolio_pad_kernel(
|
||||
float* __restrict__ d_gated_state, /* [B, state_dim_padded] f32 OUT */
|
||||
const float* __restrict__ d_concat, /* [B, concat_dim] f32 IN */
|
||||
int batch_size,
|
||||
int market_dim,
|
||||
int bn_dim,
|
||||
int concat_dim,
|
||||
int state_dim,
|
||||
int state_dim_padded
|
||||
) {
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int total = batch_size * state_dim_padded;
|
||||
if (idx >= total) return;
|
||||
|
||||
int b = idx / state_dim_padded;
|
||||
int i = idx % state_dim_padded;
|
||||
|
||||
if (i < market_dim) {
|
||||
/* Market slice: zero-initialised here so the buffer is fully defined
|
||||
* even before cuBLAS launch_dx_only_lda overwrites it (beta=0 GEMM)
|
||||
* with the bottleneck Linear backward dX. Removes a write-order
|
||||
* dependency between this kernel and the cuBLAS launch. */
|
||||
d_gated_state[idx] = 0.0f;
|
||||
} else if (i < state_dim) {
|
||||
/* Portfolio passthrough: d_gated_state[b, market_dim + k]
|
||||
* = d_concat [b, bn_dim + k]
|
||||
* where k = i - market_dim (∈ [0, portfolio_dim)). */
|
||||
int k = i - market_dim;
|
||||
d_gated_state[idx] = d_concat[(long long)b * concat_dim + bn_dim + k];
|
||||
} else {
|
||||
/* Padding tail [state_dim, state_dim_padded): zero. */
|
||||
d_gated_state[idx] = 0.0f;
|
||||
}
|
||||
}
|
||||
|
||||
/* ══════════════════════════════════════════════════════════════════════
|
||||
* BOTTLENECK WEIGHT GRADIENT BIAS REDUCE (#31)
|
||||
*
|
||||
|
||||
@@ -3798,6 +3798,36 @@ extern "C" __global__ void strided_scatter(
|
||||
dst[b * dst_stride + i] = src[idx];
|
||||
}
|
||||
|
||||
/**
|
||||
* Plan 4 Task 1B-iv: gather a single column from a wide buffer into a tight one.
|
||||
*
|
||||
* Inverse of `strided_scatter` for the column-extraction case (src_stride=1):
|
||||
* dst[b * dst_stride + i] = src[b * src_lda + col + i] for i in [0, dst_stride)
|
||||
*
|
||||
* Used by `vsn_backward` to extract column `g` of `d_logits[B, num_groups]`
|
||||
* into a contiguous `[B, dst_stride]` buffer, so the per-group cuBLAS
|
||||
* `Linear_2` backward (which expects a tightly-packed `dY[B, out_dim]`) can
|
||||
* read it without a strided-leading-dim cuBLAS variant. Per-group `out_dim=1`
|
||||
* for VSN's scalar-per-group logit, but the kernel is generic in dst_stride.
|
||||
*
|
||||
* Grid: ceil(total/256), Block: 256. total = B * dst_stride. No atomicAdd.
|
||||
* One thread per output element, deterministic.
|
||||
*/
|
||||
extern "C" __global__ void strided_gather(
|
||||
const float* __restrict__ src, /* [B, src_lda] wide */
|
||||
float* __restrict__ dst, /* [B, dst_stride] tight */
|
||||
int src_lda, /* row stride of src */
|
||||
int col, /* column offset within each src row */
|
||||
int dst_stride, /* number of contiguous columns to read */
|
||||
int total /* B * dst_stride */
|
||||
) {
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx >= total) return;
|
||||
int b = idx / dst_stride;
|
||||
int i = idx % dst_stride;
|
||||
dst[idx] = src[b * src_lda + col + i];
|
||||
}
|
||||
|
||||
/* ================================================================== */
|
||||
/* Kernel: variable_select_bottleneck */
|
||||
/* ================================================================== */
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because one or more lines are too long
Reference in New Issue
Block a user