Restore validated CUDA forward fixes on tts-arena/main

driver/cuda/src/entry.cpp

@@ -626,7 +626,8 @@ int run_impl(int argc,
         // concurrent speculation.
         const int q35_planned_slots = std::max<int>(1, mem_plan.state_slots);
         qwen3_5_runtime_rs_slots = q35_planned_slots;
-        const int q35_alloc_slots = qwen3_5_runtime_rs_slots;
+        const int q35_alloc_slots =
+            qwen3_5_runtime_rs_slots + (graph_pad_slot >= 0 ? 1 : 0);
         qwen3_5_state_cache = pie_cuda_driver::RecurrentStateCache::allocate(
             qwen3_5_layer_is_linear, conv_dim, cfg_q.linear_conv_kernel_dim,
             local_linear_value_heads,

driver/cuda/src/executor/executor.cpp

@@ -2130,7 +2130,7 @@ inline void write_probes(pie_driver::PieForwardResponseView& out,
 }
 }  // namespace
 
-void handle_fire_batch(
+bool handle_fire_batch(
     std::uint32_t req_id,
     const pie_driver::PieForwardRequestView& view,
     pie_driver::PieForwardResponseView& out_resp,
@@ -2276,6 +2276,7 @@ void handle_fire_batch(
 
         const int N = static_cast<int>(tok_view.size());
         const int num_sampling = static_cast<int>(sidx_view.size());
+        const bool has_sampling = num_sampling > 0;
         dbg_R = R; dbg_N = N;
 
         // Qwen3-VL: assemble the per-token [N,3] M-RoPE positions. Text rows
@@ -2315,13 +2316,12 @@ void handle_fire_batch(
             // Empty batch — emit a zero-request response view.
             std::vector<pie_driver::PerRequestOutput> empty(std::max(R, 0));
             executor.response_builder.build(empty, out_resp);
-            return;
+            return true;
         }
         if (N > max_workspace_tokens) {
             std::cerr << "[pie-driver-cuda] batch tokens=" << N
                       << " exceeds workspace=" << max_workspace_tokens << "\n";
-            out_resp = pie_driver::PieForwardResponseView{};
-            return;
+            return false;
         }
 
         // Compute max KV length across requests for shmem sizing.
@@ -2438,6 +2438,16 @@ void handle_fire_batch(
                 slot_ids_h[r] = executor.graph_pad_slot;
                 is_fresh_h[r] = 0u;
             }
+            const int max_rs_slots = executor.rs_cache->max_slots();
+            for (int r = 0; r < slot_count; ++r) {
+                if (slot_ids_h[r] < 0 || slot_ids_h[r] >= max_rs_slots) {
+                    throw std::runtime_error(
+                        "rs_cache slot id " + std::to_string(slot_ids_h[r]) +
+                        " for request " + std::to_string(r) +
+                        " is outside allocated slot range [0, " +
+                        std::to_string(max_rs_slots) + ")");
+                }
+            }
             pi.slot_ids.copy_from_host(std::span<const std::int32_t>(slot_ids_h));
             pi.is_fresh.copy_from_host(std::span<const std::uint8_t>(is_fresh_h));
         }
@@ -2520,16 +2530,14 @@ void handle_fire_batch(
                       << logit_rows_required
                       << " logit rows, exceeding workspace capacity "
                       << tensor_rows(ws.logits) << "\n";
-            out_resp = pie_driver::PieForwardResponseView{};
-            return;
+            return false;
         }
         if (prob_rows_required > tensor_rows(ws.probs)) {
             std::cerr << "[pie-driver-cuda] fire_batch needs "
                       << prob_rows_required
                       << " probability rows, exceeding workspace capacity "
                       << tensor_rows(ws.probs) << "\n";
-            out_resp = pie_driver::PieForwardResponseView{};
-            return;
+            return false;
         }
 
         const SamplingPlan sample_plan{
@@ -2677,7 +2685,7 @@ void handle_fire_batch(
         // pass) produces no logits and no sampled tokens — skip every argmax /
         // sampling launch (they would read the unwritten, undersized
         // `ws.logits` over all N rows and fault).
-        if (num_sampling == 0) {
+        if (!has_sampling) {
             // nothing to sample
         } else if (tp_greedy_argmax) {
             kernels::launch_select_global_argmax_pairs(
@@ -2733,19 +2741,25 @@ void handle_fire_batch(
         // `h_per_*`, `h_sample_idx`) are still in scope here for the
         // response builder.
 
-        // Only copy the first N entries — `pi.sampled` is sized for
-        // max_workspace_tokens, but only [0, N) are valid this fire.
-        // Async on the same stream the sampler ran on so it slots into
-        // the stream's FIFO; we sync immediately after because the
-        // response payload depends on these tokens. (Future work moves
-        // the sync past the host-side response-prep so the host
-        // and GPU can overlap.)
-        std::int32_t* sampled_host =
-            sampled_pinned_buf(static_cast<std::size_t>(N));
-        CUDA_CHECK(cudaMemcpyAsync(sampled_host, pi.sampled.data(),
-                                   sizeof(std::int32_t) * N,
-                                   cudaMemcpyDeviceToHost,
-                                   cublas.stream()));
+        // Only sampled fires produce valid `pi.sampled` rows. Prefill-only
+        // fires (for example Context::flush) still need the stream sync below
+        // to observe KV writes, but must not copy sampled rows that no kernel
+        // wrote.
+        std::int32_t* sampled_host = nullptr;
+        if (has_sampling) {
+            // Copy the first N entries — `pi.sampled` is sized for
+            // max_workspace_tokens, but only [0, N) are valid this fire.
+            // Async on the same stream the sampler ran on so it slots into
+            // the stream's FIFO; we sync immediately after because the
+            // response payload depends on these tokens. (Future work moves
+            // the sync past the host-side response-prep so the host
+            // and GPU can overlap.)
+            sampled_host = sampled_pinned_buf(static_cast<std::size_t>(N));
+            CUDA_CHECK(cudaMemcpyAsync(sampled_host, pi.sampled.data(),
+                                       sizeof(std::int32_t) * N,
+                                       cudaMemcpyDeviceToHost,
+                                       cublas.stream()));
+        }
         const auto t_kernel_launch_end = clock::now();
         verify_timer.finish(cublas.stream());
         CUDA_CHECK(cudaStreamSynchronize(cublas.stream()));
@@ -2758,7 +2772,7 @@ void handle_fire_batch(
         // and benchmark traffic keep the GPU sampler result.
         const std::int32_t* all_sampled = sampled_host;
         std::vector<std::int32_t> sampled_override;
-        if (!logit_masks_view.empty()) {
+        if (has_sampling && !logit_masks_view.empty()) {
             sampled_override.assign(sampled_host, sampled_host + N);
             apply_logit_mask_overrides(
                 ws, sampled_override, logit_masks_view, logit_mask_indptr_view,
@@ -2794,7 +2808,7 @@ void handle_fire_batch(
                           << " sampled=" << num_sampling
                           << " max_kv=" << max_kv_len << "\n";
             }
-            return;
+            return true;
         }
 
         // Flat-path arrays: token sampler is the only slot type allowed
@@ -3186,7 +3200,7 @@ void handle_fire_batch(
                           << " sampled=" << num_sampling
                           << " max_kv=" << max_kv_len << "\n";
             }
-            return;
+            return true;
         }
 
         executor.response_builder.build_token_only(
@@ -3202,13 +3216,12 @@ void handle_fire_batch(
                       << " sampled=" << num_sampling
                       << " max_kv=" << max_kv_len << "\n";
         }
-        return;
+        return true;
 
     } catch (const std::exception& e) {
         std::cerr << "[pie-driver-cuda] fire_batch failed for req_id="
                   << req_id << ": " << e.what() << "\n";
-        out_resp = pie_driver::PieForwardResponseView{};
-        return;
+        return false;
     }
 }
 

driver/cuda/src/executor/executor.hpp

@@ -349,7 +349,7 @@ struct Executor {
 // cumulative fire counter used as PRNG offset and logging-cadence gate.
 // The caller's inproc transport hands `out_resp` to `send_response`
 // immediately after this returns.
-void handle_fire_batch(
+bool handle_fire_batch(
     std::uint32_t req_id,
     const pie_driver::PieForwardRequestView& view,
     pie_driver::PieForwardResponseView& out_resp,

driver/cuda/src/kernels/gated_delta_net.cu

@@ -104,15 +104,14 @@ bool qwen_gdn_fused_step_enabled() {
 // FLA chunked-prefill default), so strictly less quantization than
 // the legacy per-element-round kernel.
 //
-// Default ON: +32% end-to-end throughput on Qwen/Qwen3.5-4B
-// (6924 -> 9166 tok/s). The launcher only routes here for the GQA
-// bf16 V-last decode shape (V_d==K_d==128, !k_last); everything else
-// falls back to the legacy kernel. Set PIE_QWEN35_GDN_SMEM_STEP=0 to
-// force the fallback.
+// Opt-in: this path is faster on saturated decode microbenches, but it
+// has not yet had the same forked/batched-context parity coverage as the
+// legacy BF16 GQA kernel. Keep the production default on the covered path;
+// set PIE_QWEN35_GDN_SMEM_STEP=1 to benchmark the SMEM variant.
 bool qwen_gdn_smem_step_enabled() {
     static const bool enabled = [] {
         const char* v = std::getenv("PIE_QWEN35_GDN_SMEM_STEP");
-        if (v == nullptr || v[0] == '\0') return true;
+        if (v == nullptr || v[0] == '\0') return false;
         return v[0] != '0';
     }();
     return enabled;

driver/cuda/src/model/llama_like_model.cpp

@@ -48,8 +48,11 @@ void LlamaLikeModel::body(Qwen3Workspace& ws,
                           AttentionWorkspace& attn_ws,
                           ops::CublasHandle& cublas,
                           const ForwardFn::ForwardInputs& in) {
+    LlamaLikeForwardCfg fwd = fwd_cfg_;
+    fwd.emit_logits = in.emit_logits;
+
     llama_like_forward_paged(
-        weights_, hf_config_, fwd_cfg_, plan_,
+        weights_, hf_config_, fwd, plan_,
         ws, kv, attn_ws, cublas,
         in.token_ids, in.positions,
         in.qo_indptr_d, in.kv_page_indices_d, in.kv_page_indptr_d,

driver/cuda/src/model/qwen3_5_forward.cpp

@@ -1367,12 +1367,11 @@ void qwen3_5_forward_paged(
     }
 
     // 3. Final norm.
-    // State-only forward (num_logit_rows < 0): the speculative repair re-runs
-    // the backbone purely to advance the rs_cache state (written in the layer
-    // loop above). Its logits/hidden are discarded, so skip final_norm,
-    // lm_head, and the final hidden copy entirely — that dense lm_head over all
-    // N repair tokens is the single largest piece of wasted work.
-    if (num_logit_rows < 0 || commit_advance) {
+    // State-only/no-sampler forward: the caller only needs KV/recurrent-state
+    // side effects, so skip final_norm, lm_head, and the final hidden copy.
+    // This covers speculative repair (num_logit_rows < 0) and ordinary
+    // Context::flush/KV-fill fires (emit_logits=false).
+    if (!fwd_cfg.emit_logits || num_logit_rows < 0 || commit_advance) {
         profile.end(stream);
         maybe_print_forward_profile(profile);
         return;

driver/cuda/src/model/qwen3_5_forward.hpp

@@ -87,6 +87,12 @@ struct Qwen3_5ForwardCfg {
     // cache lookup should stay pinned to the verified source prefix while
     // draft positions advance for RoPE and history masking.
     bool mtp_global_cache_uses_prefix_position = false;
+
+    // False for KV-fill/flush fires that have no sampler rows. Those fires
+    // only need cache/state side effects; materializing dense logits is both
+    // wasted work and can violate the executor's no-sampling workspace
+    // contract.
+    bool emit_logits = true;
 };
 
 // Persistent decode-plan cache. Owned by main.cpp's serving setup so

driver/cuda/src/model/qwen3_5_model.cpp

@@ -47,8 +47,11 @@ void Qwen35Model::body(Qwen3Workspace& ws,
                        AttentionWorkspace& attn_ws,
                        ops::CublasHandle& cublas,
                        const ForwardFn::ForwardInputs& in) {
+    Qwen3_5ForwardCfg fwd = fwd_cfg_;
+    fwd.emit_logits = in.emit_logits;
+
     qwen3_5_forward_paged(
-        weights_, hf_config_, fwd_cfg_, plan_state_,
+        weights_, hf_config_, fwd, plan_state_,
         ws, la_ws_, kv, state_cache_,
         attn_ws, cublas,
         in.token_ids, in.positions,

driver/cuda/src/model/qwen3_5_moe_forward.cpp

@@ -1774,10 +1774,8 @@ void qwen3_5_moe_forward_paged(
         }
     }
 
-    // State-only forward (num_logit_rows < 0) or recurrent-only commit-advance:
-    // logits/hidden are discarded, so skip final_norm + lm_head + final copy.
-    // See qwen3_5_forward.
-    if (num_logit_rows < 0 || commit_advance) {
+    // State-only/no-sampler forward: see qwen3_5_forward.
+    if (!fwd_cfg.emit_logits || num_logit_rows < 0 || commit_advance) {
         profile.end(stream);
         maybe_print_profile(profile);
         return;

driver/cuda/src/model/qwen3_5_moe_model.cpp

@@ -49,8 +49,11 @@ void Qwen35MoeModel::body(Qwen3Workspace& ws,
                           AttentionWorkspace& attn_ws,
                           ops::CublasHandle& cublas,
                           const ForwardFn::ForwardInputs& in) {
+    Qwen3_5ForwardCfg fwd = fwd_cfg_;
+    fwd.emit_logits = in.emit_logits;
+
     qwen3_5_moe_forward_paged(
-        weights_, hf_config_, fwd_cfg_, plan_state_,
+        weights_, hf_config_, fwd, plan_state_,
         ws, la_ws_, moe_ws_, kv, state_cache_,
         attn_ws, cublas,
         in.token_ids, in.positions,

driver/cuda/src/service/inproc_service.cpp

@@ -36,9 +36,27 @@ void InProcService::serve_forever(pie_driver::InProcServer& server) {
             switch (req.method) {
                 case pie_driver::PIE_METHOD_FORWARD: {
                     ++handled_;
-                    handle_fire_batch(
+                    const bool ok = handle_fire_batch(
                         req_id, req.forward, out.forward, executor_, handled_);
-                    out.status = 0;
+                    if (ok) {
+                        const auto qo = req.forward.qo_indptr.as<std::uint32_t>();
+                        const auto expected = qo.empty()
+                            ? 0u
+                            : static_cast<std::uint32_t>(qo.size() - 1);
+                        if (out.forward.num_requests == expected) {
+                            out.status = 0;
+                        } else {
+                            std::cerr << "[pie-driver-cuda] fire_batch response count "
+                                      << "mismatch for req_id=" << req_id
+                                      << ": expected " << expected
+                                      << ", got " << out.forward.num_requests << "\n";
+                            out.method = pie_driver::PIE_METHOD_HEALTH;
+                            out.status = -1;
+                        }
+                    } else {
+                        out.method = pie_driver::PIE_METHOD_HEALTH;
+                        out.status = -1;
+                    }
                     break;
                 }
                 case pie_driver::PIE_METHOD_COPY_D2H: {