@simulatte/doppler 0.1.8 → 0.2.0

package/src/inference/pipelines/text/attention/run.js

@@ -1,933 +0,0 @@
-
-
-import { isWeightBuffer, getWeightDtype } from '../../../../gpu/weight-buffer.js';
-import { getDevice } from '../../../../gpu/device.js';
-import { acquireBuffer, releaseBuffer } from '../../../../memory/buffer-pool.js';
-import {
-  runMatmul,
-  runRMSNorm,
-  runRoPE,
-  runAttention,
-  runAttentionBDPA,
-  runAttentionTieredQuant,
-  runAttentionTiered,
-  runSiLU,
-  castF16ToF32,
-  castF32ToF16,
-  runMatmulResidualFused,
-  shouldUseFusedMatmulResidual,
-} from '../../../../gpu/kernel-selector.js';
-import { createTensor } from '../../../../gpu/tensor.js';
-import { isKernelDebugEnabled, dumpTokenVector, dumpKVCache, logKernelStep } from '../debug-utils.js';
-import { applyLoRA } from '../lora-apply.js';
-import { getLoRAModule } from '../lora.js';
-import { kernelTrace, traceStep } from '../kernel-trace.js';
-import { log, trace } from '../../../../debug/index.js';
-import { selectRuleValue } from '../../../../rules/rule-registry.js';
-import { runProbes } from '../probes.js';
-import { SlidingWindowKVCache } from '../../../kv-cache.js';
-import {
-  recordAttentionInputs,
-  shouldForceF32AttentionProjectionForRoPE,
-  resolveAttentionProjectionOutputDtype,
-  projectAttentionQKV,
-  applyAttentionQKNorm,
-} from './projections.js';
-import { prepareAttentionProjectionInput } from './output-projection.js';
-
-import {
-  shouldDebugLayer,
-  markStageLogged,
-} from './types.js';
-
-const ATTENTION_DTYPE_LOGGED = new Set();
-
-function normalizeLayerType(layerType) {
-  return typeof layerType === 'string' ? layerType.trim().toLowerCase() : '';
-}
-
-function isSlidingLayerType(layerType) {
-  const normalized = normalizeLayerType(layerType);
-  return normalized === 'sliding_attention'
-    || normalized === 'local_attention'
-    || normalized === 'local'
-    || normalized === 'sliding';
-}
-
-export async function runLayerAttentionGPU(
-  input,
-  layerWeights,
-  config,
-  state,
-  debug = false,
-  debugFlags = {},
-  getWeightBuffer,
-  getNormWeightBuffer,
-  debugCheckBuffer,
-  lora
-) {
-  const {
-    layerIdx,
-    numTokens,
-    isPrefill,
-    numHeads,
-    numKVHeads,
-    headDim,
-    hiddenSize,
-    rmsNormEps,
-    currentSeqLen,
-    slidingWindow,
-    layerType,
-    residualTensor,
-    attnSoftcap,
-    queryPreAttnScalar,
-    skipInputNorm = false,
-    tokenIds = null,
-    kernelPath = null,
-    disableRoPE = false,
-  } = config;
-
-  const device = getDevice();
-
-  const desiredOutputDtype = selectRuleValue('shared', 'dtype', 'f16OrF32FromDtype', {
-    dtype: config.activationDtype,
-  });
-  const wantsF16Output = desiredOutputDtype === 'f16';
-  const useF16Activations = wantsF16Output;
-  const kvCacheFallback = selectRuleValue('inference', 'dtype', 'f16OrF32', { useF16: wantsF16Output });
-  const kvCacheDtype = state.kvCache?.kvDtype ?? kvCacheFallback;
-  const allowF16Attention = wantsF16Output && kvCacheDtype === 'f16';
-  let attentionInput = input;
-  let attentionInputTemp = false;
-  let normed = attentionInput;
-  let qTensor = null;
-  let qGateTensor = null;
-  let kTensor = null;
-  let vTensor = null;
-  let attnOutput = null;
-  let attnForProjection = null;
-  let output = null;
-  let finalOutput = null;
-  let oProjInputTemp = null;
-  if (wantsF16Output && !allowF16Attention) {
-    attentionInput = await castF16ToF32(input);
-    attentionInputTemp = true;
-    normed = attentionInput;
-  }
-
-  // Debug: attention input for configured layers
-  if (isPrefill && shouldDebugLayer(layerIdx, debugFlags.debugLayers) && !markStageLogged(layerIdx, 'attn_input', debugFlags) && debugCheckBuffer) {
-    await debugCheckBuffer(attentionInput.buffer, `L${layerIdx} attention input (GPU)`, numTokens, hiddenSize);
-  }
-
-  // Debug logging moved to debug-utils.ts (enable via setDebugConfig)
-
-  if (!layerWeights) {
-    // Return zeros if no weights
-    const bytesPerElement = wantsF16Output ? 2 : 4;
-    const outputBuf = acquireBuffer(numTokens * hiddenSize * bytesPerElement, undefined, 'attn_output');
-    const outputDtype = selectRuleValue('inference', 'dtype', 'f16OrF32', { useF16: wantsF16Output });
-    const output = createTensor(outputBuf, outputDtype, [numTokens, hiddenSize], 'attn_output');
-    return { output, residualFused: false };
-  }
-
-  const qSize = numTokens * numHeads * headDim;
-  const kvSize = numTokens * numKVHeads * headDim;
-
-  // 1. Input norm
-  
-  try {
-  if (!skipInputNorm && layerWeights.inputNorm && getNormWeightBuffer) {
-    const normWeightBuf = getNormWeightBuffer(layerWeights.inputNorm, 'input_norm');
-
-    // Debug: norm weights for configured layers
-    if (isPrefill && shouldDebugLayer(layerIdx, debugFlags.debugLayers) && !markStageLogged(layerIdx, 'norm_weights', debugFlags) && debugCheckBuffer) {
-      await debugCheckBuffer(normWeightBuf, `L${layerIdx} input norm weights (GPU)`, 1, hiddenSize);
-    }
-
-    normed = await runRMSNorm(attentionInput, normWeightBuf, rmsNormEps, {
-      batchSize: numTokens,
-      hiddenSize,
-      rmsNormWeightOffset: config.rmsNormWeightOffset,
-    });
-    if (!(layerWeights.inputNorm instanceof GPUBuffer)) releaseBuffer(normWeightBuf);
-
-    // Trace input norm output
-    if (kernelTrace.enabled) {
-      await traceStep('rmsnorm', `L${layerIdx}.input_norm`, layerIdx, normed.buffer, [numTokens, hiddenSize]);
-    }
-
-    if (isKernelDebugEnabled(layerIdx)) {
-      logKernelStep('rmsnorm', { layerIdx, label: 'input_norm', size: numTokens * hiddenSize });
-      await dumpTokenVector(normed.buffer, 'input_norm_out', {
-        layerIdx,
-        tokenIdx: Math.max(0, numTokens - 1),
-        rowSize: hiddenSize,
-        dtype: normed.dtype,
-      });
-    }
-  }
-
-  // Debug: Check normed input for L0 prefill
-  // Debug: normed input for configured layers
-  if (isPrefill && shouldDebugLayer(layerIdx, debugFlags.debugLayers) && !markStageLogged(layerIdx, 'attn_normed', debugFlags) && debugCheckBuffer) {
-    await debugCheckBuffer(normed.buffer, `L${layerIdx} normed input (GPU)`, numTokens);
-  }
-
-  const debugLayers = debugFlags.debugLayers;
-  const shouldLogLayer = debugLayers === null ? layerIdx === 0 : shouldDebugLayer(layerIdx, debugLayers);
-  if (shouldLogLayer) {
-    const phase = selectRuleValue('kernels', 'attention', 'phase', { isDecode: !isPrefill });
-    const logKey = `L${layerIdx}_${phase}_dtypes`;
-    if (!ATTENTION_DTYPE_LOGGED.has(logKey)) {
-      ATTENTION_DTYPE_LOGGED.add(logKey);
-      trace.attn(layerIdx, `dtypes: activation=${config.activationDtype ?? 'unknown'}, input=${input.dtype}, normed=${normed.dtype}`);
-    }
-  }
-
-  if (isKernelDebugEnabled(layerIdx)) {
-    await dumpTokenVector(normed.buffer, 'attn_in', {
-      layerIdx,
-      tokenIdx: Math.max(0, numTokens - 1),
-      rowSize: hiddenSize,
-      dtype: normed.dtype,
-    });
-  }
-
-  // 2. Q/K/V projections
-  const matmulOutputDtype = resolveAttentionProjectionOutputDtype(desiredOutputDtype, {
-    forceF32: shouldForceF32AttentionProjectionForRoPE({
-      attentionInputDtype: desiredOutputDtype,
-      headDim,
-      rotaryDim: config.ropeRotaryDim,
-      interleaved: config.ropeInterleaved,
-    }),
-  });
-  let usedFusedQKV = false;
-  ({ qTensor, qGateTensor, kTensor, vTensor, usedFusedQKV } = await projectAttentionQKV({
-    recorder: null,
-    normed,
-    layerWeights,
-    numTokens,
-    numHeads,
-    numKVHeads,
-    headDim,
-    hiddenSize,
-    layerIdx,
-    kernelPath,
-    matmulOutputDtype,
-    getWeightBuffer,
-    lora,
-    attentionOutputGate: config.attentionOutputGate === true,
-    releaseTemporary: (buffer) => releaseBuffer(buffer),
-    onFusedQKV: layerIdx === 0 && isPrefill
-      ? ({ qSize: qSizeFused, kSize: kSizeFused, vSize: vSizeFused, totalSize }) => {
-        trace.attn(layerIdx, `Using fused QKV path: ${qSizeFused}+${kSizeFused}+${vSizeFused}=${totalSize}`);
-      }
-      : null,
-  }));
-
-  // Trace Q/K/V projections
-  if (kernelTrace.enabled) {
-    await traceStep('matmul', `L${layerIdx}.q_proj`, layerIdx, qTensor.buffer, [numTokens, numHeads * headDim]);
-    await traceStep('matmul', `L${layerIdx}.k_proj`, layerIdx, kTensor.buffer, [numTokens, numKVHeads * headDim]);
-    await traceStep('matmul', `L${layerIdx}.v_proj`, layerIdx, vTensor.buffer, [numTokens, numKVHeads * headDim]);
-  }
-  await runProbes('q_proj', qTensor.buffer, {
-    layerIdx,
-    numTokens,
-    hiddenSize: numHeads * headDim,
-    probes: state.debugProbes,
-    dtype: qTensor.dtype,
-  });
-  await runProbes('k_proj', kTensor.buffer, {
-    layerIdx,
-    numTokens,
-    hiddenSize: numKVHeads * headDim,
-    probes: state.debugProbes,
-    dtype: kTensor.dtype,
-  });
-  await runProbes('v_proj', vTensor.buffer, {
-    layerIdx,
-    numTokens,
-    hiddenSize: numKVHeads * headDim,
-    probes: state.debugProbes,
-    dtype: vTensor.dtype,
-  });
-
-  // Kernel step debug: Q/K/V projections
-  if (isKernelDebugEnabled(layerIdx)) {
-    logKernelStep('matmul', { layerIdx, label: 'Q_proj', M: numTokens, N: numHeads * headDim, K: hiddenSize });
-    await dumpTokenVector(qTensor.buffer, 'Q_proj', {
-      layerIdx,
-      tokenIdx: Math.max(0, numTokens - 1),
-      rowSize: numHeads * headDim,
-      dtype: qTensor.dtype,
-    });
-    logKernelStep('matmul', { layerIdx, label: 'K_proj', M: numTokens, N: numKVHeads * headDim, K: hiddenSize });
-    await dumpTokenVector(kTensor.buffer, 'K_proj', {
-      layerIdx,
-      tokenIdx: Math.max(0, numTokens - 1),
-      rowSize: numKVHeads * headDim,
-      dtype: kTensor.dtype,
-    });
-    logKernelStep('matmul', { layerIdx, label: 'V_proj', M: numTokens, N: numKVHeads * headDim, K: hiddenSize });
-    await dumpTokenVector(vTensor.buffer, 'V_proj', {
-      layerIdx,
-      tokenIdx: Math.max(0, numTokens - 1),
-      rowSize: numKVHeads * headDim,
-      dtype: vTensor.dtype,
-    });
-  }
-
-  // Debug: Check Q/K/V after projections for L0 prefill
-  // Debug: Q/K/V projections for configured layers
-  if (isPrefill && shouldDebugLayer(layerIdx, debugFlags.debugLayers) && !markStageLogged(layerIdx, 'qkv_proj', debugFlags) && debugCheckBuffer) {
-    await debugCheckBuffer(qTensor.buffer, `L${layerIdx} Q after proj (GPU)`, numTokens, numHeads * headDim);
-    await debugCheckBuffer(kTensor.buffer, `L${layerIdx} K after proj (GPU)`, numTokens, numKVHeads * headDim);
-    await debugCheckBuffer(vTensor.buffer, `L${layerIdx} V after proj (GPU)`, numTokens, numKVHeads * headDim);
-  }
-
-  // Optional per-head Q/K normalization
-  const wantsQKNorm = config.queryKeyNorm === true;
-  const hasQNorm = !!layerWeights.qNorm;
-  const hasKNorm = !!layerWeights.kNorm;
-  if (isKernelDebugEnabled(layerIdx)) {
-    logKernelStep('qk_norm', { layerIdx, label: `hasQ=${hasQNorm} hasK=${hasKNorm} wants=${wantsQKNorm}` });
-  }
-  if (wantsQKNorm && layerIdx === 0 && (!hasQNorm || !hasKNorm)) {
-    log.warn('Attention', `Q/K norm requested but weights missing (hasQ=${hasQNorm}, hasK=${hasKNorm}); skipping QK norm.`);
-  }
-
-  ({ qTensor, kTensor } = await applyAttentionQKNorm({
-    recorder: null,
-    qTensor,
-    kTensor,
-    layerWeights,
-    getNormWeightBuffer,
-    rmsNormEps,
-    numTokens,
-    numHeads,
-    numKVHeads,
-    headDim,
-    rmsNormWeightOffset: config.rmsNormWeightOffset,
-    releaseTemporary: (buffer) => releaseBuffer(buffer),
-    onQNormApplied: isKernelDebugEnabled(layerIdx)
-      ? async (tensor) => {
-        await dumpTokenVector(tensor.buffer, 'Q_norm', {
-          layerIdx,
-          tokenIdx: Math.max(0, numTokens - 1),
-          rowSize: numHeads * headDim,
-          dtype: tensor.dtype,
-        });
-      }
-      : null,
-    onKNormApplied: isKernelDebugEnabled(layerIdx)
-      ? async (tensor) => {
-        await dumpTokenVector(tensor.buffer, 'K_norm', {
-          layerIdx,
-          tokenIdx: Math.max(0, numTokens - 1),
-          rowSize: numKVHeads * headDim,
-          dtype: tensor.dtype,
-        });
-      }
-      : null,
-  }));
-
-  if (normed !== attentionInput) releaseBuffer(normed.buffer);
-  if (attentionInputTemp) releaseBuffer(attentionInput.buffer);
-
-  // 3. RoPE (modifies tensor in-place)
-
-  if (!disableRoPE && state.ropeFreqsCos && state.ropeFreqsSin) {
-    await runRoPE(qTensor, state.ropeFreqsCos, state.ropeFreqsSin, numTokens, {
-      numHeads,
-      headDim,
-      rotaryDim: config.ropeRotaryDim,
-      interleaved: config.ropeInterleaved,
-      startPos: currentSeqLen,
-    });
-    await runRoPE(kTensor, state.ropeFreqsCos, state.ropeFreqsSin, numTokens, {
-      numHeads: numKVHeads,
-      headDim,
-      rotaryDim: config.ropeRotaryDim,
-      interleaved: config.ropeInterleaved,
-      startPos: currentSeqLen,
-    });
-
-    // Trace RoPE outputs
-    if (kernelTrace.enabled) {
-      await traceStep('rope', `L${layerIdx}.q_rope`, layerIdx, qTensor.buffer, [numTokens, numHeads * headDim]);
-      await traceStep('rope', `L${layerIdx}.k_rope`, layerIdx, kTensor.buffer, [numTokens, numKVHeads * headDim]);
-    }
-  }
-  await runProbes('q_rope', qTensor.buffer, {
-    layerIdx,
-    numTokens,
-    hiddenSize: numHeads * headDim,
-    probes: state.debugProbes,
-    dtype: qTensor.dtype,
-  });
-  await runProbes('k_rope', kTensor.buffer, {
-    layerIdx,
-    numTokens,
-    hiddenSize: numKVHeads * headDim,
-    probes: state.debugProbes,
-    dtype: kTensor.dtype,
-  });
-  if (isKernelDebugEnabled(layerIdx)) {
-    logKernelStep('rope', { layerIdx, label: `startPos=${currentSeqLen}` });
-    await dumpTokenVector(qTensor.buffer, 'Q_rope', {
-      layerIdx,
-      tokenIdx: Math.max(0, numTokens - 1),
-      rowSize: numHeads * headDim,
-      dtype: qTensor.dtype,
-    });
-    await dumpTokenVector(kTensor.buffer, 'K_rope', {
-      layerIdx,
-      tokenIdx: Math.max(0, numTokens - 1),
-      rowSize: numKVHeads * headDim,
-      dtype: kTensor.dtype,
-    });
-  }
-
-  // Debug: Check Q/K after RoPE for L0 prefill
-  // Debug: Q/K after RoPE for configured layers
-  if (isPrefill && shouldDebugLayer(layerIdx, debugFlags.debugLayers) && !markStageLogged(layerIdx, 'qk_rope', debugFlags) && debugCheckBuffer) {
-    await debugCheckBuffer(qTensor.buffer, `L${layerIdx} Q after RoPE (GPU)`, numTokens, numHeads * headDim);
-    await debugCheckBuffer(kTensor.buffer, `L${layerIdx} K after RoPE (GPU)`, numTokens, numKVHeads * headDim);
-  }
-
-  // 4. Update KV cache (cache stores raw GPUBuffers for memory efficiency)
-  
-  let cachedK;
-  
-  let cachedV;
-  let kvLenForAttention = currentSeqLen + numTokens;
-  let causalForAttention = config.causalAttention !== false;
-  let startPosForMask = currentSeqLen;
-  let kvStart = 0;
-  let kvLayout = 'contiguous';
-  let kvPageTable = null;
-  let kvPageSize = 0;
-  let cachedKHot;
-  let cachedVHot;
-  let cachedKCold;
-  let cachedVCold;
-  let coldScalesK = null;
-  let coldScalesV = null;
-  let coldPackedStride = 0;
-  let coldQuantMode = 'none';
-  let coldLen = 0;
-  let hotLen = 0;
-  let hotStart = 0;
-  let hotWindow = 0;
-  let coldPageTable = null;
-  let coldPageSize = 0;
-  let bdpaBasisK = null;
-  let bdpaBasisV = null;
-  let bdpaPagedK = null;
-  let bdpaPagedV = null;
-  let bdpaIndex = null;
-  let bdpaBasisCount = 0;
-  const totalSeqLen = currentSeqLen + numTokens;
-
-  const hasCache = state.kvCache?.hasGPUCache?.();
-
-  if (hasCache) {
-    if (state.kvCache.kvDtype === 'f16') {
-      // Use tensor dtype to determine if cast is needed
-      const kCasted = kTensor.dtype === 'f16' ? kTensor : await castF32ToF16(kTensor);
-      const vCasted = vTensor.dtype === 'f16' ? vTensor : await castF32ToF16(vTensor);
-
-      await state.kvCache.updateFromGPU(layerIdx, kCasted.buffer, vCasted.buffer, currentSeqLen, numTokens, tokenIds);
-
-      // Only release if we created new buffers
-      if (kTensor.dtype !== 'f16') releaseBuffer(kCasted.buffer);
-      if (vTensor.dtype !== 'f16') releaseBuffer(vCasted.buffer);
-    } else {
-      await state.kvCache.updateFromGPU(layerIdx, kTensor.buffer, vTensor.buffer, currentSeqLen, numTokens, tokenIds);
-    }
-    const gpuBuffers = state.kvCache.getGPUBuffers(layerIdx);
-    if (gpuBuffers?.layout === 'tiered') {
-      cachedKHot = gpuBuffers.hotKeysGPU;
-      cachedVHot = gpuBuffers.hotValuesGPU;
-      cachedKCold = gpuBuffers.coldKeysGPU;
-      cachedVCold = gpuBuffers.coldValuesGPU;
-      coldScalesK = gpuBuffers.coldScalesKGPU ?? null;
-      coldScalesV = gpuBuffers.coldScalesVGPU ?? null;
-      coldPackedStride = gpuBuffers.coldPackedStride ?? 0;
-      coldQuantMode = gpuBuffers.coldQuantMode ?? 'none';
-      hotLen = gpuBuffers.hotSeqLen ?? 0;
-      coldLen = gpuBuffers.coldSeqLen ?? 0;
-      hotStart = gpuBuffers.hotStart ?? 0;
-      hotWindow = gpuBuffers.hotWindow ?? 0;
-      coldPageTable = gpuBuffers.coldPageTableGPU ?? null;
-      coldPageSize = gpuBuffers.coldPageSize ?? state.kvCache.coldPageSize ?? 0;
-      kvLenForAttention = coldLen + hotLen;
-      kvLayout = 'tiered';
-    } else if (gpuBuffers?.layout === 'bdpa') {
-      kvLayout = 'bdpa';
-      kvLenForAttention = gpuBuffers.seqLen;
-      bdpaBasisK = gpuBuffers.basisGPU.k;
-      bdpaBasisV = gpuBuffers.basisGPU.v;
-      bdpaPagedK = gpuBuffers.pagedGPU.k;
-      bdpaPagedV = gpuBuffers.pagedGPU.v;
-      bdpaIndex = gpuBuffers.indexGPU;
-      bdpaBasisCount = gpuBuffers.numBasisVectors ?? state.kvCache.basisVocabSize;
-    } else {
-      cachedK = gpuBuffers.keysGPU;
-      cachedV = gpuBuffers.valuesGPU;
-      kvLenForAttention = gpuBuffers.seqLen;
-      kvPageTable = gpuBuffers.pageTableGPU ?? null;
-      kvPageSize = gpuBuffers.pageSize ?? state.kvCache.pageSize ?? 0;
-      if (state.kvCache instanceof SlidingWindowKVCache) {
-        kvLayout = 'ring';
-      } else if (state.kvCache.layout === 'paged') {
-        kvLayout = 'paged';
-      }
-    }
-
-    // Kernel step debug: KV cache state after update
-    if (isKernelDebugEnabled(layerIdx)) {
-      trace.kv(layerIdx, `KV cache updated: kvLen=${kvLenForAttention}, startPos=${currentSeqLen}, numTokens=${numTokens}`);
-      await dumpKVCache(( (state.kvCache)), layerIdx);
-    }
-  } else {
-    cachedK = kTensor.buffer;
-    cachedV = vTensor.buffer;
-    kvLenForAttention = numTokens;
-    startPosForMask = 0;
-  }
-
-  const attentionKernelVariant = selectRuleValue('inference', 'attention', 'attentionKernelVariant', {
-    kvLayout,
-    numTokens,
-    coldQuantMode,
-  });
-
-  if (attentionKernelVariant === 'contiguous' && kvLayout === 'tiered') {
-    kvLayout = 'contiguous';
-    cachedK = kTensor.buffer;
-    cachedV = vTensor.buffer;
-    kvLenForAttention = numTokens;
-    startPosForMask = 0;
-    cachedKHot = null;
-    cachedVHot = null;
-    cachedKCold = null;
-    cachedVCold = null;
-    coldQuantMode = 'none';
-  }
-
-  // Sliding window attention for specific layers
-  // The kernel now handles both causal AND sliding window masking together.
-  // We no longer need to disable causal masking for sliding layers.
-  const hasSlidingWindow = Number.isFinite(slidingWindow) && slidingWindow > 0;
-  const hasLayerTypes = Array.isArray(config.layerTypes);
-  const isLayerSliding = isSlidingLayerType(layerType) || (!hasLayerTypes && hasSlidingWindow);
-  const candidateSlidingWindow = isLayerSliding ? slidingWindow : null;
-  const effectiveSlidingWindow = candidateSlidingWindow;
-
-  const canWindow = hasCache && effectiveSlidingWindow;
-  if (attentionKernelVariant !== 'tiered' && attentionKernelVariant !== 'tieredQuant') {
-    if (canWindow && kvLenForAttention > effectiveSlidingWindow) {
-      kvLenForAttention = effectiveSlidingWindow;
-    }
-    if (hasCache && (kvLayout === 'ring' || (canWindow && kvLenForAttention < totalSeqLen))) {
-      kvStart = Math.max(0, totalSeqLen - kvLenForAttention);
-    }
-  }
-
-  if (kvLenForAttention <= 0) {
-    throw new Error(`Invalid kvLen ${kvLenForAttention} at layer ${layerIdx}`);
-  }
-
-  // 5. Attention (uses raw GPUBuffers)
-  // query_pre_attn_scalar is used as: scale = scalar^(-0.5) = 1/sqrt(scalar)
-  // When scalar equals headDim (e.g., 256): scale = 1/sqrt(256) = 1/16 (standard head_dim scaling)
-  const attnScale = queryPreAttnScalar ? 1.0 / Math.sqrt(queryPreAttnScalar) : 1.0 / Math.sqrt(headDim);
-  // Debug: log scale on layer 0
-  if (layerIdx === 0 && isPrefill) {
-    trace.attn(layerIdx, `Attention scale=${attnScale.toFixed(6)}, queryPreAttnScalar=${queryPreAttnScalar ?? 'undefined'}, headDim=${headDim}`);
-  }
-  // Wrap cached K/V in Tensors (dtype from cache or input tensor)
-  const cachedKDtype = selectRuleValue('inference', 'dtype', 'f16OrFallback', {
-    kvDtype: state.kvCache?.kvDtype,
-    fallback: kTensor.dtype,
-  });
-  const cachedVDtype = selectRuleValue('inference', 'dtype', 'f16OrFallback', {
-    kvDtype: state.kvCache?.kvDtype,
-    fallback: vTensor.dtype,
-  });
-  const cachedKTensor = attentionKernelVariant === 'tiered' || attentionKernelVariant === 'tieredQuant'
-    ? null
-    : createTensor(cachedK, cachedKDtype, [kvLenForAttention, numKVHeads * headDim], 'cached_K');
-  const cachedVTensor = attentionKernelVariant === 'tiered' || attentionKernelVariant === 'tieredQuant'
-    ? null
-    : createTensor(cachedV, cachedVDtype, [kvLenForAttention, numKVHeads * headDim], 'cached_V');
-
-  recordAttentionInputs(state, {
-    phase: isPrefill ? 'prefill' : 'decode',
-    layerIdx,
-    numTokens,
-    kvLen: kvLenForAttention,
-    numHeads,
-    numKVHeads,
-    headDim,
-    activationDtype: config.activationDtype ?? null,
-    inputDtype: input.dtype,
-    normedDtype: normed.dtype,
-    useF16Activations,
-    matmulOutputDtype,
-    kvCacheDtype: state.kvCache?.kvDtype ?? null,
-    cachedKDtype,
-    cachedVDtype,
-    qDtype: qTensor?.dtype ?? null,
-    kDtype: kTensor?.dtype ?? null,
-    vDtype: vTensor?.dtype ?? null,
-    useFusedQKV: usedFusedQKV,
-    kvStart,
-    kvLayout,
-    kvPageSize: kvLayout === 'tiered' ? (coldPageSize || null) : (kvPageSize || null),
-    hotLen: kvLayout === 'tiered' ? hotLen : null,
-    coldLen: kvLayout === 'tiered' ? coldLen : null,
-    hotWindow: kvLayout === 'tiered' ? hotWindow : null,
-    hotStart: kvLayout === 'tiered' ? hotStart : null,
-    coldQuantMode: kvLayout === 'tiered' ? coldQuantMode : null,
-  });
-
-  const attentionKernelRunners = {
-    bdpa: async () => {
-      const BDPA_MAX_HEAD_DIM = 256;
-      const BDPA_MAX_KV_LEN = 2048;
-      if (numTokens !== 1) {
-        throw new Error(`BDPA attention supports decode-only seqLen=1; got seqLen=${numTokens}.`);
-      }
-      if (headDim > BDPA_MAX_HEAD_DIM) {
-        throw new Error(
-          `BDPA attention kernel supports headDim <= ${BDPA_MAX_HEAD_DIM}; got headDim=${headDim}.`
-        );
-      }
-      if (kvLenForAttention > BDPA_MAX_KV_LEN) {
-        throw new Error(
-          `BDPA attention kernel supports kvLen <= ${BDPA_MAX_KV_LEN}; got kvLen=${kvLenForAttention}.`
-        );
-      }
-      const basisKDtype = 'f16';
-      const basisVDtype = 'f16';
-      const basisCount = Math.max(1, bdpaBasisCount);
-      const basisKTensor = createTensor(bdpaBasisK, basisKDtype, [basisCount, numKVHeads * headDim], 'bdpa_basis_k');
-      const basisVTensor = createTensor(bdpaBasisV, basisVDtype, [basisCount, numKVHeads * headDim], 'bdpa_basis_v');
-      let qForBDPA = qTensor;
-      if (qForBDPA.dtype !== 'f16') {
-        qForBDPA = await castF32ToF16(qTensor);
-      }
-      const output = await runAttentionBDPA(qForBDPA, basisKTensor, basisVTensor, bdpaPagedK, bdpaPagedV, bdpaIndex, numHeads, headDim, {
-        seqLen: numTokens,
-        kvLen: kvLenForAttention,
-        numKVHeads,
-        causal: causalForAttention,
-        startPos: startPosForMask,
-        layerIdx,
-        slidingWindow: effectiveSlidingWindow,
-        attnSoftcap,
-        scale: attnScale,
-        ropeCos: state.ropeFreqsCos,
-        ropeSin: state.ropeFreqsSin,
-      });
-      if (qForBDPA !== qTensor) {
-        releaseBuffer(qForBDPA.buffer);
-      }
-      return output;
-    },
-    tieredQuant: async () => {
-      let qForAttention = qTensor;
-      let qTemp = null;
-      if (coldQuantMode !== 'none' && qTensor.dtype !== 'f32') {
-        qForAttention = await castF16ToF32(qTensor);
-        qTemp = qForAttention;
-      }
-      const cachedHotKTensor = createTensor(cachedKHot, cachedKDtype, [hotLen, numKVHeads * headDim], 'cached_K_hot');
-      const cachedHotVTensor = createTensor(cachedVHot, cachedVDtype, [hotLen, numKVHeads * headDim], 'cached_V_hot');
-
-      if (coldQuantMode === 'none') {
-        throw new Error('Tiered quant attention requires cold quant mode.');
-      }
-      if (!coldScalesK || !coldScalesV) {
-        throw new Error('Tiered quant attention requires cold scale buffers.');
-      }
-
-      const output = await runAttentionTieredQuant(qForAttention, cachedHotKTensor, cachedHotVTensor, cachedKCold, cachedVCold, coldScalesK, coldScalesV, numHeads, headDim, {
-        seqLen: numTokens,
-        coldLen,
-        hotLen,
-        numKVHeads,
-        causal: causalForAttention,
-        startPos: startPosForMask,
-        slidingWindow: effectiveSlidingWindow ?? 0,
-        attnSoftcap,
-        scale: attnScale,
-        hotWindow,
-        hotStart,
-        packedStride: coldPackedStride,
-        mode: coldQuantMode,
-      });
-
-      if (qTemp) {
-        releaseBuffer(qTemp.buffer);
-      }
-      return output;
-    },
-    tiered: async () => {
-      const qForAttention = qTensor;
-      const cachedHotKTensor = createTensor(cachedKHot, cachedKDtype, [hotLen, numKVHeads * headDim], 'cached_K_hot');
-      const cachedHotVTensor = createTensor(cachedVHot, cachedVDtype, [hotLen, numKVHeads * headDim], 'cached_V_hot');
-      const cachedColdKTensor = createTensor(cachedKCold, cachedKDtype, [coldLen, numKVHeads * headDim], 'cached_K_cold');
-      const cachedColdVTensor = createTensor(cachedVCold, cachedVDtype, [coldLen, numKVHeads * headDim], 'cached_V_cold');
-      return runAttentionTiered(qForAttention, cachedHotKTensor, cachedHotVTensor, cachedColdKTensor, cachedColdVTensor, numHeads, headDim, {
-        seqLen: numTokens,
-        coldLen,
-        hotLen,
-        numKVHeads,
-        causal: causalForAttention,
-        startPos: startPosForMask,
-        slidingWindow: effectiveSlidingWindow ?? 0,
-        attnSoftcap,
-        scale: attnScale,
-        hotWindow,
-        hotStart,
-        coldPageTable,
-        coldPageSize,
-        coldLayout: coldPageTable ? 2 : 0,
-        hotLayout: hotWindow > 0 ? 1 : 0,
-      });
-    },
-    contiguous: async () => runAttention(qTensor, cachedKTensor, cachedVTensor, null, numHeads, headDim, {
-      seqLen: numTokens,
-      kvLen: kvLenForAttention,
-      numKVHeads,
-      causal: causalForAttention,
-      startPos: startPosForMask,
-      layerIdx,
-      slidingWindow: effectiveSlidingWindow,
-      attnSoftcap,
-      scale: attnScale,
-      kvStart,
-      kvLayout,
-      kvPageTable,
-      kvPageSize,
-      kernelPath,
-    }),
-  };
-
-  const runAttentionKernel = attentionKernelRunners[attentionKernelVariant];
-  if (!runAttentionKernel) {
-    throw new Error(`Unsupported attention kernel variant "${attentionKernelVariant}" at layer ${layerIdx}`);
-  }
-
-  attnOutput = await runAttentionKernel();
-
-  // Trace attention output
-  if (kernelTrace.enabled) {
-    await traceStep('attention', `L${layerIdx}.attention`, layerIdx, attnOutput.buffer, [numTokens, numHeads * headDim]);
-  }
-
-  // Kernel step debug: attention output
-  if (isKernelDebugEnabled(layerIdx)) {
-    logKernelStep('attention', { layerIdx, label: `seqLen=${numTokens} kvLen=${kvLenForAttention}` });
-    await dumpTokenVector(attnOutput.buffer, 'attn_out', {
-      layerIdx,
-      tokenIdx: Math.max(0, numTokens - 1),
-      rowSize: numHeads * headDim,
-      dtype: attnOutput.dtype,
-    });
-  }
-
-  // Debug: attention output for configured layers
-  if (isPrefill && shouldDebugLayer(layerIdx, debugFlags.debugLayers) && !markStageLogged(layerIdx, 'attn_out', debugFlags) && debugCheckBuffer) {
-    await debugCheckBuffer(attnOutput.buffer, `L${layerIdx} attention output (before o_proj, GPU)`, numTokens, numHeads * headDim);
-  }
-
-  attnForProjection = attnOutput;
-  if (qGateTensor) {
-    attnForProjection = await runSiLU(attnOutput, {
-      size: numTokens * numHeads * headDim,
-      gate: qGateTensor,
-      gateActivation: 'sigmoid',
-      inputActivation: 'identity',
-      swigluLimit: null,
-    });
-    releaseBuffer(attnOutput.buffer);
-  }
-
-  // 6. Output projection (with optional fused residual for decode)
-  
-  output = null;
-  let residualFused = false;
-  let oProjInput = attnForProjection;
-  oProjInputTemp = null;
-  if (layerWeights.oProj && getWeightBuffer) {
-    ({ oProjInput, oProjInputTemp } = await prepareAttentionProjectionInput(
-      attnForProjection,
-      matmulOutputDtype,
-      castF32ToF16
-    ));
-    const oProjBuf = getWeightBuffer(layerWeights.oProj, 'o_proj');
-    const loraO = getLoRAModule(lora, layerIdx, 'o_proj');
-
-    // Use fused o_proj + residual for decode when possible
-    // Note: dtype from WeightBuffer metadata (buffer-dtypes WeakMap removed)
-    const oProjDtype = getWeightDtype(oProjBuf);
-    const canUseFused = selectRuleValue('inference', 'attention', 'useFusedOProjResidual', {
-      allowFusedResidual: shouldUseFusedMatmulResidual(numTokens),
-      hasResidual: Boolean(residualTensor),
-      residualMatches: Boolean(residualTensor && residualTensor.dtype === oProjInput.dtype),
-      attnIsF32: oProjInput.dtype === 'f32',
-      attnIsF16: oProjInput.dtype === 'f16',
-      hasLoRA: Boolean(loraO),
-      oProjIsF16: oProjDtype === 'f16',
-    });  // GEMV kernel expects f16 weights
-
-    if (canUseFused && residualTensor) {
-      // FUSED PATH: o_proj matmul + residual add in one dispatch
-      output = await runMatmulResidualFused(oProjInput, oProjBuf, residualTensor, {
-        N: hiddenSize,
-        K: numHeads * headDim,
-      });
-      residualFused = true;
-
-      if (layerIdx === 0 && !isPrefill) {
-        trace.attn(layerIdx, `Using fused o_proj+residual path`);
-      }
-    } else {
-      // STANDARD PATH: o_proj matmul only (residual will be added by layer.ts)
-      output = await runMatmul(oProjInput, oProjBuf, numTokens, hiddenSize, numHeads * headDim, {
-        transposeB: 'auto',
-        role: 'o_proj',
-        layerIdx,
-        kernelPath,
-        outputDtype: matmulOutputDtype,
-      });
-    }
-    // Release temporary buffer if we created it (original was not already on GPU)
-    if (!(layerWeights.oProj instanceof GPUBuffer) && !isWeightBuffer(layerWeights.oProj)) {
-      releaseBuffer(isWeightBuffer(oProjBuf) ? oProjBuf.buffer : oProjBuf);
-    }
-
-    // Trace output projection
-    if (kernelTrace.enabled) {
-      await traceStep('matmul', `L${layerIdx}.o_proj${residualFused ? '+residual' : ''}`, layerIdx, output.buffer, [numTokens, hiddenSize]);
-    }
-
-    // Kernel step debug: output projection
-    if (isKernelDebugEnabled(layerIdx)) {
-      logKernelStep('matmul', { layerIdx, label: residualFused ? 'O_proj+residual' : 'O_proj', M: numTokens, N: hiddenSize, K: numHeads * headDim });
-      await dumpTokenVector(output.buffer, 'o_proj_out', {
-        layerIdx,
-        tokenIdx: Math.max(0, numTokens - 1),
-        rowSize: hiddenSize,
-        dtype: output.dtype,
-      });
-    }
-
-  } else {
-    output = oProjInput;
-  }
-
-  // Apply LoRA to output projection if present (only if not using fused path)
-  if (!residualFused) {
-    const loraO = getLoRAModule(lora, layerIdx, 'o_proj');
-    if (loraO && getWeightBuffer) {
-      const combined = await applyLoRA(
-        oProjInput,
-        output,
-        loraO,
-        { M: numTokens, N: hiddenSize, K: numHeads * headDim },
-        getWeightBuffer,
-        undefined,
-        { kernelPath }
-      );
-      if (combined.buffer !== output.buffer) {
-        releaseBuffer(output.buffer);
-        output = combined;
-      }
-    }
-  }
-
-  if (oProjInputTemp) {
-    releaseBuffer(oProjInputTemp.buffer);
-  }
-
-  // Debug: o_proj output for configured layers
-  if (isPrefill && shouldDebugLayer(layerIdx, debugFlags.debugLayers) && !markStageLogged(layerIdx, 'o_proj', debugFlags) && debugCheckBuffer) {
-    await debugCheckBuffer(output.buffer, `L${layerIdx} attention output (after o_proj, GPU)`, numTokens, hiddenSize);
-  }
-
-  finalOutput = output;
-  
-  const buffersToRelease = [];
-  if (output.buffer !== attnForProjection.buffer) {
-    buffersToRelease.push(attnForProjection.buffer);
-  }
-
-  if (wantsF16Output && output.dtype !== 'f16') {
-    const f16Output = await castF32ToF16(output);
-    buffersToRelease.push(output.buffer);
-    finalOutput = f16Output;
-  }
-
-  // Cleanup
-  releaseBuffer(qTensor.buffer);
-  if (qGateTensor) {
-    releaseBuffer(qGateTensor.buffer);
-  }
-  releaseBuffer(kTensor.buffer);
-  releaseBuffer(vTensor.buffer);
-  for (const buffer of buffersToRelease) {
-    releaseBuffer(buffer);
-  }
-
-  return { output: finalOutput, residualFused };
-  } catch (error) {
-    const released = new Set();
-    const releaseOnce = (buffer) => {
-      if (!buffer || released.has(buffer)) return;
-      released.add(buffer);
-      releaseBuffer(buffer);
-    };
-    if (finalOutput?.buffer && finalOutput.buffer !== output?.buffer) {
-      releaseOnce(finalOutput.buffer);
-    }
-    if (output?.buffer && output.buffer !== attnForProjection?.buffer) {
-      releaseOnce(output.buffer);
-    }
-    if (oProjInputTemp?.buffer) {
-      releaseOnce(oProjInputTemp.buffer);
-    }
-    if (attnForProjection?.buffer && attnForProjection.buffer !== attnOutput?.buffer) {
-      releaseOnce(attnForProjection.buffer);
-    }
-    if (attnOutput?.buffer) {
-      releaseOnce(attnOutput.buffer);
-    }
-    if (qGateTensor?.buffer) {
-      releaseOnce(qGateTensor.buffer);
-    }
-    if (qTensor?.buffer) {
-      releaseOnce(qTensor.buffer);
-    }
-    if (kTensor?.buffer) {
-      releaseOnce(kTensor.buffer);
-    }
-    if (vTensor?.buffer) {
-      releaseOnce(vTensor.buffer);
-    }
-    if (normed?.buffer && normed.buffer !== attentionInput?.buffer) {
-      releaseOnce(normed.buffer);
-    }
-    if (attentionInputTemp && attentionInput?.buffer) {
-      releaseOnce(attentionInput.buffer);
-    }
-    throw error;
-  }
-}