Restore DeepSeek4 HIP backend safely

server/CMakeLists.txt

@@ -401,19 +401,6 @@ if(DFLASH27B_GPU_BACKEND STREQUAL "hip")
     # rms_norm_hip.cu is needed by the HIP chunk-B graph path regardless of SM80_EQUIV.
     target_sources(dflash_common PRIVATE src/rms_norm_hip.cu)
     set_source_files_properties(src/rms_norm_hip.cu PROPERTIES LANGUAGE HIP)
-    # GPU draft top-K + log-prob kernel (DFlash). The shared body in
-    # geometric_draft_topk_cuda.cu compiles unchanged for HIP through the
-    # hip_compat <cuda_runtime.h> shim, so it is added directly with LANGUAGE HIP
-    # (same shared-.cu pattern as deepseek4_hc_cuda.cu above — no separate HIP
-    # translation unit). This makes DFLASH_GPU_DRAFT_TOPK a real path on ROCm
-    # instead of a no-op, so AMD stops paying the per-step vocab x n_tokens D2H +
-    # CPU heap extract. The hip_compat include dir is already on dflash_common.
-    target_sources(dflash_common PRIVATE src/common/geometric_draft_topk_cuda.cu)
-    set_source_files_properties(src/common/geometric_draft_topk_cuda.cu
-        PROPERTIES LANGUAGE HIP)
-    # PUBLIC so test consumers (test_dflash / test_draft_topk_cuda) also take the
-    # GPU draft top-K path instead of the CPU fallback.
-    target_compile_definitions(dflash_common PUBLIC DFLASH27B_HAVE_DRAFT_TOPK=1)
     if(DFLASH27B_HIP_SM80_EQUIV)
         find_path(DFLASH27B_ROCWMMA_INCLUDE_DIR rocwmma/rocwmma.hpp
             HINTS "${_dflash_rocm_root}/include" /opt/rocm/include
@@ -447,9 +434,8 @@ elseif(DFLASH27B_GPU_BACKEND STREQUAL "cuda")
         src/flashprefill.cpp
         src/common/geometric_draft_topk_cuda.cu)
     # PUBLIC so consumers (e.g. the test_dflash executable) also see the macro
-    # and take the GPU draft top-K path instead of the CPU fallback. Same macro
-    # name as the HIP branch above (backend-neutral).
-    target_compile_definitions(dflash_common PUBLIC DFLASH27B_HAVE_DRAFT_TOPK=1)
+    # and take the GPU draft top-K path instead of the CPU fallback.
+    target_compile_definitions(dflash_common PUBLIC DFLASH27B_HAVE_DRAFT_TOPK_CUDA=1)
     # GPU port of the sample_logits chain. Compiled in by default; the path is
     # then opted into at runtime via the DFLASH_GPU_SAMPLE env var. Turn the
     # whole thing off at configure time with -DDFLASH_GPU_SAMPLER=OFF.
@@ -677,40 +663,13 @@ if(DFLASH27B_TESTS)
         target_link_libraries(test_rms_norm_hip PRIVATE dflash_common ${DFLASH27B_GGML_BACKEND_TARGET})
         add_test(NAME rms_norm_hip COMMAND test_rms_norm_hip)
     endif()
-    # DS4 Hierarchical-Controller pre-mix kernels vs CPU reference. HIP only:
-    # deepseek4_hc_cuda.cu is compiled into dflash_common on this backend. Validates
-    # the HC mix matvec + NHC=4 Sinkhorn path against the device's actual wave width.
-    if(DFLASH27B_GPU_BACKEND STREQUAL "hip" AND EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_deepseek4_hc_cuda.cpp")
-        add_executable(test_deepseek4_hc_cuda test/test_deepseek4_hc_cuda.cpp)
-        set_source_files_properties(test/test_deepseek4_hc_cuda.cpp PROPERTIES LANGUAGE HIP)
-        set_target_properties(test_deepseek4_hc_cuda PROPERTIES HIP_ARCHITECTURES "${_dflash_archs}")
-        target_include_directories(test_deepseek4_hc_cuda PRIVATE
-            ${CMAKE_CURRENT_SOURCE_DIR}/src
-            ${CMAKE_CURRENT_SOURCE_DIR}/hip_compat)
-        target_link_libraries(test_deepseek4_hc_cuda PRIVATE dflash_common ${DFLASH27B_GGML_BACKEND_TARGET})
-        add_test(NAME deepseek4_hc_cuda COMMAND test_deepseek4_hc_cuda)
-    endif()
-    # GPU draft top-K kernel vs CPU reference (extract_draft_topk). Built on both
-    # backends: geometric_draft_topk_cuda.cu is compiled into dflash_common on the
-    # cuda backend directly and on hip via LANGUAGE HIP + the hip_compat shim.
+    # GPU draft top-K kernel vs CPU reference (extract_draft_topk). CUDA only:
+    # geometric_draft_topk_cuda.cu is compiled into dflash_common solely on the cuda backend.
     if(DFLASH27B_GPU_BACKEND STREQUAL "cuda" AND EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_draft_topk_cuda.cpp")
         add_executable(test_draft_topk_cuda test/test_draft_topk_cuda.cpp)
         target_include_directories(test_draft_topk_cuda PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src)
         target_link_libraries(test_draft_topk_cuda PRIVATE dflash_common CUDA::cudart)
         add_test(NAME draft_topk_cuda COMMAND test_draft_topk_cuda)
-    elseif(DFLASH27B_GPU_BACKEND STREQUAL "hip" AND EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_draft_topk_cuda.cpp")
-        # HIP build of the same GPU-vs-CPU parity test. The test source uses CUDA
-        # spellings (<cuda_runtime.h> + cudaMalloc/cudaMemcpy/...); the hip_compat
-        # shim maps them onto HIP, exactly as the kernel TU does. Lets the draft
-        # top-K port be validated on AMD (gfx1100 / gfx1151 / gfx12xx).
-        add_executable(test_draft_topk_cuda test/test_draft_topk_cuda.cpp)
-        set_source_files_properties(test/test_draft_topk_cuda.cpp PROPERTIES LANGUAGE HIP)
-        set_target_properties(test_draft_topk_cuda PROPERTIES HIP_ARCHITECTURES "${_dflash_archs}")
-        target_include_directories(test_draft_topk_cuda PRIVATE
-            ${CMAKE_CURRENT_SOURCE_DIR}/src
-            ${CMAKE_CURRENT_SOURCE_DIR}/hip_compat)
-        target_link_libraries(test_draft_topk_cuda PRIVATE dflash_common ${DFLASH27B_GGML_BACKEND_TARGET})
-        add_test(NAME draft_topk_cuda COMMAND test_draft_topk_cuda)
     endif()
     # GPU port of the sample_logits chain vs the CPU reference. CUDA only:
     # geometric_sampler_cuda.cu is compiled into dflash_common solely on the cuda backend.

server/src/common/backend_factory.cpp

@@ -220,29 +220,26 @@ std::unique_ptr<ModelBackend> create_backend(const BackendArgs & args) {
                 ? compiled_placement_backend()
                 : args.device.backend;
 
-        // HIP single-device launches cannot rely on the CUDA/Halo auto-split
-        // path; use the single-backend loader, which can fall back to hybrid
-        // expert placement when a full monolithic load does not fit.
         if (target_backend == PlacementBackend::Hip &&
             !args.device.is_layer_split() &&
             !args.remote_target_shard.enabled()) {
             DeepSeek4BackendConfig cfg;
             cfg.model_path = args.model_path;
-            cfg.device     = args.device;
-            cfg.stream_fd  = args.stream_fd;
-            cfg.max_ctx    = args.device.max_ctx;
-            cfg.chunk      = args.chunk;
+            cfg.device = args.device;
+            cfg.max_ctx = args.device.max_ctx;
+            cfg.chunk = args.chunk;
 
             auto backend = std::make_unique<DeepSeek4Backend>(cfg);
             if (!backend->init()) {
-                std::fprintf(stderr, "[backend_factory] DeepSeek4Backend init failed\n");
+                std::fprintf(stderr, "[backend_factory] DeepSeek4Backend full-weight init failed\n");
                 return nullptr;
             }
             return backend;
         }
 
-        // CUDA builds keep the layer-split backend so they can auto-split
-        // across CUDA and remote HIP target shards.
+        // CUDA keeps the layer-split backend so prefill stays on the batched
+        // layer-range path. HIP single-device can still load all weights
+        // monolithically, without the removed hybrid expert fallback.
         DeepSeek4LayerSplitAdapterConfig cfg;
         cfg.target_path        = args.model_path;
         cfg.device             = args.device;

server/src/common/geometric_draft_topk_cuda.h

@@ -14,10 +14,8 @@
 //
 // Returns false (caller must fall back to the CPU path) when the GPU runtime is
 // unavailable, the pointer is not device memory, K is out of range, or any
-// device call fails. Compiled on both CUDA and HIP/ROCm builds — this same .cu
-// is compiled directly with LANGUAGE HIP on ROCm, the cuda_runtime.h spellings
-// mapped by the hip_compat shim; guarded by DFLASH27B_HAVE_DRAFT_TOPK. See
-// CMakeLists.txt.
+// device call fails. Compiled on the CUDA backend only, where this .cu is added
+// to dflash_common; guarded by DFLASH27B_HAVE_DRAFT_TOPK_CUDA. See CMakeLists.txt.
 
 #pragma once
 

server/src/common/layer_split_backend.cpp

@@ -85,7 +85,11 @@ GenerateResult LayerSplitBackend::run_from_state(const GenerateRequest & req,
         }
         std::vector<int32_t> chunk(req.prompt.begin() + consumed,
                                    req.prompt.begin() + consumed + n_tokens);
-        if (!adapter_->prefill(chunk, base_pos + consumed, last_tok)) {
+        const bool ends_prompt = consumed + n_tokens == prompt_len;
+        const bool ends_snap = req.snap_pos >= 0 && req.snap_slot >= 0 &&
+            base_pos + consumed + n_tokens == req.snap_pos;
+        const bool need_logits = ends_prompt || ends_snap;
+        if (!adapter_->prefill(chunk, base_pos + consumed, last_tok, need_logits)) {
             result.fail(GenerateErrorCode::PrefillFailed);
             return result;
         }

server/src/common/layer_split_backend.h

@@ -27,7 +27,8 @@ class LayerSplitAdapter {
     virtual void reset_request_state() = 0;
     virtual int prefill_chunk_tokens() const { return 0; }
     virtual bool prefill(const std::vector<int32_t> & prompt,
-                         int base_pos, int & last_tok) = 0;
+                         int base_pos, int & last_tok,
+                         bool need_logits = true) = 0;
     virtual bool decode_ar(int last_tok, int committed, int n_gen,
                            std::vector<int32_t> & out_tokens,
                            const DaemonIO & io) = 0;