feat(qwen3): add DFlash2 Phase 1 path selector

docs/index.md

@@ -35,6 +35,7 @@ Organized by domain (model line / subsystem / playbook / lesson) instead of by l
 | `models/qwen3/prefix-cache.md` | Prefix caching on by default for Qwen3-4B: full-block kvbm radix matching at the executor, suffix-only prefill. Repeated ~1900-token prompt TTFT 141.8 → 16.3ms p50 (8.7×); warm TTFT ≈ TPOT + ~5ms setup. Includes the RoPE scalar-path corruption fix and the drain-the-stream TTFT measurement pitfall. |
 | `models/qwen3/dspark-integration.md` | DeepSeek **DSpark** Phase 1 is implemented for Qwen3-4B: DFlash backbone + rank-256 Markov head, anchor-first DeepSpec layout, one strided argmax-with-bias kernel, PDL polish, and one D2H per draft block. Greedy losslessness passes; 5090 block7 A/B vs matched DFlash shows DSpark +3.6% geomean output tok/s overall (+3–16% on text/code, random synthetic exception) and better accepted-draft distribution (2.52 vs 2.30 draft tokens/round). |
 | `models/qwen3/dflash-speculative-decoding.md` | DFlash speculative decoding behind `--dflash-draft-model-path`, modelled as an optimistic transaction (propose K → verify K+1 span → accept longest argmax prefix + 1 bonus → commit/roll back KV). Lossless up to bf16 tie-flips (bit-identical multi-token accepts; lm-eval gsm8k strict-match identical spec on/off). Single-stream decode 1.82× on 5070 Ti, 1.56× on 5090. Concurrent throughput fixed by batching the draft forward, then a piecewise verify CUDA Graph (dense ops captured, attention eager) closed single-stream: 5090 greedy c1 274 ≈ vLLM 278, c8 1525 > 1240, c16 1834 ≈ 1846 — all batch sizes now ≥ vLLM. Accept measured equal (9.1% vs 8.85%, same drafter); draft-side piecewise graph tracked next. Proposer trait deferred to EAGLE. |
+| `models/qwen3/dflash2-phase1-930.md` | Issue #930 Phase 1 record: bounded top-16 DFlash2 candidate selection; dynamic convolution, sliding-window execution, and sampled rejection remain out of scope. |
 | `models/qwen3/accuracy-gate.md` | Qwen3 size-keyed logits golden gate (all six sizes 0.6B–32B committed) (`tests/hf_golden_gate.rs`): 48 teacher-forced sequences / 816 positions vs a stored HF bf16 golden, replayed over bs=1 / batched eager / CUDA-graph. Strict guards: regret check + mean ≤ 0.06 + p99 ≤ 0.20; absolute max printed but not asserted (coverage-unstable). Methodology in `subsystems/correctness/`. |
 | `models/qwen3/kernels-crate.md` | Phase 1 split implemented and 5090-verified: Qwen3-4B kernel surface lives in `pegainfer-kernels`; release build, test-target compile, accuracy gate, and bench snapshot pass. |
 | `models/qwen3/tp-design.md` | Qwen3 tensor-parallel design: `TP=2` milestone scope plus the controller/worker broadcast execution model, request identity, and coarse-grained step protocol for future TP/MoE work. |

docs/models/qwen3/dflash2-phase1-930.md

@@ -0,0 +1,150 @@
+# DFlash2 Phase 1 Selector
+
+> **TL;DR:** Issue #930 Phase 1 adds a bounded, deterministic top-16 candidate selector on top of the existing Qwen3 DFlash backbone; dynamic convolution, sliding-window execution, and sampled rejection are deliberately out of scope.
+>
+> **Last touched:** 2026-09
+
+## Preparation
+
+- **Read**:
+  - `docs/index.md` - routes Qwen3 model and kernel design records.
+  - `docs/models/qwen3/dflash-speculative-decoding.md` - defines the existing proposer/verify/KV transaction contract and the batch layout.
+  - `docs/models/qwen3/dspark-integration.md` - documents the legacy anchor-first/Markov path that must remain unchanged.
+  - `docs/models/qwen3/kernels-crate.md` - assigns CUDA primitives and FFI ownership to `pegainfer-kernels`.
+  - `docs/models/qwen3/model-crate.md` - documents Qwen3 model-crate boundaries and single-GPU speculative decoding.
+  - `docs/conventions/coding-style.md` - requires focused tests and project logging conventions.
+  - `CLAUDE.md` - defines build, branch, and AI-assisted contribution requirements.
+- **Relevant history**:
+  - `docs/models/qwen3/dflash-speculative-decoding.md` - the existing DFlash lane owns proposal while the shared verify and KV transaction contracts stay method-agnostic.
+  - No prior DFlash2 Phase 1 task record exists in this checkout.
+- **Plan**:
+  1. Audit the current configuration and loader scaffold; keep legacy DFlash and DSpark behavior unchanged, load an independent native output head when the checkpoint declares untied embeddings, and reject hybrid capabilities that Phase 1 cannot execute.
+  2. Load and validate the selector projection/codebooks, add a fixed-size GPU selector primitive and Rust wrapper, and account for its persistent and scratch allocations.
+  3. Dispatch `TopKSelector` from the DFlash draft lane without changing draft span, verify, KV transaction, or CUDA-Graph shapes.
+  4. Run formatting, compile, focused selector/reference checks, GPU-vs-reference checks, and legacy DFlash/DSpark regression checks; record actual results and limitations.
+- **Risks / open questions**:
+  - The only discovered DFlash2 checkpoint also declares Phase 2 convolution and sliding-window capabilities; it must fail closed until those execution paths exist.
+  - Selector tie-breaking, anchor mapping, request-major row offsets, and scratch reservation must be deterministic and shape-safe.
+
+## Execution Log
+
+### Step 1: Normalize the DFlash2 capability contract
+
+- Added a `DFlashProposal::TopKSelector` capability and an explicit
+  `DFlashLayout` in `pegainfer-qwen3/src/config.rs`.
+- Legacy DFlash and DSpark schemas remain on their existing proposal paths.
+- Native DFlash2 configurations parse their root or nested
+  `tie_word_embeddings` field into a head-source contract. Legacy and tied
+  checkpoints reuse the verifier embedding/output projection; untied native
+  checkpoints load only their separate `lm_head.weight`.
+- Phase 2 convolution, sliding-window attention, and anchor-first selector
+  layouts still fail closed before GPU weight allocation.
+
+### Step 2: Load selector weights and wire the proposer
+
+- Added SafeTensors manifest checks for the hidden projection and predecessor /
+  successor codebooks.
+- Added persistent selector scratch and a two-launch CUDA implementation:
+  deterministic top-16 candidate extraction followed by a request-local path
+  walk using the predecessor/successor codebooks.
+- Kept the existing full-block draft result contract, verify span, KV updates,
+  and CUDA-Graph shapes unchanged.
+- The native verifier embedding is intentionally reused instead of loading a
+  duplicate `embed_tokens.weight`; the downloaded Qwen3-4B DFlash2 checkpoint's
+  embedding bytes match the verifier exactly, while its `lm_head.weight` is
+  loaded when the schema is untied.
+
+### Step 3: Fix anchor-drop row mapping
+
+- The DFlash backbone emits an anchor-inclusive block. For the current
+  anchor-drop layout, row 0 is discarded by the executor and rows 1..N-1 are
+  the real proposal positions.
+- The selector now uses compact candidate/output rows for those real positions,
+  while reading the corresponding rows from the original anchor-inclusive
+  logits/hidden buffers. Every request-local walk starts from the verified
+  anchor token, so no draft depends on a candidate from the discarded row 0.
+- The host wrapper reconstructs `[anchor, selected_1, ..., selected_N-1]` for
+  the unchanged executor contract and rejects an invalid GPU token id before
+  it can reach token lookup.
+
+### Step 4: Lightweight verification
+
+Commands were run in the Linux feature checkout
+`/database/ricardo.zheng/projects/open-access/pegainfer`
+with `/usr/bin` present in `PATH` (the build script invokes `git`):
+
+| Command | Result |
+| --- | --- |
+| `cargo fmt --all -- --check` | Passed |
+| `git diff --check` | Passed |
+| `cargo check --release -p pegainfer-qwen3 --tests` | Passed; CUDA `sm_89` build |
+| `cargo test --release -p pegainfer-qwen3 --lib` | 88 passed, 0 failed |
+| `cargo test --release -p pegainfer-build --lib` | 8 passed, 0 failed |
+
+The native selector gate was also run with the real Qwen3-4B DFlash2 tensors
+and a selector-only config overlay. The overlay removes only the checkpoint's
+Phase 2 convolution/sliding-window declarations; it does not replace selector,
+backbone, embedding, or output-head weights:
+
+```bash
+export PATH=/home/ricardo.zheng/.cargo/bin:/usr/local/cuda/bin:/usr/bin:/bin:$PATH
+export CUDA_HOME=/usr/local/cuda
+export PEGAINFER_CUDA_SM=89
+PEGAINFER_TEST_MODEL_PATH=/database/ricardo.zheng/models/Qwen3/Qwen3-4B \
+PEGAINFER_DFLASH2_TEST_MODEL_PATH=/tmp/dflash2-phase1-native-overlay \
+cargo test --release -p pegainfer-qwen3 \
+  --test dflash_speculative_gate \
+  dflash2_native_selector_untied_head_greedy_gate \
+  -- --ignored --nocapture --test-threads=1
+```
+
+Result: `1 passed, 0 failed`; the selector-only native launch dispatched the
+CUDA top-k/path-walk path and matched the plain Qwen3 greedy continuation.
+
+### Step 5: Remove redundant scaffolding
+
+- Kept the selector tensor preflight because the shared loader does not check
+  SafeTensors dtype or malformed rank; removed its unused `SelectorManifest`
+  wrapper and duplicate positive-value checks.
+- Removed the unused selector scratch accessor and ABI-only bf16 assertion.
+- Shortened comments to the anchor mapping, two-launch dependency, and
+  unsupported-capability boundaries.
+- Re-ran formatting, Qwen3/build tests, and the server release build.
+
+| Cleanup verification | Result |
+| --- | --- |
+| `cargo fmt --all -- --check` | Passed |
+| `git diff --check` | Passed |
+| `cargo check --release -p pegainfer-qwen3 --tests` | Passed |
+| `cargo test --release -p pegainfer-qwen3 --lib` | 88 passed, 0 failed |
+| `cargo test --release -p pegainfer-build --lib` | 8 passed, 0 failed |
+| `cargo build --release -p pegainfer-server --bin pegainfer` | Passed |
+
+### Step 6: Address review feedback
+
+- Removed the native-schema-wide head rejection. The loader now distinguishes
+  verifier-owned tied heads from an untied native `lm_head.weight`, so a
+  selector-only DFlash2 checkpoint can reach the selector path.
+- Added an ignored GPU gate that imports a native untied checkpoint, runs the
+  selector CUDA launches with real weights, and checks greedy losslessness.
+  The public Qwen3-4B DFlash2 artifact is currently Phase 2-capable, so the
+  gate uses a config-only overlay while preserving every model tensor.
+
+## Debrief
+
+- **Outcome:** Phase 1 selector wiring, anchor-drop mapping, native untied-head
+  loading, and a focused cleanup of redundant scaffolding are complete in the
+  feature branch. The checkout is based on upstream main; no changes are
+  staged or committed.
+- **Pitfalls encountered:** The first verification command omitted system
+  directories from `PATH`, so `pegainfer-kernels/build.rs` could not spawn
+  `git`. Re-running with `/usr/bin:/bin` succeeded. The row mapping bug was a
+  real semantic issue that compilation alone could not detect.
+- **Lessons learned:** Selector buffers must distinguish the input block shape
+  from the compact set of positions actually proposed. The anchor is a
+  request-level predecessor, not a selector candidate when the executor drops
+  row 0.
+- **Follow-ups:** Run the full native checkpoint import only after Phase 2
+  convolution and sliding-window execution land, because the public checkpoint
+  intentionally advertises those capabilities and Phase 1 rejects them. Phase 3
+  remains responsible for sampled losslessness/rejection sampling.

pegainfer-kernels/csrc/shared/dflash2_selector.cu

@@ -0,0 +1,222 @@
+#include "common.cuh"
+
+#include <climits>
+#include <cmath>
+#include <cstdint>
+#include <cuda.h>
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+
+// DFlash2 keeps a small candidate set per draft position. The selector is
+// deliberately split into two kernels: top-k is embarrassingly parallel over
+// logits, while the path walk has a request-local dependency on the previously
+// selected token. Keeping that dependency out of the top-k kernel makes both
+// launches easy to reason about and keeps the temporary layout graph-safe.
+namespace {
+
+constexpr int SELECTOR_TOP_K = 16;
+constexpr int SELECTOR_TOPK_THREADS = 256;
+constexpr int SELECTOR_WALK_THREADS = 512;
+
+__device__ __forceinline__ bool selector_better(float lhs_value, int lhs_id,
+                                                float rhs_value, int rhs_id) {
+  return lhs_value > rhs_value ||
+         (lhs_value == rhs_value && lhs_id < rhs_id);
+}
+
+__device__ __forceinline__ void selector_insert(float value, int id, float* values, int* ids) {
+  if (!selector_better(value, id, values[SELECTOR_TOP_K - 1],
+                       ids[SELECTOR_TOP_K - 1])) {
+    return;
+  }
+  int slot = SELECTOR_TOP_K - 1;
+  while (slot > 0 &&
+         selector_better(value, id, values[slot - 1], ids[slot - 1])) {
+    values[slot] = values[slot - 1];
+    ids[slot] = ids[slot - 1];
+    --slot;
+  }
+  values[slot] = value;
+  ids[slot] = id;
+}
+
+__global__ void dflash2_selector_topk_kernel(
+    const __nv_bfloat16* __restrict__ logits, uint32_t* __restrict__ ids,
+    float* __restrict__ scores, int rows, int input_block_size,
+    int position_offset, int positions_per_request, int vocab) {
+  // Output rows are compact, while the source logits retain the anchor row.
+  // Translate each compact row back to its request-major input row.
+  const int row = blockIdx.x;
+  if (row >= rows || positions_per_request <= 0) {
+    return;
+  }
+  const int request = row / positions_per_request;
+  const int position = row % positions_per_request;
+  const size_t source_row = static_cast<size_t>(request) * input_block_size +
+                            position_offset + position;
+
+  // Each thread keeps a private top-16 list. The lists occupy 32 KiB of
+  // shared memory and are merged by thread zero in canonical score/id order.
+  __shared__ float local_values[SELECTOR_TOPK_THREADS][SELECTOR_TOP_K];
+  __shared__ int local_ids[SELECTOR_TOPK_THREADS][SELECTOR_TOP_K];
+  float* my_values = local_values[threadIdx.x];
+  int* my_ids = local_ids[threadIdx.x];
+  for (int j = 0; j < SELECTOR_TOP_K; ++j) {
+    my_values[j] = -INFINITY;
+    my_ids[j] = INT_MAX;
+  }
+
+  const __nv_bfloat16* row_logits = logits + source_row * vocab;
+  for (int token = threadIdx.x; token < vocab;
+       token += SELECTOR_TOPK_THREADS) {
+    selector_insert(__bfloat162float(row_logits[token]), token, my_values,
+                    my_ids);
+  }
+  __syncthreads();
+
+  if (threadIdx.x == 0) {
+    float best_values[SELECTOR_TOP_K];
+    int best_ids[SELECTOR_TOP_K];
+    for (int j = 0; j < SELECTOR_TOP_K; ++j) {
+      best_values[j] = -INFINITY;
+      best_ids[j] = INT_MAX;
+    }
+    for (int thread = 0; thread < SELECTOR_TOPK_THREADS; ++thread) {
+      for (int j = 0; j < SELECTOR_TOP_K; ++j) {
+        selector_insert(local_values[thread][j], local_ids[thread][j],
+                        best_values, best_ids);
+      }
+    }
+    for (int j = 0; j < SELECTOR_TOP_K; ++j) {
+      ids[static_cast<size_t>(row) * SELECTOR_TOP_K + j] =
+          static_cast<uint32_t>(best_ids[j]);
+      scores[static_cast<size_t>(row) * SELECTOR_TOP_K + j] = best_values[j];
+    }
+  }
+}
+
+__global__ void dflash2_selector_walk_kernel(
+    const __nv_bfloat16* __restrict__ projected_hidden,
+    const __nv_bfloat16* __restrict__ predecessor,
+    const __nv_bfloat16* __restrict__ successor,
+    const uint32_t* __restrict__ anchor_tokens,
+    const uint32_t* __restrict__ candidate_ids,
+    const float* __restrict__ candidate_unary, uint32_t* __restrict__ output,
+    int requests, int input_block_size, int position_offset,
+    int positions_per_request, int vocab, int rank) {
+  const int request = blockIdx.x;
+  if (request >= requests) {
+    return;
+  }
+
+  __shared__ float edge_scores[SELECTOR_TOP_K];
+  __shared__ uint32_t edge_ids[SELECTOR_TOP_K];
+  __shared__ uint32_t previous;
+  if (threadIdx.x == 0) {
+    previous = anchor_tokens[request];
+  }
+  __syncthreads();
+
+  const int lane = threadIdx.x & 31;
+  const int candidate = threadIdx.x >> 5;
+  for (int position = 0; position < positions_per_request; ++position) {
+    // Candidate/output rows are compact; hidden rows retain the anchor slot.
+    const int row = request * positions_per_request + position;
+    const size_t source_row = static_cast<size_t>(request) * input_block_size +
+                              position_offset + position;
+    if (candidate < SELECTOR_TOP_K) {
+      const uint32_t candidate_id =
+          candidate_ids[static_cast<size_t>(row) * SELECTOR_TOP_K + candidate];
+      float dot = 0.0f;
+      if (candidate_id < static_cast<uint32_t>(vocab) &&
+          previous < static_cast<uint32_t>(vocab)) {
+        const __nv_bfloat16* hidden_row =
+            projected_hidden + source_row * rank;
+        const __nv_bfloat16* predecessor_row =
+            predecessor + static_cast<size_t>(previous) * rank;
+        const __nv_bfloat16* successor_row =
+            successor + static_cast<size_t>(candidate_id) * rank;
+        for (int component = lane; component < rank; component += 32) {
+          dot += __bfloat162float(predecessor_row[component]) *
+                 __bfloat162float(hidden_row[component]) *
+                 __bfloat162float(successor_row[component]);
+        }
+      }
+      dot = warp_reduce_sum(dot);
+      if (lane == 0) {
+        edge_ids[candidate] = candidate_id;
+        edge_scores[candidate] =
+            candidate_unary[static_cast<size_t>(row) * SELECTOR_TOP_K +
+                           candidate] +
+            dot;
+      }
+    }
+    __syncthreads();
+
+    if (threadIdx.x == 0) {
+      uint32_t best_id = edge_ids[0];
+      float best_score = edge_scores[0];
+      for (int j = 1; j < SELECTOR_TOP_K; ++j) {
+        if (selector_better(edge_scores[j], static_cast<int>(edge_ids[j]),
+                            best_score, static_cast<int>(best_id))) {
+          best_score = edge_scores[j];
+          best_id = edge_ids[j];
+        }
+      }
+      output[row] = best_id;
+      previous = best_id;
+    }
+    __syncthreads();
+  }
+}
+
+}  // namespace
+
+extern "C" int dflash2_selector_topk_cuda(
+    const __nv_bfloat16* logits, uint32_t* candidate_ids,
+    float* candidate_scores, int rows, int input_block_size,
+    int position_offset, int positions_per_request, int vocab,
+    cudaStream_t stream) {
+  if (logits == nullptr || candidate_ids == nullptr || candidate_scores == nullptr ||
+      rows <= 0 || input_block_size <= 0 || position_offset < 0 ||
+      positions_per_request <= 0 ||
+      position_offset > input_block_size - positions_per_request ||
+      vocab < SELECTOR_TOP_K) {
+    return static_cast<int>(cudaErrorInvalidValue);
+  }
+  dflash2_selector_topk_kernel<<<rows, SELECTOR_TOPK_THREADS, 0, stream>>>(
+      logits, candidate_ids, candidate_scores, rows, input_block_size,
+      position_offset, positions_per_request, vocab);
+  return static_cast<int>(cudaGetLastError());
+}
+
+extern "C" int dflash2_selector_walk_cuda(
+    const __nv_bfloat16* projected_hidden,
+    const __nv_bfloat16* predecessor,
+    const __nv_bfloat16* successor,
+    const uint32_t* anchor_tokens,
+    const uint32_t* candidate_ids,
+    const float* candidate_unary,
+    uint32_t* output,
+    int requests,
+    int input_block_size,
+    int position_offset,
+    int positions_per_request,
+    int vocab,
+    int rank,
+    cudaStream_t stream) {
+  if (projected_hidden == nullptr || predecessor == nullptr ||
+      successor == nullptr || anchor_tokens == nullptr || candidate_ids == nullptr ||
+      candidate_unary == nullptr || output == nullptr || requests <= 0 ||
+      input_block_size <= 0 || position_offset < 0 ||
+      positions_per_request <= 0 ||
+      position_offset > input_block_size - positions_per_request ||
+      vocab < SELECTOR_TOP_K || rank <= 0) {
+    return static_cast<int>(cudaErrorInvalidValue);
+  }
+  dflash2_selector_walk_kernel<<<requests, SELECTOR_WALK_THREADS, 0, stream>>>(
+      projected_hidden, predecessor, successor, anchor_tokens,
+      candidate_ids, candidate_unary, output, requests, input_block_size,
+      position_offset, positions_per_request, vocab, rank);
+  return static_cast<int>(cudaGetLastError());
+}