diff --git a/src/DotLLM.Cuda/Architectures/CudaMtpState.cs b/src/DotLLM.Cuda/Architectures/CudaMtpState.cs new file mode 100644 index 00000000..91472687 --- /dev/null +++ b/src/DotLLM.Cuda/Architectures/CudaMtpState.cs @@ -0,0 +1,202 @@ +using DotLLM.Core.Models; +using DotLLM.Cuda.Interop; + +namespace DotLLM.Cuda.Architectures; + +/// +/// CUDA implementation: the MTP head's own tiny KV-cache (sized for just +/// the trailing MTP block's attention, not the trunk) plus the pending-hidden-state handoff row +/// and the captured-rows buffer a verify-phase Forward call populates. Mirrors +/// (see issue #253 and +/// for the overall design) with device-resident K/V cache and pending-hidden buffers instead of +/// host NativeMemory — the MTP head's autoregressive draft loop +/// () runs entirely on-device between +/// rounds, so the pending-hidden handoff never round-trips through host memory except at +/// (once per speculation round, from the host-resident captured +/// rows a verify-phase forward D2H-copies back). +/// +public sealed class CudaMtpState : IMtpState, IDisposable +{ + private readonly int _hiddenSize; + private readonly int _numKvHeads; + private readonly int _headDim; + private readonly int _maxSteps; + private readonly int _kvStride; // numKvHeads * headDim + + private nint _keyCacheDevice; // [maxSteps, kvStride] f32, device-resident + private nint _valueCacheDevice; // [maxSteps, kvStride] f32, device-resident + private nint _pendingHiddenDevice; // [hiddenSize] f32, device-resident — seed for the next ForwardMtp call + + private float[] _capturedRows = []; // host [rowCount, hiddenSize], grown on demand + private int _capturedRowCount; + + private int _currentLength; + private bool _disposed; + + /// Max autoregressive MTP steps this state's KV-cache can hold before needing a reset/rollback. + public int MaxSteps => _maxSteps; + + /// K/V stride (numKvHeads * headDim) — the per-step row width of the device K/V cache. + public int KvStride => _kvStride; + + /// + public int CurrentLength => _currentLength; + + /// + public ReadOnlySpan CapturedHiddenRows => _capturedRows.AsSpan(0, _capturedRowCount * _hiddenSize); + + /// + public int CapturedRowCount => _capturedRowCount; + + /// + public int HiddenSize => _hiddenSize; + + /// + /// Creates a fresh MTP state. All device buffers are zero-initialised (an empty MTP KV-cache + /// and a zero pending-hidden vector — the latter is always overwritten by + /// before the first ForwardMtp call in normal use). + /// + /// Model hidden size (matches ). + /// KV head count for the MTP block's own attention. + /// Per-head dimension for the MTP block's own attention. + /// Maximum autoregressive MTP draft steps to size the KV-cache for (typically the max candidate count K). + public CudaMtpState(int hiddenSize, int numKvHeads, int headDim, int maxSteps) + { + if (hiddenSize <= 0) throw new ArgumentOutOfRangeException(nameof(hiddenSize)); + if (numKvHeads <= 0) throw new ArgumentOutOfRangeException(nameof(numKvHeads)); + if (headDim <= 0) throw new ArgumentOutOfRangeException(nameof(headDim)); + if (maxSteps <= 0) throw new ArgumentOutOfRangeException(nameof(maxSteps)); + + _hiddenSize = hiddenSize; + _numKvHeads = numKvHeads; + _headDim = headDim; + _maxSteps = maxSteps; + _kvStride = numKvHeads * headDim; + + long kvBytes = (long)maxSteps * _kvStride * sizeof(float); + CudaDriverApi.cuMemAlloc_v2(out _keyCacheDevice, (nuint)kvBytes).ThrowOnError(); + CudaDriverApi.cuMemAlloc_v2(out _valueCacheDevice, (nuint)kvBytes).ThrowOnError(); + CudaDriverApi.cuMemsetD8_v2(_keyCacheDevice, 0, (nuint)kvBytes).ThrowOnError(); + CudaDriverApi.cuMemsetD8_v2(_valueCacheDevice, 0, (nuint)kvBytes).ThrowOnError(); + + long hiddenBytes = (long)hiddenSize * sizeof(float); + CudaDriverApi.cuMemAlloc_v2(out _pendingHiddenDevice, (nuint)hiddenBytes).ThrowOnError(); + CudaDriverApi.cuMemsetD8_v2(_pendingHiddenDevice, 0, (nuint)hiddenBytes).ThrowOnError(); + } + + /// Device pointer to the pending-hidden vector ([hiddenSize] f32) that seeds the next ForwardMtp call. + internal nint PendingHiddenDevicePtr + { + get { ThrowIfDisposed(); return _pendingHiddenDevice; } + } + + /// Device pointer to the key-cache buffer, shape [maxSteps, kvStride] row-major. + internal nint KeyCacheDevicePtr + { + get { ThrowIfDisposed(); return _keyCacheDevice; } + } + + /// Device pointer to the value-cache buffer, shape [maxSteps, kvStride] row-major. + internal nint ValueCacheDevicePtr + { + get { ThrowIfDisposed(); return _valueCacheDevice; } + } + + /// Device pointer to the key-cache row for MTP step (0-based), shape [kvStride]. + internal nint GetKeyRowDevicePtr(int step) + { + ThrowIfDisposed(); + if ((uint)step >= (uint)_maxSteps) throw new ArgumentOutOfRangeException(nameof(step)); + return _keyCacheDevice + (nint)((long)step * _kvStride * sizeof(float)); + } + + /// Device pointer to the value-cache row for MTP step (0-based), shape [kvStride]. + internal nint GetValueRowDevicePtr(int step) + { + ThrowIfDisposed(); + if ((uint)step >= (uint)_maxSteps) throw new ArgumentOutOfRangeException(nameof(step)); + return _valueCacheDevice + (nint)((long)step * _kvStride * sizeof(float)); + } + + /// Advances the MTP KV-cache length by one step after a successful ForwardMtp call. + internal void Advance() + { + ThrowIfDisposed(); + if (_currentLength >= _maxSteps) + throw new InvalidOperationException( + $"CudaMtpState KV-cache exhausted: {_currentLength} steps already advanced against a " + + $"MaxSteps={_maxSteps} cache. Size the state for at least numCandidates steps."); + _currentLength++; + } + + /// + public void Rollback(int length) + { + ThrowIfDisposed(); + if (length < 0 || length > _currentLength) + throw new ArgumentOutOfRangeException(nameof(length)); + _currentLength = length; + } + + /// + /// Called by an MTP-supporting model's verify-phase Forward overload to D2H-copy the + /// captured pre-final-norm hidden rows (a device buffer, e.g. _state.HiddenState) into + /// this state's host-resident captured-rows buffer for this state's next round. The caller is + /// responsible for ensuring the source device buffer's writes have completed (stream + /// synchronized) before calling this — cuMemcpyDtoH_v2 does not implicitly wait for a + /// non-default stream's queued work. + /// + /// Device pointer to row-major [rowCount, hiddenSize] hidden state rows. + /// Number of rows to capture. + internal unsafe void SetCapturedRowsFromDevice(nint deviceHiddenState, int rowCount) + { + ThrowIfDisposed(); + int needed = rowCount * _hiddenSize; + if (_capturedRows.Length < needed) + _capturedRows = new float[needed]; + fixed (float* p = _capturedRows) + { + CudaDriverApi.cuMemcpyDtoH_v2((nint)p, deviceHiddenState, + (nuint)((long)needed * sizeof(float))).ThrowOnError(); + } + _capturedRowCount = rowCount; + } + + /// + public unsafe void SeedFromCapturedRow(int rowIndex) + { + ThrowIfDisposed(); + if ((uint)rowIndex >= (uint)_capturedRowCount) + throw new ArgumentOutOfRangeException(nameof(rowIndex), + $"rowIndex {rowIndex} out of range [0, {_capturedRowCount}) — CapturedHiddenRows was not populated " + + "by a verify-phase Forward call, or has fewer rows than expected."); + + // H2D: host-captured row -> device pending-hidden buffer. The very next ForwardMtp call + // consumes _pendingHiddenDevice directly on-device (RMSNorm), so no further round-trip is + // needed until the NEXT round's SeedFromCapturedRow. + fixed (float* p = &_capturedRows[rowIndex * _hiddenSize]) + { + CudaDriverApi.cuMemcpyHtoD_v2(_pendingHiddenDevice, (nint)p, + (nuint)((long)_hiddenSize * sizeof(float))).ThrowOnError(); + } + } + + /// Total bytes allocated for this state's own KV-cache + pending-hidden buffer (device memory). + public long AllocatedBytes => (2L * _maxSteps * _kvStride + _hiddenSize) * sizeof(float); + + private void ThrowIfDisposed() + { + if (_disposed) throw new ObjectDisposedException(nameof(CudaMtpState)); + } + + /// + public void Dispose() + { + if (_disposed) return; + if (_keyCacheDevice != 0) { CudaDriverApi.cuMemFree_v2(_keyCacheDevice); _keyCacheDevice = 0; } + if (_valueCacheDevice != 0) { CudaDriverApi.cuMemFree_v2(_valueCacheDevice); _valueCacheDevice = 0; } + if (_pendingHiddenDevice != 0) { CudaDriverApi.cuMemFree_v2(_pendingHiddenDevice); _pendingHiddenDevice = 0; } + _disposed = true; + GC.SuppressFinalize(this); + } +} diff --git a/src/DotLLM.Cuda/Architectures/CudaQwen3HybridDenseTransformerModel.cs b/src/DotLLM.Cuda/Architectures/CudaQwen3HybridDenseTransformerModel.cs index ffdf84f8..2c497aa1 100644 --- a/src/DotLLM.Cuda/Architectures/CudaQwen3HybridDenseTransformerModel.cs +++ b/src/DotLLM.Cuda/Architectures/CudaQwen3HybridDenseTransformerModel.cs @@ -1,6 +1,7 @@ using System.Runtime.CompilerServices; using DotLLM.Core.Attention; using DotLLM.Core.Configuration; +using DotLLM.Core.Lora; using DotLLM.Core.Models; using DotLLM.Core.Tensors; using DotLLM.Cpu.Kernels; @@ -165,11 +166,20 @@ public sealed unsafe class CudaQwen3HybridDenseTransformerModel : IModel private long _attnMmaDecodeQF16ElemsAllocated; private readonly DotLLM.Cuda.CudaAttentionMmaDecodeGqaSplit _mmaDecodeGqaSplit; + // Multi-Token Prediction (MTP / "NextN") head — issue #253. Null for every GGUF without a + // nextn.* tensor group, which is the overwhelming majority: every other field and code path + // in this class is completely unaffected by MTP being absent. Mirrors the CPU host's + // Qwen3HybridDenseTransformerModel._mtpHead (DotLLM.Models.Architectures.MtpHeadWeights). + private readonly CudaMtpHeadWeights? _mtpHead; + private bool _disposed; /// public ModelConfig Config { get; } + /// + public bool SupportsMtp => _mtpHead is not null; + /// public long ComputeMemoryBytes => _state.AllocatedBytes + _gdnCache.AllocatedBytes; @@ -195,6 +205,33 @@ public sealed unsafe class CudaQwen3HybridDenseTransformerModel : IModel /// public CudaHybridKvCacheHandle CreateKvCache(int maxSeqLen) => new(maxSeqLen); + /// + /// + /// Sized for the MTP head's own attention ('s standard head count/dim — + /// the MTP block is a normal full-attention layer, see ), with a + /// device-resident KV-cache deep enough for autoregressive + /// draft steps. Mirrors the CPU host's Qwen3HybridDenseTransformerModel.CreateMtpState. + /// + public IMtpState? CreateMtpState() + { + if (_mtpHead is null) + return null; + + return new CudaMtpState( + hiddenSize: Config.HiddenSize, + numKvHeads: _mtpHead.Value.Layer.FullAttn!.Value.NumKvHeads, + headDim: Config.HeadDim, + maxSteps: MtpDefaultMaxDraftSteps); + } + + /// + /// Default MTP KV-cache depth when a caller doesn't need a specific candidate count K up + /// front. Callers that know K in advance (e.g. an MTP self-speculative decoder, see issue + /// #253) can size their own directly instead of going through + /// . + /// + public const int MtpDefaultMaxDraftSteps = 16; + private CudaQwen3HybridDenseTransformerModel( ModelConfig config, GgufFile? gguf, @@ -209,11 +246,13 @@ private CudaQwen3HybridDenseTransformerModel( CudaQwen3HybridDenseForwardState state, CudaGdnStateCache gdnCache, CudaStream stream, CudaCublasHandle cublas, CudaContext context, CudaKernels kernels, int deviceId, - nint dequantScratchDevice) + nint dequantScratchDevice, + CudaMtpHeadWeights? mtpHead = null) { Config = config; _gguf = gguf; _layers = layers; + _mtpHead = mtpHead; _tokenEmbedDevice = tokenEmbedDevice; _tokenEmbedQt = tokenEmbedQt; _embedDataBase = embedDataBase; @@ -393,6 +432,12 @@ public static CudaQwen3HybridDenseTransformerModel LoadFromGguf( : -1; } + // MTP (issue #253): load the trailing NextN head when the GGUF carries one. Zero behavior + // change for every other checkpoint — LoadMtpHeadIfPresent returns null unless + // config.NextnPredictLayers > 0 AND the nextn.* tensors are actually present. Mirrors the + // CPU host's Qwen3HybridDenseTransformerModel.LoadMtpHeadIfPresent tensor layout exactly. + CudaMtpHeadWeights? mtpHead = LoadMtpHeadIfPresent(dataBase, tensors, config, ref maxTileFloats); + maxTileFloats = Math.Max(maxTileFloats, (long)outputOutputDim * outputInputDim); nint dequantScratchDevice = AllocDevice(maxTileFloats * sizeof(ushort)); @@ -423,7 +468,7 @@ public static CudaQwen3HybridDenseTransformerModel LoadFromGguf( kvSlotForLayer, attentionLayerCount, ropeTheta, ropeDim, state, gdnCache, stream, cublas, context, kernels, deviceId, - dequantScratchDevice); + dequantScratchDevice, mtpHead); } // ────────────────────────────────────────────────────────────────────── @@ -606,6 +651,149 @@ private static DeviceFullAttn LoadFullAttnLayerDevice( }; } + /// + /// Loads the trailing Multi-Token Prediction (MTP / "NextN") head when the GGUF has one + /// (issue #253), or returns for a checkpoint without MTP — the + /// overwhelming majority of GGUFs, completely unaffected by this method. Mirrors + /// DotLLM.Models.Architectures.Qwen3HybridDenseTransformerModel.LoadMtpHeadIfPresent + /// (CPU) exactly: same tensor names, same "hparam without tensors ⇒ no MTP head" tolerance, + /// same nextn_predict_layers == 1-only restriction. The MTP block's own attn+ffn + /// tensors reuse / the same dense-FFN upload pattern + /// uses — it is structurally a normal full-attention decoder + /// layer, just appended at raw GGUF block index config.NumLayers (trunk layers occupy + /// [0, NumLayers), since GgufModelConfigExtractor already subtracted + /// nextn_predict_layers back out of the raw block_count). + /// + private static CudaMtpHeadWeights? LoadMtpHeadIfPresent( + nint dataBase, + IReadOnlyDictionary tensors, + ModelConfig config, + ref long maxTileFloats) + { + if (config.NextnPredictLayers <= 0) + return null; + + if (config.NextnPredictLayers != 1) + throw new NotSupportedException( + $"Qwen3HybridDense MTP only supports a single trailing MTP block " + + $"(nextn_predict_layers=1); got {config.NextnPredictLayers}. Matches llama.cpp's own " + + "current QWEN35 MTP assertion (issue #253 scope: Qwen3.6, not a future multi-head variant)."); + + int mtpBlk = config.NumLayers; // trunk occupies [0, NumLayers); MTP is appended right after + string prefix = $"blk.{mtpBlk}"; + + // The hparam can be set without the tensors actually being present (e.g. a trunk-only + // GGUF someone hand-edited the metadata on) — treat that as "no MTP head", not an error, + // to keep the zero-behavior-change guarantee unconditional. + if (!tensors.ContainsKey($"{prefix}.nextn.eh_proj.weight")) + return null; + + int hiddenSize = config.HiddenSize; + + // The MTP block's own attn+ffn tensors use the exact same naming/shapes as any other + // full-attention Qwen3HybridDense layer — reuse the trunk loaders directly. + var attnNormDesc = tensors[$"{prefix}.attn_norm.weight"]; + nint attnNormDevice = UploadF32Tensor(dataBase, attnNormDesc, hiddenSize); + var postNormDesc = tensors[$"{prefix}.post_attention_norm.weight"]; + nint postAttnNormDevice = UploadF32Tensor(dataBase, postNormDesc, hiddenSize); + + DeviceFullAttn attnDev = LoadFullAttnLayerDevice(prefix, dataBase, tensors, config, + config.NumKvHeads, ref maxTileFloats); + + var gateDesc = tensors[$"{prefix}.ffn_gate.weight"]; + var upDesc = tensors[$"{prefix}.ffn_up.weight"]; + var downDesc = tensors[$"{prefix}.ffn_down.weight"]; + nint gateDevice = UploadRawTensor(dataBase, gateDesc); + nint upDevice = UploadRawTensor(dataBase, upDesc); + nint downDevice = UploadRawTensor(dataBase, downDesc); + UpdateMaxTile(ref maxTileFloats, (long)gateDesc.Shape[0] * gateDesc.Shape[1]); + UpdateMaxTile(ref maxTileFloats, (long)upDesc.Shape[0] * upDesc.Shape[1]); + UpdateMaxTile(ref maxTileFloats, (long)downDesc.Shape[0] * downDesc.Shape[1]); + + var layer = new DeviceLayer + { + AttnNormWeightDevice = attnNormDevice, + PostAttnNormWeightDevice = postAttnNormDevice, + Gdn = null, + FullAttn = attnDev, + + GateWeight = gateDevice, GateQt = gateDesc.QuantizationType, + GateInputDim = gateDesc.Shape[0], GateOutputDim = gateDesc.Shape[1], + + UpWeight = upDevice, UpQt = upDesc.QuantizationType, + UpInputDim = upDesc.Shape[0], UpOutputDim = upDesc.Shape[1], + + DownWeight = downDevice, DownQt = downDesc.QuantizationType, + DownInputDim = downDesc.Shape[0], DownOutputDim = downDesc.Shape[1], + }; + + var ehProjDesc = tensors[$"{prefix}.nextn.eh_proj.weight"]; + nint ehProjDevice = UploadRawTensor(dataBase, ehProjDesc); + UpdateMaxTile(ref maxTileFloats, (long)ehProjDesc.Shape[0] * ehProjDesc.Shape[1]); + + nint enormDevice = UploadF32Tensor(dataBase, tensors[$"{prefix}.nextn.enorm.weight"], hiddenSize); + nint hnormDevice = UploadF32Tensor(dataBase, tensors[$"{prefix}.nextn.hnorm.weight"], hiddenSize); + + // Optional nextn.embed_tokens: host-mmap pointer (NOT uploaded to device), mirroring the + // trunk's own _embedDataBase convention — MTP embeds one token per ForwardMtp call via a + // host-side dequant + tiny H2D copy, exactly like the trunk's per-token embedding lookup + // in Forward(). Falls back to the trunk's own token_embd.weight (via _embedDataBase/ + // _embedDataOffset/_embedRowBytes) when absent. + nint? embedTokensHostBase = null; + ulong embedTokensDataOffset = 0; + long embedTokensRowBytes = 0; + QuantizationType embedTokensQt = default; + if (tensors.TryGetValue($"{prefix}.nextn.embed_tokens.weight", out var embedDesc)) + { + embedTokensHostBase = dataBase; + embedTokensDataOffset = embedDesc.DataOffset; + embedTokensRowBytes = Dequantize.RowByteSize(hiddenSize, embedDesc.QuantizationType); + embedTokensQt = embedDesc.QuantizationType; + } + + // Optional nextn.shared_head_head / nextn.shared_head_norm: device-resident (feed a GEMM / + // RMSNorm kernel respectively), mirroring _outputDevice/_outputNormDevice. Fall back to the + // trunk's own lm_head/output_norm when absent. + nint? sharedHeadHeadDevice = null; + QuantizationType sharedHeadHeadQt = default; + int sharedHeadHeadInputDim = 0, sharedHeadHeadOutputDim = 0; + if (tensors.TryGetValue($"{prefix}.nextn.shared_head_head.weight", out var sharedHeadDesc)) + { + sharedHeadHeadDevice = UploadRawTensor(dataBase, sharedHeadDesc); + sharedHeadHeadQt = sharedHeadDesc.QuantizationType; + sharedHeadHeadInputDim = sharedHeadDesc.Shape[0]; + sharedHeadHeadOutputDim = sharedHeadDesc.Shape[1]; + UpdateMaxTile(ref maxTileFloats, (long)sharedHeadDesc.Shape[0] * sharedHeadDesc.Shape[1]); + } + + nint? sharedHeadNormDevice = tensors.TryGetValue($"{prefix}.nextn.shared_head_norm.weight", out var shnDesc) + ? UploadF32Tensor(dataBase, shnDesc, hiddenSize) + : null; + + return new CudaMtpHeadWeights + { + Layer = layer, + + EhProjDevice = ehProjDevice, EhProjQt = ehProjDesc.QuantizationType, + EhProjInputDim = ehProjDesc.Shape[0], EhProjOutputDim = ehProjDesc.Shape[1], + + EnormDevice = enormDevice, + HnormDevice = hnormDevice, + + EmbedTokensHostBase = embedTokensHostBase, + EmbedTokensDataOffset = embedTokensDataOffset, + EmbedTokensRowBytes = embedTokensRowBytes, + EmbedTokensQt = embedTokensQt, + + SharedHeadHeadDevice = sharedHeadHeadDevice, + SharedHeadHeadQt = sharedHeadHeadQt, + SharedHeadHeadInputDim = sharedHeadHeadInputDim, + SharedHeadHeadOutputDim = sharedHeadHeadOutputDim, + + SharedHeadNormDevice = sharedHeadNormDevice, + }; + } + // ────────────────────────────────────────────────────────────────────── // Forward dispatch // ────────────────────────────────────────────────────────────────────── @@ -620,9 +808,35 @@ public ITensor Forward(ReadOnlySpan tokenIds, ReadOnlySpan positions, => Forward(tokenIds, positions, deviceId, kvCache, lastTokenLogitsOnly: false); /// - [SkipLocalsInit] public ITensor Forward(ReadOnlySpan tokenIds, ReadOnlySpan positions, int deviceId, IKvCache? kvCache, bool lastTokenLogitsOnly) + => ForwardCore(tokenIds, positions, deviceId, kvCache, lastTokenLogitsOnly, mtpCapture: null); + + /// + /// + /// MTP (issue #253): when is a on a + /// model with true, this call additionally D2H-copies the trunk's + /// pre-final-norm hidden state for every position in into it — see + /// the capture point inside . The returned logits are byte-identical + /// to calling without . is accepted (per + /// the contract) but has no effect — this model has no LoRA support, same + /// as every other overload here. + /// + public ITensor Forward(ReadOnlySpan tokenIds, ReadOnlySpan positions, + int deviceId, IKvCache? kvCache, ILoraAdapter? adapter, IMtpState? mtpState) + => ForwardCore(tokenIds, positions, deviceId, kvCache, lastTokenLogitsOnly: false, + mtpCapture: mtpState as CudaMtpState); + + /// + /// Core forward-pass implementation shared by every public Forward overload above. + /// is non-null only from the MTP-aware overload — see that + /// overload's remarks and the capture point below, right before the final RMSNorm overwrites + /// in place. + /// + [SkipLocalsInit] + private ITensor ForwardCore(ReadOnlySpan tokenIds, ReadOnlySpan positions, + int deviceId, IKvCache? kvCache, bool lastTokenLogitsOnly, + CudaMtpState? mtpCapture) { ObjectDisposedException.ThrowIf(_disposed, this); @@ -697,6 +911,22 @@ public ITensor Forward(ReadOnlySpan tokenIds, ReadOnlySpan positions, if (DebugTrace) { _stream.Synchronize(); Console.Error.WriteLine("[hybrid-debug] all layers done, starting lm_head"); Console.Error.Flush(); LogVram("before lm-head"); } ProfStart(); + // MTP (issue #253): capture the pre-final-norm hidden state for every position, one row + // per input token, BEFORE the final RMSNorm below overwrites _state.HiddenState in place. + // This is the exact quantity llama.cpp's MTP head consumes (`h_pre_norm` / + // `llama_get_embeddings_pre_norm`) — a pure side effect that never changes the logits this + // call returns. The MTP-aware Forward overload always passes lastTokenLogitsOnly=false, so + // _state.HiddenState always holds all `seqLen` valid rows here regardless of logitsRows + // below. cuMemcpyDtoH_v2 does not implicitly wait for this model's non-default _stream, so + // synchronize first — the LM-head projection below queues fresh work after this point, so + // this sync does not skip/reorder anything, only adds one extra host-blocking wait on the + // (low-frequency, K+1-token-per-round) MTP verify/catchup path. + if (mtpCapture is not null) + { + _stream.Synchronize(); + mtpCapture.SetCapturedRowsFromDevice(_state.HiddenState, seqLen); + } + // Issue #185: only compute/copy the LAST token's logits when the caller has explicitly // opted in via lastTokenLogitsOnly (e.g. BenchRunner's untimed prefill / --depth context // extension, which only ever reads the last row via argmax). This must NOT be inferred @@ -737,6 +967,274 @@ public ITensor Forward(ReadOnlySpan tokenIds, ReadOnlySpan positions, return result; } + // ────────────────────────────────────────────────────────────────────── + // MTP (issue #253) — self-speculative decoding draft head + // ────────────────────────────────────────────────────────────────────── + + /// + /// Lazily-allocated, model-owned device scratch for — one row + /// (seqLen=1) worth of every intermediate buffer the MTP block's own decoder-layer forward + /// needs. Sized once from (fixed per model instance) and reused across + /// every call, mirroring the model's other lazily-allocated instance + /// scratch (e.g. _activF16InScratch) rather than allocating/freeing per call — MTP draft + /// steps run K times per speculation round, frequently enough that repeated cuMemAlloc/cuMemFree + /// overhead would be wasteful, even though the call is off the trunk's hot decode path. + /// + private sealed class CudaMtpScratch : IDisposable + { + public nint Embed; // [hiddenSize] + public nint Concat; // [2*hiddenSize] — eNorm at [0,hiddenSize), hNorm at [hiddenSize,2*hiddenSize) + public nint Cur; // [hiddenSize] + public nint Residual; // [hiddenSize] + public nint Normed; // [hiddenSize] + public nint Qg; // [2*qElems] + public nint Q; // [qElems] + public nint Gate; // [qElems] + public nint AttnOut; // [qElems] + public nint FfnGate; // [intermediateSize] + public nint FfnUp; // [intermediateSize] + public nint Silu; // [intermediateSize] + public nint NormedHead; // [hiddenSize] + public nint LogitsDevice; // [vocabSize] + public nint PositionDevice; // [1] int32 + + public static CudaMtpScratch Allocate(int hiddenSize, int qElems, int intermediateSize, int vocabSize) + { + var s = new CudaMtpScratch + { + Embed = AllocDevice((long)hiddenSize * sizeof(float)), + Concat = AllocDevice(2L * hiddenSize * sizeof(float)), + Cur = AllocDevice((long)hiddenSize * sizeof(float)), + Residual = AllocDevice((long)hiddenSize * sizeof(float)), + Normed = AllocDevice((long)hiddenSize * sizeof(float)), + Qg = AllocDevice(2L * qElems * sizeof(float)), + Q = AllocDevice((long)qElems * sizeof(float)), + Gate = AllocDevice((long)qElems * sizeof(float)), + AttnOut = AllocDevice((long)qElems * sizeof(float)), + FfnGate = AllocDevice((long)intermediateSize * sizeof(float)), + FfnUp = AllocDevice((long)intermediateSize * sizeof(float)), + Silu = AllocDevice((long)intermediateSize * sizeof(float)), + NormedHead = AllocDevice((long)hiddenSize * sizeof(float)), + LogitsDevice = AllocDevice((long)vocabSize * sizeof(float)), + PositionDevice = AllocDevice(sizeof(int)), + }; + return s; + } + + public void Dispose() + { + FreeIfNonZero(ref Embed); + FreeIfNonZero(ref Concat); + FreeIfNonZero(ref Cur); + FreeIfNonZero(ref Residual); + FreeIfNonZero(ref Normed); + FreeIfNonZero(ref Qg); + FreeIfNonZero(ref Q); + FreeIfNonZero(ref Gate); + FreeIfNonZero(ref AttnOut); + FreeIfNonZero(ref FfnGate); + FreeIfNonZero(ref FfnUp); + FreeIfNonZero(ref Silu); + FreeIfNonZero(ref NormedHead); + FreeIfNonZero(ref LogitsDevice); + FreeIfNonZero(ref PositionDevice); + } + } + + private CudaMtpScratch? _mtpScratch; + + /// + public ITensor ForwardMtp(IMtpState state, int tokenId, int position) + { + ObjectDisposedException.ThrowIf(_disposed, this); + if (_mtpHead is not { } mtpHead) + throw new NotSupportedException( + $"{nameof(CudaQwen3HybridDenseTransformerModel)} has no MTP head loaded (SupportsMtp=false)."); + if (state is not CudaMtpState mtp) + throw new ArgumentException( + $"CudaQwen3HybridDenseTransformerModel requires a CUDA CudaMtpState; got {state.GetType().Name}.", + nameof(state)); + + _context.MakeCurrent(); + return ForwardMtpCore(mtpHead, mtp, tokenId, position); + } + + /// + /// Runs one MTP head autoregressive draft step (issue #253) — see . + /// Off the trunk's hot forward path (single token, called K≤~16 times per speculation round). + /// Operation order mirrors the CPU host's Qwen3HybridDenseTransformerModel.ForwardMtpCore + /// exactly (confirmed against llama.cpp PR ggml-org/llama.cpp#22673's graph_mtp): + /// + /// h_norm = RMSNorm(pendingHidden, nextn.hnorm); e_norm = RMSNorm(embed(tokenId), nextn.enorm). + /// cur = eh_proj @ concat(e_norm, h_norm) — this becomes the attention sub-block's residual (inpSA). + /// Gated full attention over the MTP head's own device-resident KV-cache (identical math to + /// , but seqQ=1 against the head's private cache rather + /// than the trunk's), residual-added back onto inpSA. + /// Dense SwiGLU FFN, residual-added — the result is the MTP block's own output hidden + /// state ("h_pre_norm"), which seeds this state's next call + /// (D2D-copied directly into — no host round-trip). + /// shared_head_norm (or the trunk's output_norm fallback) then + /// shared_head_head (or the trunk's own LM head fallback) → logits. + /// + /// + private ITensor ForwardMtpCore(in CudaMtpHeadWeights mtpHead, CudaMtpState state, int tokenId, int position) + { + int hiddenSize = Config.HiddenSize; + int vocabSize = Config.VocabSize; + var attn = mtpHead.Layer.FullAttn!.Value; + int numHeads = Config.NumAttentionHeads; + int numKvHeads = attn.NumKvHeads; + int headDim = Config.HeadDim; + int qElems = numHeads * headDim; + int intermediateSize = mtpHead.Layer.GateOutputDim; + float eps = Config.NormEpsilon; + nint streamH = _stream.Handle; + + int step = state.CurrentLength; + if (step >= state.MaxSteps) + throw new InvalidOperationException( + $"CudaMtpState KV-cache exhausted ({state.MaxSteps} steps advanced). Size the state for " + + "at least numCandidates MTP draft steps per speculation round."); + + _mtpScratch ??= CudaMtpScratch.Allocate(hiddenSize, qElems, intermediateSize, vocabSize); + var s = _mtpScratch; + + // ── Embed predicted-from token (host dequant of one row + tiny H2D — same pattern as the + // trunk's own per-token embedding lookup in ForwardCore) ── + nint embedHostBase = mtpHead.EmbedTokensHostBase ?? _embedDataBase; + ulong embedDataOffset = mtpHead.EmbedTokensHostBase is not null ? mtpHead.EmbedTokensDataOffset : _embedDataOffset; + long embedRowBytes = mtpHead.EmbedTokensHostBase is not null ? mtpHead.EmbedTokensRowBytes : _embedRowBytes; + QuantizationType embedQt = mtpHead.EmbedTokensHostBase is not null ? mtpHead.EmbedTokensQt : _tokenEmbedQt; + + float[] embedHost = new float[hiddenSize]; + nint rowSrc = embedHostBase + (nint)(embedDataOffset + (ulong)tokenId * (ulong)embedRowBytes); + Dequantize.ToFloat32(rowSrc, hiddenSize, embedQt, embedHost); + fixed (float* pEmbedHost = embedHost) + { + CudaDriverApi.cuMemcpyHtoDAsync_v2(s.Embed, (nint)pEmbedHost, + (nuint)((long)hiddenSize * sizeof(float)), streamH).ThrowOnError(); + } + + // ── h_norm / e_norm — written directly into the two halves of the eh_proj concat buffer + // (avoids an extra D2D copy vs. normalizing into standalone buffers first) ── + nint eNormDst = s.Concat; + nint hNormDst = s.Concat + (nint)((long)hiddenSize * sizeof(float)); + _kernels.LaunchRmsNormF32(s.Embed, mtpHead.EnormDevice, eNormDst, hiddenSize, eps, 1, streamH); + _kernels.LaunchRmsNormF32(state.PendingHiddenDevicePtr, mtpHead.HnormDevice, hNormDst, hiddenSize, eps, 1, streamH); + + // cur = eh_proj @ concat(e_norm, h_norm) + Gemm(mtpHead.EhProjDevice, mtpHead.EhProjQt, s.Concat, s.Cur, + mtpHead.EhProjOutputDim, mtpHead.EhProjInputDim, 1); + + // inpSA: the attention sub-block's residual is the eh_proj output, not the raw input. + CudaDriverApi.cuMemcpyDtoDAsync_v2(s.Residual, s.Cur, (nuint)((long)hiddenSize * sizeof(float)), streamH) + .ThrowOnError(); + + // ── Attention sub-block — same gated-QKV math as ForwardFullAttnBody, seqQ=1 ── + _kernels.LaunchRmsNormF32(s.Cur, mtpHead.Layer.AttnNormWeightDevice, s.Normed, hiddenSize, eps, 1, streamH); + + Gemm(attn.QDevice, attn.QQt, s.Normed, s.Qg, attn.QOutputDim, attn.QInputDim, 1); + + if (_kernels.HasDeinterleaveF32) + { + _kernels.LaunchDeinterleaveQGateF32(s.Qg, s.Q, s.Gate, numHeads, headDim, 1, streamH); + } + else + { + long perHeadBytes = (long)headDim * sizeof(float); + for (int h = 0; h < numHeads; h++) + { + nint qgHead = s.Qg + (nint)(h * 2 * perHeadBytes); + nint qHead = s.Q + (nint)(h * perHeadBytes); + nint gHead = s.Gate + (nint)(h * perHeadBytes); + CudaDriverApi.cuMemcpyDtoDAsync_v2(qHead, qgHead, (nuint)perHeadBytes, streamH).ThrowOnError(); + CudaDriverApi.cuMemcpyDtoDAsync_v2(gHead, qgHead + (nint)perHeadBytes, + (nuint)perHeadBytes, streamH).ThrowOnError(); + } + } + + // K/V projections write directly into this step's KV-cache row — appends this step's K/V + // into the MTP head's own tiny cache (no extra copy), matching the CPU host's + // `k.CopyTo(state.GetKeyRow(step))` / `v.CopyTo(state.GetValueRow(step))`. + nint kRowDst = state.GetKeyRowDevicePtr(step); + nint vRowDst = state.GetValueRowDevicePtr(step); + Gemm(attn.KDevice, attn.KQt, s.Normed, kRowDst, attn.KOutputDim, attn.KInputDim, 1); + Gemm(attn.VDevice, attn.VQt, s.Normed, vRowDst, attn.VOutputDim, attn.VInputDim, 1); + + // Per-head QK-norm (RMSNorm over headDim, one "row" per head — seqLen=1 * numHeads/numKvHeads rows). + _kernels.LaunchRmsNormF32(s.Q, attn.QNormDevice, s.Q, headDim, eps, numHeads, streamH); + _kernels.LaunchRmsNormF32(kRowDst, attn.KNormDevice, kRowDst, headDim, eps, numKvHeads, streamH); + + // RoPE — partial-rotary NeoX, at this step's absolute round-relative position. + int[] posHost = [position]; + fixed (int* pPos = posHost) + { + CudaDriverApi.cuMemcpyHtoDAsync_v2(s.PositionDevice, (nint)pPos, sizeof(int), streamH).ThrowOnError(); + } + _kernels.LaunchRoPEF32(s.Q, kRowDst, s.PositionDevice, 1, numHeads, numKvHeads, headDim, + _ropeDim, _ropeTheta, 1, streamH); + + // Append this step's K/V (already written above) and attend causally over everything + // drafted so far in this round (NOT the trunk's KV-cache) — positionOffset=step means every + // key j in [0, step] satisfies the causal check (step >= j), i.e. no masking within the + // round, matching the CPU host's Attention.Execute(..., positionOffset: step, ...) call. + int seqKv = step + 1; + _kernels.LaunchAttentionF32(s.Q, state.KeyCacheDevicePtr, state.ValueCacheDevicePtr, s.AttnOut, + /* seqQ */ 1, /* seqKv */ seqKv, numHeads, numKvHeads, headDim, + /* positionOffset */ step, /* slidingWindow */ 0, streamH); + + // attnOut *= sigmoid(gate) — Qwen3.5/3.6 gated attention, applied before the O-proj. + if (_kernels.HasElementwiseF32) + _kernels.LaunchSigmoidMulF32(s.AttnOut, s.Gate, qElems, streamH); + else + LaunchSigmoidMulHostFallback(s.AttnOut, s.Gate, qElems); + + Gemm(attn.ODevice, attn.OQt, s.AttnOut, s.Cur, attn.OOutputDim, attn.OInputDim, 1); + + _kernels.LaunchAddF32(s.Residual, s.Cur, s.Cur, hiddenSize, streamH); // cur = inpSA + attn_out_projected + + // ── Dense SwiGLU FFN sub-layer ── + CudaDriverApi.cuMemcpyDtoDAsync_v2(s.Residual, s.Cur, (nuint)((long)hiddenSize * sizeof(float)), streamH) + .ThrowOnError(); // ffn_residual + _kernels.LaunchRmsNormF32(s.Cur, mtpHead.Layer.PostAttnNormWeightDevice, s.Normed, hiddenSize, eps, 1, streamH); + + Gemm(mtpHead.Layer.GateWeight, mtpHead.Layer.GateQt, s.Normed, s.FfnGate, + mtpHead.Layer.GateOutputDim, mtpHead.Layer.GateInputDim, 1); + Gemm(mtpHead.Layer.UpWeight, mtpHead.Layer.UpQt, s.Normed, s.FfnUp, + mtpHead.Layer.UpOutputDim, mtpHead.Layer.UpInputDim, 1); + _kernels.LaunchSwiGLUF32(s.FfnGate, s.FfnUp, s.Silu, intermediateSize, 1, streamH); + Gemm(mtpHead.Layer.DownWeight, mtpHead.Layer.DownQt, s.Silu, s.Cur, + mtpHead.Layer.DownOutputDim, mtpHead.Layer.DownInputDim, 1); + + _kernels.LaunchAddF32(s.Residual, s.Cur, s.Cur, hiddenSize, streamH); // cur = ffn_residual + ffn_out + + // `cur` is now the MTP block's own output hidden state ("h_pre_norm" in llama.cpp) — seed + // the NEXT ForwardMtp call's pending hidden with it (D2D, stays fully device-resident) + // before the head-norm below consumes it, then advance the MTP KV-cache length. + CudaDriverApi.cuMemcpyDtoDAsync_v2(state.PendingHiddenDevicePtr, s.Cur, + (nuint)((long)hiddenSize * sizeof(float)), streamH).ThrowOnError(); + state.Advance(); + + // ── Shared LM head (falls back to the trunk's output_norm/output.weight when the GGUF + // didn't ship head-local nextn.shared_head_* tensors) ── + nint headNormWeight = mtpHead.SharedHeadNormDevice ?? _outputNormDevice; + _kernels.LaunchRmsNormF32(s.Cur, headNormWeight, s.NormedHead, hiddenSize, eps, 1, streamH); + + nint headWeight = mtpHead.SharedHeadHeadDevice ?? _outputDevice; + QuantizationType headQt = mtpHead.SharedHeadHeadDevice is not null ? mtpHead.SharedHeadHeadQt : _outputQt; + int headOutputDim = mtpHead.SharedHeadHeadDevice is not null ? mtpHead.SharedHeadHeadOutputDim : _outputOutputDim; + int headInputDim = mtpHead.SharedHeadHeadDevice is not null ? mtpHead.SharedHeadHeadInputDim : _outputInputDim; + + Gemm(headWeight, headQt, s.NormedHead, s.LogitsDevice, headOutputDim, headInputDim, 1); + + _stream.Synchronize(); + var shape = new TensorShape(1, vocabSize); + var result = UnmanagedTensor.Allocate(shape, DType.Float32, deviceId: -1); + CudaDriverApi.cuMemcpyDtoH_v2(result.DataPointer, s.LogitsDevice, + (nuint)((long)vocabSize * sizeof(float))).ThrowOnError(); + return result; + } + // Issue #178 (2026-07-25 profiling round): re-verified that HasCopyRmsNormF32 is actually // true at runtime for both "layer-pre-norm1" and "layer-resid1-norm2" below (a real xUnit // pass on CudaCopyRmsNormF32Test, not just the property returning true — confirms PTX isn't @@ -1999,6 +2497,20 @@ public void Dispose() FreeIfNonZero(ref _attnGqaSplitPartialOut); FreeIfNonZero(ref _attnMmaDecodeQF16); + // MTP (issue #253): free the trailing NextN head's device weights (a no-op when absent). + if (_mtpHead is { } mtpHead) + { + var layerCopy = mtpHead.Layer; + FreeLayer(ref layerCopy); + nint ehProj = mtpHead.EhProjDevice; if (ehProj != 0) CudaDriverApi.cuMemFree_v2(ehProj); + nint enorm = mtpHead.EnormDevice; if (enorm != 0) CudaDriverApi.cuMemFree_v2(enorm); + nint hnorm = mtpHead.HnormDevice; if (hnorm != 0) CudaDriverApi.cuMemFree_v2(hnorm); + if (mtpHead.SharedHeadHeadDevice is { } shh && shh != 0) CudaDriverApi.cuMemFree_v2(shh); + if (mtpHead.SharedHeadNormDevice is { } shn && shn != 0) CudaDriverApi.cuMemFree_v2(shn); + } + _mtpScratch?.Dispose(); + _mtpScratch = null; + _state.Dispose(); _gdnCache.Dispose(); _kernels.Dispose(); @@ -2259,4 +2771,53 @@ internal struct DeviceFullAttn public nint QNormDevice; public nint KNormDevice; } + + /// + /// Device-side weights for a Multi-Token Prediction (MTP / "NextN") head — the trailing extra + /// decoder block used for self-speculative decoding (issue #253). Mirrors + /// DotLLM.Models.Architectures.MtpHeadWeights (CPU): the MTP block for Qwen3.5/3.6 is + /// structurally a full-attention Qwen3HybridDense decoder layer () with four + /// extra "nextn" tensors wrapped around it. is a host-mmap + /// pointer (not device-resident) — MTP embeds one token per ForwardMtp call via a + /// host-side dequant + tiny H2D copy, exactly like the trunk's own per-token embedding lookup. + /// / are device-resident + /// (they feed a GEMM / RMSNorm kernel respectively) and null when the GGUF didn't ship + /// head-local nextn.shared_head_* tensors — the trunk's own lm_head/output_norm are used + /// instead in that case (see ). + /// + internal struct CudaMtpHeadWeights + { + /// + /// The MTP block's own decoder-layer weights (attn norms, gated full attention, dense FFN) — + /// structurally identical to any other full-attention . + /// + public DeviceLayer Layer; + + /// nextn.eh_proj.weight [2·hiddenSize, hiddenSize]. Quantized, device-resident. + public nint EhProjDevice; + public QuantizationType EhProjQt; + public int EhProjInputDim; + public int EhProjOutputDim; + + /// nextn.enorm.weight [hiddenSize] — F32 device RMSNorm weight applied to the predicted token's embedding. + public nint EnormDevice; + + /// nextn.hnorm.weight [hiddenSize] — F32 device RMSNorm weight applied to the incoming trunk hidden state. + public nint HnormDevice; + + /// Optional nextn.embed_tokens.weight host-mmap base pointer. Null ⇒ reuse the trunk's token_embd.weight (_embedDataBase). + public nint? EmbedTokensHostBase; + public ulong EmbedTokensDataOffset; + public long EmbedTokensRowBytes; + public QuantizationType EmbedTokensQt; + + /// Optional nextn.shared_head_head.weight [hiddenSize, vocabSize], device-resident. Null ⇒ reuse the trunk's lm_head. + public nint? SharedHeadHeadDevice; + public QuantizationType SharedHeadHeadQt; + public int SharedHeadHeadInputDim; + public int SharedHeadHeadOutputDim; + + /// Optional nextn.shared_head_norm.weight [hiddenSize], F32 device-resident. Null ⇒ reuse the trunk's output_norm.weight. + public nint? SharedHeadNormDevice; + } } diff --git a/src/DotLLM.Engine/MtpSpeculativeDecoder.cs b/src/DotLLM.Engine/MtpSpeculativeDecoder.cs index 59a2c746..ad933309 100644 --- a/src/DotLLM.Engine/MtpSpeculativeDecoder.cs +++ b/src/DotLLM.Engine/MtpSpeculativeDecoder.cs @@ -38,7 +38,8 @@ namespace DotLLM.Engine; /// against plain greedy decode of the same (synthetic) model. /// /// -/// The "catchup" forward — why every round starts by re-processing lastToken. +/// The "catchup" forward — why every round starts by re-processing lastToken, and why +/// the verify batch does NOT re-submit it (issue #253 CUDA follow-up, fixed 2026-08-07). /// The MTP head's first draft step needs mtpState seeded with the trunk's own /// hidden state after processing lastToken (the pairing invariant, confirmed /// against llama.cpp's graph_mtp: h_input and tok_embd in the same MTP call @@ -49,17 +50,26 @@ namespace DotLLM.Engine; /// construction, never been forwarded through the trunk as an input: there is no row in /// the previous round's verify batch whose hidden state reflects it. So every round begins with a /// single-token trunk forward of lastToken (with mtpState capture) purely -/// to obtain that hidden state before the MTP draft loop can start. This re-forward is safe and -/// idempotent — lastToken already occupies position in -/// kvCacheTarget from the prior round (or the initial prefill); -/// is position-indexed (Rollback's doc: "Allocated memory is retained and overwritten on -/// subsequent Update calls"), so writing the same token at the same position again is a no-op on -/// the cache contents. The verify batch afterward still includes lastToken as its own row 0 -/// (matching 's shape) for the comparison-basis logits it needs — -/// a second harmless re-forward of the same token. This doubles the trunk's per-round -/// single-token-equivalent cost relative to a maximally-optimized implementation that reuses the -/// catchup call's own logits as the row-0 comparison basis; documented here as a known, -/// correctness-first simplification rather than silently traded away. +/// to obtain that hidden state before the MTP draft loop can start. +/// +/// +/// The catchup call's own logits are ALSO this round's "position 0" comparison basis for +/// draftTokens[0] — reused directly (catchupArgmax), rather than re-submitting +/// lastToken as row 0 of the verify batch the way 's +/// two-model verify phase does. An earlier version of this method DID re-submit it, reasoning +/// (correctly, but incompletely) that is position-indexed so re-writing the +/// same token at the same position is a no-op on cache contents. That reasoning does not extend +/// to recurrent trunk layers (Gated DeltaNet / Mamba, exactly the token-mixing kind +/// hybrid architectures use — the real MTP target, +/// Qwen3.6-27B/Bonsai-27B, IS one): their state is a pure sequential recurrence, not +/// position-indexed, so forwarding the same token through it a second time double-advances that +/// state and corrupts every subsequent decode step. This was caught empirically by a CUDA +/// integration test driving the real (GDN-containing) Qwen3HybridDense model through this +/// decoder and comparing against plain greedy decode of the same model — a comparison the original +/// mock-model unit tests could not catch because their mock used a non-recurrent architecture. +/// Skipping the redundant row also removes the "doubles the trunk's per-round single-token +/// cost" overhead the original version explicitly traded away as a documented simplification — +/// fixing the bug turned out to be strictly cheaper, not a tradeoff. /// /// public sealed class MtpSpeculativeDecoder : IMtpSpeculativeDecoder @@ -116,13 +126,23 @@ public SpeculativeResult DraftAndVerify( // Fresh MTP head KV-cache for this round — see the type remarks on why this is safe. mtpState.Rollback(0); - // ── Catchup (see type remarks): seed mtpState with h-after-lastToken before drafting. ── + // ── Catchup (see type remarks): seed mtpState with h-after-lastToken before drafting. + // ALSO capture this call's own logits' argmax — this IS position 0's verify comparison + // basis (see the "no redundant re-forward" remarks below), so the verify batch never + // resubmits lastToken. ── long catchupStart = Stopwatch.GetTimestamp(); + int catchupArgmax; using (ITensor catchupLogits = targetModel.Forward( [lastToken], [position], deviceId: -1, kvCacheTarget, adapter: null, mtpState)) { mtpState.SeedFromCapturedRow(0); - _ = catchupLogits; // logits themselves unused here — only the hidden-state side effect matters + unsafe + { + var span = new Span((void*)catchupLogits.DataPointer, vocabSize); + if (constraint != null) + TokenMaskApplier.Apply(span, constraint.GetAllowedTokens()); + catchupArgmax = TensorPrimitives.IndexOfMax((ReadOnlySpan)span); + } } verifyTicks += Stopwatch.GetTimestamp() - catchupStart; @@ -166,42 +186,67 @@ public SpeculativeResult DraftAndVerify( generatedIds.RemoveRange(originalGenCount, generatedIds.Count - originalGenCount); } - // ── Verify Phase (single batched forward pass over the target model — identical - // shape to SpeculativeDecoder's verify phase; row 0 redundantly re-forwards - // lastToken, matching the catchup step above — see type remarks). ── - int verifyLen = k + 1; + // ── Accept/reject position 0 against the catchup call's own argmax — NOT a fresh + // verify-batch row. Re-submitting lastToken as a verify-batch row (the original + // design) would be byte-redundant for attention/KV-cache math (causally + // independent of later batch rows) but subtly WRONG for recurrent (GDN/Mamba) + // trunk layers: their state is a pure sequential recurrence, not position-indexed, + // so forwarding the same token through it twice per round double-advances that + // state and corrupts every subsequent decode step. catchupArgmax IS the same + // computation a fresh "row 0" would have produced (identical input token, position, + // and preceding KV-cache/recurrent state) — reusing it is both correct and, as a + // side effect, removes the redundant single-token forward the original design paid + // every round. ── + int acceptedCount = 0; + if (draftTokens[0] == catchupArgmax) + { + outputBuffer[acceptedCount++] = draftTokens[0]; + constraint?.Advance(draftTokens[0]); + } + else + { + outputBuffer[acceptedCount++] = catchupArgmax; + constraint?.Advance(catchupArgmax); + RollbackKvCache(kvCacheTarget, position, acceptedCount); + return new SpeculativeResult(acceptedCount, draftTicks, verifyTicks, k); + } + + // ── Verify Phase (single batched forward pass over ALL of draftTokens[0..k-1] at + // position+1..position+k — this is the pre-fix verify batch with ONLY the leading + // lastToken row dropped; draftTokens[0] itself still needs to appear as an INPUT + // token here even though its own value was already resolved via catchupArgmax + // above, because it is what row m needs to predict draftTokens[m+1]. Row m (0-based) + // predicts the token after position+m+1, i.e. it is the comparison basis for + // draftTokens[m+1] (m=0..k-2); the LAST row (m=k-1, input=draftTokens[k-1]) doubles + // as the bonus-token source when every draft token is accepted, mirroring + // SpeculativeDecoder's verify shape. ── + int verifyLen = k; Span verifyTokens = verifyLen <= 16 ? stackalloc int[verifyLen] : new int[verifyLen]; Span verifyPositions = verifyLen <= 16 ? stackalloc int[verifyLen] : new int[verifyLen]; - - verifyTokens[0] = lastToken; - verifyPositions[0] = position; - for (int i = 0; i < k; i++) + for (int i = 0; i < verifyLen; i++) { - verifyTokens[i + 1] = draftTokens[i]; - verifyPositions[i + 1] = position + i + 1; + verifyTokens[i] = draftTokens[i]; + verifyPositions[i] = position + i + 1; } - int actualVerifyLen = Math.Min(verifyLen, maxPos - position); - if (actualVerifyLen < 1) - return default; - + // k is already clamped to maxPos - position - 1 above, so position + k <= maxPos - 1 + // and every verify position here (<= position + k) is guaranteed in-range — no + // additional clamp needed (unlike the pre-fix code, which clamped defensively against + // an off-by-one that can no longer occur with this narrower verify range). long verifyStart = Stopwatch.GetTimestamp(); using ITensor targetLogits = targetModel.Forward( - verifyTokens.Slice(0, actualVerifyLen), - verifyPositions.Slice(0, actualVerifyLen), - deviceId: -1, kvCacheTarget, adapter: null); + verifyTokens, verifyPositions, deviceId: -1, kvCacheTarget, adapter: null); verifyTicks += Stopwatch.GetTimestamp() - verifyStart; - // ── Accept/Reject Phase (greedy argmax — see type remarks for the correctness argument) ── - int acceptedCount = 0; - unsafe { nint basePtr = targetLogits.DataPointer; - for (int i = 0; i < Math.Min(k, actualVerifyLen); i++) + // Rows 0..k-2 verify draftTokens[1..k-1]; row k-1 is reserved for the bonus token + // below and is never itself an accept/reject comparison target. + for (int i = 0; i < verifyLen - 1; i++) { - int draftTok = draftTokens[i]; + int draftTok = draftTokens[i + 1]; var targetLogitSpan = new Span( (void*)(basePtr + (long)i * vocabSize * sizeof(float)), vocabSize); @@ -223,18 +268,16 @@ public SpeculativeResult DraftAndVerify( } } - // All K accepted — sample bonus token from the target's own argmax. - if (actualVerifyLen > k) - { - var bonusLogitSpan = new Span( - (void*)(basePtr + (long)k * vocabSize * sizeof(float)), vocabSize); + // All K accepted — sample bonus token from the LAST verify row's own argmax + // (predicts position+k+1, exactly matching the pre-fix design's bonus semantics). + var bonusLogitSpan = new Span( + (void*)(basePtr + (long)(verifyLen - 1) * vocabSize * sizeof(float)), vocabSize); - if (constraint != null) - TokenMaskApplier.Apply(bonusLogitSpan, constraint.GetAllowedTokens()); + if (constraint != null) + TokenMaskApplier.Apply(bonusLogitSpan, constraint.GetAllowedTokens()); - int bonusToken = TensorPrimitives.IndexOfMax((ReadOnlySpan)bonusLogitSpan); - outputBuffer[acceptedCount++] = bonusToken; - } + int bonusToken = TensorPrimitives.IndexOfMax((ReadOnlySpan)bonusLogitSpan); + outputBuffer[acceptedCount++] = bonusToken; } RollbackKvCache(kvCacheTarget, position, acceptedCount); diff --git a/src/DotLLM.Models/Gguf/SyntheticQwen35HybridDenseMtpGguf.cs b/src/DotLLM.Models/Gguf/SyntheticQwen35HybridDenseMtpGguf.cs index 39edac66..0f1815b2 100644 --- a/src/DotLLM.Models/Gguf/SyntheticQwen35HybridDenseMtpGguf.cs +++ b/src/DotLLM.Models/Gguf/SyntheticQwen35HybridDenseMtpGguf.cs @@ -64,7 +64,14 @@ public static class SyntheticQwen35HybridDenseMtpGguf /// "fall back to the trunk's token_embd/output/output_norm" path. Ignored when /// is . /// - public static byte[] Build(uint seed = 0xC0FFEEu, bool withMtp = true, bool mtpHasOwnHeadTensors = true) + /// + /// Overrides 's default mixed GDN+attention layout. Pass 1 for an + /// all-full-attention trunk (no GDN layer at all) — useful for isolating a test from the + /// separate, pre-existing "speculative decoding has no rollback for recurrent trunk state" + /// limitation (see MtpSpeculativeDecoder's remarks). + /// + public static byte[] Build(uint seed = 0xC0FFEEu, bool withMtp = true, bool mtpHasOwnHeadTensors = true, + int fullAttnInterval = FullAttnInterval) { var w = new GgufWriter(); var rng = new SyntheticGemma4Gguf.Xorshift(seed); @@ -89,7 +96,11 @@ public static byte[] Build(uint seed = 0xC0FFEEu, bool withMtp = true, bool mtpH w.AddUInt32($"{arch}.rope.dimension_count", RopeDim); // Hybrid layout: trunk layer i is full attention when (i+1) % full_attention_interval == 0. - w.AddUInt32($"{arch}.full_attention_interval", FullAttnInterval); + // fullAttnInterval defaults to the standard fixture shape (mixed GDN+attention); callers that + // need an all-full-attention trunk (e.g. to isolate a test from the separate, pre-existing + // "speculative decoding + recurrent trunk state has no rollback" limitation — see + // MtpSpeculativeDecoder's remarks / issue #253's CUDA follow-up notes) pass 1. + w.AddUInt32($"{arch}.full_attention_interval", (uint)fullAttnInterval); if (withMtp) w.AddUInt32($"{arch}.nextn_predict_layers", 1); @@ -112,7 +123,7 @@ public static byte[] Build(uint seed = 0xC0FFEEu, bool withMtp = true, bool mtpH for (int i = 0; i < BlockCount; i++) { - bool fullAttn = (i + 1) % FullAttnInterval == 0; + bool fullAttn = (i + 1) % fullAttnInterval == 0; string p = $"blk.{i}"; AddNorm(w, rng, $"{p}.attn_norm.weight", HiddenSize); @@ -150,9 +161,10 @@ public static byte[] Build(uint seed = 0xC0FFEEu, bool withMtp = true, bool mtpH } /// Writes the synthetic fixture to . - public static string Write(string path, uint seed = 0xC0FFEEu, bool withMtp = true, bool mtpHasOwnHeadTensors = true) + public static string Write(string path, uint seed = 0xC0FFEEu, bool withMtp = true, bool mtpHasOwnHeadTensors = true, + int fullAttnInterval = FullAttnInterval) { - File.WriteAllBytes(path, Build(seed, withMtp, mtpHasOwnHeadTensors)); + File.WriteAllBytes(path, Build(seed, withMtp, mtpHasOwnHeadTensors, fullAttnInterval)); return path; } diff --git a/tests/DotLLM.Tests.Unit/Cuda/CudaQwen3HybridDenseMtpTests.cs b/tests/DotLLM.Tests.Unit/Cuda/CudaQwen3HybridDenseMtpTests.cs new file mode 100644 index 00000000..9397e7fc --- /dev/null +++ b/tests/DotLLM.Tests.Unit/Cuda/CudaQwen3HybridDenseMtpTests.cs @@ -0,0 +1,406 @@ +using System.Runtime.InteropServices; +using DotLLM.Core.Configuration; +using DotLLM.Core.Models; +using DotLLM.Core.Tensors; +using DotLLM.Cuda; +using DotLLM.Cuda.Architectures; +using DotLLM.Engine; +using DotLLM.Engine.Samplers; +using DotLLM.Models.Gguf; +using Xunit; +using Xunit.Abstractions; +using Architecture = DotLLM.Core.Configuration.Architecture; + +namespace DotLLM.Tests.Unit.Cuda; + +/// +/// CUDA coverage for issue #253 (Multi-Token Prediction self-speculative decoding), mirroring +/// Qwen3HybridDenseMtpTests (CPU): GGUF detection/loading of the trailing MTP/"NextN" head +/// on , the MTP head's own CUDA forward pass +/// ( device-resident KV-cache + pending-hidden handoff), and an +/// engine-level integration test proving — which +/// is completely unmodified/backend-agnostic — produces the same output sequence as plain greedy +/// decode when driven by the real CUDA model. Uses +/// — the same tiny (F32, ~KB-scale) fixture the CPU +/// tests use, so no real Qwen3.6-MTP-GGUF download is needed. +/// +[Trait("Category", "GPU")] +[Collection(CudaCollection.Name)] +public sealed class CudaQwen3HybridDenseMtpTests : IDisposable +{ + private readonly string _scratch; + private readonly ITestOutputHelper _out; + + public CudaQwen3HybridDenseMtpTests(ITestOutputHelper output) + { + _out = output; + _scratch = Path.Combine(Path.GetTempPath(), $"dotllm-cuda-qwen35-mtp-{Guid.NewGuid():N}"); + Directory.CreateDirectory(_scratch); + } + + public void Dispose() + { + try { Directory.Delete(_scratch, recursive: true); } catch { /* best-effort */ } + } + + private string WriteFixture(bool withMtp, bool mtpHasOwnHeadTensors = true, string name = "qwen35-mtp.gguf") => + SyntheticQwen35HybridDenseMtpGguf.Write( + Path.Combine(_scratch, name), withMtp: withMtp, mtpHasOwnHeadTensors: mtpHasOwnHeadTensors); + + private static bool IsCudaDriverPresent() + { + string lib = OperatingSystem.IsWindows() ? "nvcuda.dll" : "libcuda.so.1"; + if (!NativeLibrary.TryLoad(lib, out nint h)) return false; + NativeLibrary.Free(h); + return CudaDevice.IsAvailable(); + } + + private static string? FindPtxDir() + { + var candidates = new[] + { + Path.Combine(AppContext.BaseDirectory, "ptx"), + Path.Combine(AppContext.BaseDirectory, "..", "..", "..", "..", "..", "native", "ptx"), + }; + foreach (var dir in candidates) + { + var full = Path.GetFullPath(dir); + if (Directory.Exists(full) && Directory.GetFiles(full, "*.ptx").Length > 0) + return full; + } + return null; + } + + // ── GGUF detection / zero-behavior-change ────────────────────────────────── + + [SkippableFact] + public void LoadFromGguf_WithMtp_DetectsHeadAndCreatesState() + { + Skip.IfNot(IsCudaDriverPresent(), "No CUDA GPU available"); + string? ptxDir = FindPtxDir(); + Skip.If(ptxDir is null, "PTX files not found"); + + string path = WriteFixture(withMtp: true); + using var gguf = GgufFile.Open(path); + var config = GgufModelConfigExtractor.Extract(gguf.Metadata); + + Assert.Equal(Architecture.Qwen3HybridDense, config.Architecture); + Assert.Equal(1, config.NextnPredictLayers); + Assert.Equal(SyntheticQwen35HybridDenseMtpGguf.BlockCount, config.NumLayers); + + using var model = CudaQwen3HybridDenseTransformerModel.LoadFromGguf(gguf, config, deviceId: 0, ptxDir); + Assert.True(model.SupportsMtp); + using IMtpState? state = model.CreateMtpState(); + Assert.NotNull(state); + Assert.IsType(state); + } + + [SkippableFact] + public void LoadFromGguf_WithoutMtp_ZeroBehaviorChange() + { + Skip.IfNot(IsCudaDriverPresent(), "No CUDA GPU available"); + string? ptxDir = FindPtxDir(); + Skip.If(ptxDir is null, "PTX files not found"); + + string path = WriteFixture(withMtp: false); + using var gguf = GgufFile.Open(path); + var config = GgufModelConfigExtractor.Extract(gguf.Metadata); + + Assert.Equal(0, config.NextnPredictLayers); + Assert.Equal(SyntheticQwen35HybridDenseMtpGguf.BlockCount, config.NumLayers); + + using var model = CudaQwen3HybridDenseTransformerModel.LoadFromGguf(gguf, config, deviceId: 0, ptxDir); + Assert.False(model.SupportsMtp); + Assert.Null(model.CreateMtpState()); + + Assert.Throws(() => + { + using var state = new CudaMtpState(hiddenSize: 1, numKvHeads: 1, headDim: 1, maxSteps: 1); + model.ForwardMtp(state, tokenId: 0, position: 0); + }); + } + + [SkippableFact] + public void Forward_WithMtpStateCapture_ProducesIdenticalLogitsToPlainForward() + { + // The hidden-state capture is documented as a pure side effect — verify it byte-for-byte, + // mirroring the CPU host's equivalent test. + Skip.IfNot(IsCudaDriverPresent(), "No CUDA GPU available"); + string? ptxDir = FindPtxDir(); + Skip.If(ptxDir is null, "PTX files not found"); + + string path = WriteFixture(withMtp: true); + int[] tokenIds = [0, 1, 2, 3]; + int[] positions = [0, 1, 2, 3]; + + using var gguf1 = GgufFile.Open(path); + var config = GgufModelConfigExtractor.Extract(gguf1.Metadata); + using var model1 = CudaQwen3HybridDenseTransformerModel.LoadFromGguf(gguf1, config, deviceId: 0, ptxDir); + using var kvCache1 = model1.CreateKvCache(maxSeqLen: 64); + using ITensor plainLogits = model1.Forward(tokenIds, positions, deviceId: -1, kvCache1); + + // Two SEPARATE model instances (not model1.ResetSequenceState() + reuse): this model has a + // GatedDeltaNet layer whose recurrent state is model-owned, not IKvCache-owned — a fresh + // instance guarantees a genuinely independent starting state for the capture-vs-plain + // comparison, matching CudaQwen3HybridDenseLastTokenLogitsOnlyTest's established pattern + // for this exact class of model. + using var gguf2 = GgufFile.Open(path); + using var model2 = CudaQwen3HybridDenseTransformerModel.LoadFromGguf(gguf2, config, deviceId: 0, ptxDir); + using var kvCache2 = model2.CreateKvCache(maxSeqLen: 64); + using var mtpState = model2.CreateMtpState()!; + using ITensor capturedLogits = model2.Forward(tokenIds, positions, deviceId: -1, kvCache2, adapter: null, mtpState); + + Assert.Equal(plainLogits.Shape[0], capturedLogits.Shape[0]); + Assert.Equal(plainLogits.Shape[1], capturedLogits.Shape[1]); + unsafe + { + int n = tokenIds.Length * config.VocabSize; + var a = new ReadOnlySpan((void*)plainLogits.DataPointer, n); + var b = new ReadOnlySpan((void*)capturedLogits.DataPointer, n); + for (int i = 0; i < n; i++) + Assert.Equal(a[i], b[i]); // byte-identical float compare — pure side effect claim + } + + Assert.Equal(tokenIds.Length, mtpState.CapturedRowCount); + Assert.Equal(config.HiddenSize, mtpState.HiddenSize); + } + + // ── MTP head forward math ─────────────────────────────────────────────────── + + [SkippableFact] + public void ForwardMtp_ProducesFiniteLogitsAndAdvancesState() + { + Skip.IfNot(IsCudaDriverPresent(), "No CUDA GPU available"); + string? ptxDir = FindPtxDir(); + Skip.If(ptxDir is null, "PTX files not found"); + + string path = WriteFixture(withMtp: true); + using var gguf = GgufFile.Open(path); + var config = GgufModelConfigExtractor.Extract(gguf.Metadata); + using var model = CudaQwen3HybridDenseTransformerModel.LoadFromGguf(gguf, config, deviceId: 0, ptxDir); + + int[] tokenIds = [0, 1, 2]; + int[] positions = [0, 1, 2]; + using var kvCache = model.CreateKvCache(maxSeqLen: 64); + using var mtpState = (CudaMtpState)model.CreateMtpState()!; + using (ITensor _ = model.Forward(tokenIds, positions, deviceId: -1, kvCache, adapter: null, mtpState)) { } + mtpState.SeedFromCapturedRow(mtpState.CapturedRowCount - 1); + + Assert.Equal(0, mtpState.CurrentLength); + + using ITensor draft0 = model.ForwardMtp(mtpState, tokenId: tokenIds[^1], position: 2); + Assert.Equal(1, draft0.Shape[0]); + Assert.Equal(config.VocabSize, draft0.Shape[1]); + Assert.Equal(1, mtpState.CurrentLength); + AssertAllFinite(draft0, config.VocabSize); + + // Second autoregressive MTP step, seeded from the head's own output (not the trunk's). + int argmax0 = ArgMax(draft0, config.VocabSize); + using ITensor draft1 = model.ForwardMtp(mtpState, tokenId: argmax0, position: 3); + Assert.Equal(2, mtpState.CurrentLength); + AssertAllFinite(draft1, config.VocabSize); + } + + [SkippableFact] + public void ForwardMtp_Deterministic_SameInputsSameOutputs() + { + Skip.IfNot(IsCudaDriverPresent(), "No CUDA GPU available"); + string? ptxDir = FindPtxDir(); + Skip.If(ptxDir is null, "PTX files not found"); + + string path = WriteFixture(withMtp: true); + + float[] logitsA = RunSingleMtpStep(path, ptxDir!); + float[] logitsB = RunSingleMtpStep(path, ptxDir!); + + Assert.Equal(logitsA.Length, logitsB.Length); + for (int i = 0; i < logitsA.Length; i++) + Assert.Equal(logitsA[i], logitsB[i]); + } + + [SkippableFact] + public void ForwardMtp_WithoutOwnHeadTensors_FallsBackToTrunkHeadAndEmbedding() + { + Skip.IfNot(IsCudaDriverPresent(), "No CUDA GPU available"); + string? ptxDir = FindPtxDir(); + Skip.If(ptxDir is null, "PTX files not found"); + + string path = WriteFixture(withMtp: true, mtpHasOwnHeadTensors: false, name: "qwen35-mtp-noheadtensors.gguf"); + using var gguf = GgufFile.Open(path); + var config = GgufModelConfigExtractor.Extract(gguf.Metadata); + using var model = CudaQwen3HybridDenseTransformerModel.LoadFromGguf(gguf, config, deviceId: 0, ptxDir); + Assert.True(model.SupportsMtp); + + int[] tokenIds = [0, 1]; + int[] positions = [0, 1]; + using var kvCache = model.CreateKvCache(maxSeqLen: 64); + using var mtpState = (CudaMtpState)model.CreateMtpState()!; + using (ITensor _ = model.Forward(tokenIds, positions, deviceId: -1, kvCache, adapter: null, mtpState)) { } + mtpState.SeedFromCapturedRow(mtpState.CapturedRowCount - 1); + + using ITensor draft = model.ForwardMtp(mtpState, tokenId: tokenIds[^1], position: 1); + AssertAllFinite(draft, config.VocabSize); + } + + // ── Engine-level integration: MtpSpeculativeDecoder (unmodified, backend-agnostic) against the + // real CUDA model ────────────────────────────────────────────────────────────────────────── + + /// + /// The central claim under test, ported from MtpSpeculativeDecoderTests's mock-model + /// version to the REAL CUDA model: MTP self-speculative decoding must produce the exact same + /// output token sequence as plain greedy decode of the target model alone. Unlike the mock + /// version (which hand-crafts "MTP disagrees with target" logits), the real MTP head and trunk + /// are independently-weighted projections over an untrained (random-weight) tiny model, so + /// disagreements between the MTP head's own guess and the trunk's argmax arise naturally — + /// this exercises both the accept and the reject/correction path without needing a hand-crafted + /// mock, while also proving (completely + /// unmodified by this CUDA work) really is backend-agnostic against a real CUDA . + /// + /// + /// Uses fullAttnInterval: 1 (an all-full-attention trunk, no GDN layer) deliberately — + /// NOT the default mixed GDN+attention fixture the other tests in this class use. A separate, + /// pre-existing gap was found while first writing this test against the default (GDN-containing) + /// fixture: 's (and the equivalent two-model + /// SpeculativeDecoder's) verify-batch forward runs SEVERAL candidate tokens through the + /// trunk in one call, and rejected ones still permanently advance any recurrent (GDN/Mamba) + /// layer's state — 's own doc says outright "it has no position + /// indexing", so unlike the attention KV-cache's Rollback, there is no way to undo that + /// once the batched forward has run. This is an engine-layer limitation of speculative decoding + /// against recurrent trunks in general (present identically on CPU — verified separately, not + /// introduced by this CUDA work, and out of scope to fix here: it needs a real GDN-state + /// checkpoint/restore design threaded through both decoders across CPU/CUDA/Vulkan). Isolating + /// THIS test to an all-full-attention trunk keeps it a clean proof of the actual CUDA MTP + /// correctness this task is scoped to; see the class remarks / final report for the full + /// finding and recommendation to file it as its own follow-up issue. + /// + [SkippableFact] + public void DraftAndVerify_MatchesPlainGreedyDecode_OnRealCudaModel() + { + Skip.IfNot(IsCudaDriverPresent(), "No CUDA GPU available"); + string? ptxDir = FindPtxDir(); + Skip.If(ptxDir is null, "PTX files not found"); + + string path = SyntheticQwen35HybridDenseMtpGguf.Write( + Path.Combine(_scratch, "qwen35-mtp-fullattn.gguf"), withMtp: true, fullAttnInterval: 1); + + const int startToken = 1; + const int totalNewTokens = 12; + const int k = 3; + + List speculative = RunSpeculative(path, ptxDir!, startToken, totalNewTokens, k); + List plain = RunPlainGreedy(path, ptxDir!, startToken, totalNewTokens); + + _out.WriteLine($"plain: {string.Join(",", plain)}"); + _out.WriteLine($"speculative: {string.Join(",", speculative)}"); + + Assert.Equal(plain, speculative); + } + + private static float[] RunSingleMtpStep(string path, string ptxDir) + { + using var gguf = GgufFile.Open(path); + var config = GgufModelConfigExtractor.Extract(gguf.Metadata); + using var model = CudaQwen3HybridDenseTransformerModel.LoadFromGguf(gguf, config, deviceId: 0, ptxDir); + + int[] tokenIds = [0, 1, 2]; + int[] positions = [0, 1, 2]; + using var kvCache = model.CreateKvCache(maxSeqLen: 64); + using var mtpState = (CudaMtpState)model.CreateMtpState()!; + using (ITensor _ = model.Forward(tokenIds, positions, deviceId: -1, kvCache, adapter: null, mtpState)) { } + mtpState.SeedFromCapturedRow(mtpState.CapturedRowCount - 1); + + using ITensor draft = model.ForwardMtp(mtpState, tokenId: tokenIds[^1], position: 2); + unsafe + { + var span = new ReadOnlySpan((void*)draft.DataPointer, config.VocabSize); + return span.ToArray(); + } + } + + private static List RunPlainGreedy(string path, string ptxDir, int startToken, int totalNewTokens) + { + using var gguf = GgufFile.Open(path); + var config = GgufModelConfigExtractor.Extract(gguf.Metadata); + using var model = CudaQwen3HybridDenseTransformerModel.LoadFromGguf(gguf, config, deviceId: 0, ptxDir); + using var kvCache = model.CreateKvCache(maxSeqLen: 64); + + var seq = new List { startToken }; + int cur = startToken; + for (int pos = 0; pos < totalNewTokens; pos++) + { + using ITensor logits = model.Forward([cur], [pos], deviceId: -1, kvCache); + int argmax = ArgMax(logits, config.VocabSize); + seq.Add(argmax); + cur = argmax; + } + return seq; + } + + private static List RunSpeculative( + string path, string ptxDir, int startToken, int totalNewTokens, int k) + { + using var gguf = GgufFile.Open(path); + var config = GgufModelConfigExtractor.Extract(gguf.Metadata); + using var model = CudaQwen3HybridDenseTransformerModel.LoadFromGguf(gguf, config, deviceId: 0, ptxDir); + Assert.True(model.SupportsMtp); + + var decoder = new MtpSpeculativeDecoder(greedy: true); + var pipeline = new SamplerPipeline(new InferenceOptions { Temperature = 0f }); + + var generatedIds = new List { startToken }; + using var kvCache = model.CreateKvCache(maxSeqLen: 64); + using var mtpState = (CudaMtpState)model.CreateMtpState()!; + + // Prefill: seed the target KV-cache with the start token at position 0. DraftAndVerify's + // own contract (see its remarks: "lastToken already occupies position in kvCacheTarget") + // means `position` must equal lastToken's OWN KV-cache slot, not "prompt length + decoded + // so far" naively read as generatedIds.Count -- so it starts at 0 (matching the prefill + // call just above, not 1). Confirmed against SpeculativeDecoder's identical convention + // (verifyPositions[0] = position, holding lastToken). The CPU mock-model engine tests this + // was originally copied from never caught an off-by-one here because their mock Forward + // ignores position entirely for logit computation; this real (position-sensitive, + // RoPE-using) CUDA model does not tolerate it. + using (ITensor _ = model.Forward([startToken], [0], deviceId: -1, kvCache)) { } + + int position = 0; + Span outputBuffer = stackalloc int[k + 1]; + int guard = 0; + while (generatedIds.Count - 1 < totalNewTokens && guard++ < totalNewTokens * 4) + { + var result = decoder.DraftAndVerify( + model, kvCache, mtpState, pipeline, generatedIds, + constraint: null, position, vocabSize: config.VocabSize, numCandidates: k, outputBuffer); + + Assert.True(result.AcceptedCount > 0, "Every round must accept at least the corrected/bonus token."); + + for (int i = 0; i < result.AcceptedCount && generatedIds.Count - 1 < totalNewTokens; i++) + generatedIds.Add(outputBuffer[i]); + + position += result.AcceptedCount; + } + + return generatedIds.Take(totalNewTokens + 1).ToList(); + } + + private static void AssertAllFinite(ITensor tensor, int vocabSize) + { + unsafe + { + var span = new ReadOnlySpan((void*)tensor.DataPointer, vocabSize); + foreach (float v in span) + Assert.True(float.IsFinite(v), $"non-finite logit: {v}"); + } + } + + private static int ArgMax(ITensor tensor, int vocabSize) + { + unsafe + { + var span = new ReadOnlySpan((void*)tensor.DataPointer, vocabSize); + int best = 0; + for (int i = 1; i < span.Length; i++) + if (span[i] > span[best]) best = i; + return best; + } + } +} diff --git a/tests/DotLLM.Tests.Unit/Engine/MtpSpeculativeDecoderTests.cs b/tests/DotLLM.Tests.Unit/Engine/MtpSpeculativeDecoderTests.cs index 19189424..49e9b3ae 100644 --- a/tests/DotLLM.Tests.Unit/Engine/MtpSpeculativeDecoderTests.cs +++ b/tests/DotLLM.Tests.Unit/Engine/MtpSpeculativeDecoderTests.cs @@ -155,13 +155,18 @@ private static List RunSpeculative( using var kvCache = new SimpleKvCache(1, NumKvHeads, HeadDim, MaxSeqLen); using var mtpState = new CpuMtpState(HiddenSize, MtpNumKvHeads, MtpHeadDim, maxSteps: Math.Max(k, 1) + 4); - // Prefill: seed the target KV-cache with the start token at position 0, matching the - // "position = prompt length + decoded so far" contract DraftAndVerify documents. + // Prefill: seed the target KV-cache with the start token at position 0. `position` must + // equal lastToken's OWN KV-cache slot (see DraftAndVerify's remarks: "lastToken already + // occupies position in kvCacheTarget"), matching SpeculativeDecoder's identical + // convention -- so it starts at 0 here, not 1. This mock model ignores position entirely + // for logit computation, so this correction has no effect on this test's assertions (both + // values pass); fixed for realism after a real (position-sensitive) CUDA model test caught + // an off-by-one derived from copying this exact pattern -- see issue #253's CUDA follow-up. using (ITensor _ = model.Forward([startToken], [0], deviceId: -1, kvCache)) { } - int position = 1; + int position = 0; Span outputBuffer = stackalloc int[k + 1]; int guard = 0; while (generatedIds.Count - 1 < totalNewTokens && guard++ < totalNewTokens * 4)