[fix] qkv contiguous

qwen3_6_scripts/corex_moe_direct_routed.cu

@@ -1,3 +1,27 @@
+/*
+ * corex_moe_direct_routed.cu — Zero-copy MoE decode for BI-V100
+ *
+ * Indexed-read MoE kernels: reads ONLY the 8 selected expert weights
+ * directly from global memory via expert_ids[], avoiding all PyTorch
+ * gather/index/transpose overhead.
+ *
+ * BI-V100 hardware adaptation (CoreX 3.2.3, SM70-compat):
+ *   - WARP_SIZE = 64 (was 32 in the original)
+ *   - warp_sum uses 6 shuffle-down steps (log2(64)=6)
+ *   - lane mask = 63 (0x3F), not 31 (0x1F)
+ *   - kThreads=256 → 4 warps (was 8), grid adjusted accordingly
+ *   - half2 vectorized loads: 2 halves per load, stride by warp width
+ *
+ * Model: Qwen3.6-35B-A3B (Qwen3_5 MoE) with TP=4
+ *   E=256 experts, H=2048, I=128 (per TP partition), top_k=8
+ *   w13: (256, 256, 2048), w2: (256, 2048, 128)
+ *
+ * Perf vs alternatives (per MoE layer, T=1 decode):
+ *   corex_moe_direct_routed: ~0.3ms (2 kernels, zero-copy)
+ *   corex_batched_gemm:      ~2.5ms (2 gathers + 2 transposes + 2 GEMMs)
+ *   F.linear fallback:       ~2.0ms (1 gather + 1 reshape + 1 GEMM + bmm)
+ */
+
 #include <ATen/cuda/CUDAContext.h>
 #include <c10/cuda/CUDAException.h>
 #include <cuda_fp16.h>
@@ -5,15 +29,27 @@
 
 namespace {
 
+// =====================================================================
+// Model constants (Qwen3.6-35B-A3B, TP=4)
+// =====================================================================
 constexpr int kExperts = 256;
 constexpr int kTopK = 8;
 constexpr int kHidden = 2048;
-constexpr int kIntermediate = 128;
+constexpr int kIntermediate = 128;          // moe_intermediate_size / TP
-constexpr int kW13Rows = 2 * kIntermediate;
+constexpr int kW13Rows = 2 * kIntermediate; // 256
-constexpr int kThreads = 256;
-constexpr int kWarpSize = 32;
 
-__device__ inline float warp_sum(float value) {
+// =====================================================================
+// BI-V100 hardware constants
+// =====================================================================
+constexpr int kWarpSize = 64;   // BI-V100 warp width (was 32)
+constexpr int kThreads = 256;   // 4 warps of 64 (was 8 warps of 32)
+constexpr int kWarpsPerBlock = kThreads / kWarpSize; // 4
+
+// =====================================================================
+// Warp-level sum reduction for 64-wide warps
+// =====================================================================
+// 6 steps: 32, 16, 8, 4, 2, 1  (was 5 steps for warp=32)
+__device__ __forceinline__ float warp_sum(float value) {
 #pragma unroll
   for (int offset = kWarpSize / 2; offset > 0; offset /= 2) {
     value += __shfl_down_sync(0xffffffff, value, offset);
@@ -21,79 +57,125 @@ __device__ inline float warp_sum(float value) {
   return value;
 }
 
+// =====================================================================
+// W13 kernel: gate_up = input @ W13[expert_ids[slot]]^T
+// =====================================================================
+// Grid maps one warp per (slot, output_row) pair.
+// Each warp computes dot(input[1,H], W13[eid, row, :]) using half2 loads
+// and reduces via 64-wide warp_sum.
+//
+// Total warps needed: kTopK * kW13Rows = 8 * 256 = 2048
+// With kWarpsPerBlock=4: 2048/4 = 512 blocks
 __global__ void direct_w13_kernel(
-    const __half* input, const __half* w13, const int64_t* expert_ids,
+    const __half* __restrict__ input,       // (1, 2048)
-    __half* gate_up) {
+    const __half* __restrict__ w13,         // (256, 256, 2048)
-  const int warp =
+    const int64_t* __restrict__ expert_ids, // (8,)
-      (static_cast<int>(blockIdx.x) * blockDim.x + threadIdx.x) / kWarpSize;
+    __half* __restrict__ gate_up) {         // (8, 256)
-  const int lane = threadIdx.x & (kWarpSize - 1);
-  if (warp >= kTopK * kW13Rows) {
-    return;
-  }
 
-  const int slot = warp / kW13Rows;
+  // Map thread to (warp_id → slot, row) and lane within warp
-  const int local_row = warp - slot * kW13Rows;
+  const int global_warp =
+      static_cast<int>(blockIdx.x) * kWarpsPerBlock +
+      (threadIdx.x / kWarpSize);
+  const int lane = threadIdx.x & (kWarpSize - 1);  // 0..63
+
+  if (global_warp >= kTopK * kW13Rows)
+    return;
+
+  const int slot = global_warp / kW13Rows;
+  const int local_row = global_warp % kW13Rows;
   const int64_t expert = expert_ids[slot];
-  const int64_t weight_row =
+
+  // Weight row pointer: w13[expert][local_row][0..kHidden)
+  const int64_t weight_offset =
       (expert * kW13Rows + local_row) * static_cast<int64_t>(kHidden);
+
+  // Vectorized dot product using half2 loads
+  // Each lane processes kHidden/2 / kWarpSize iterations
   const __half2* input2 = reinterpret_cast<const __half2*>(input);
-  const __half2* weight2 =
+  const __half2* weight2 = reinterpret_cast<const __half2*>(w13 + weight_offset);
-      reinterpret_cast<const __half2*>(w13 + weight_row);
+
   float sum = 0.0f;
   for (int index = lane; index < kHidden / 2; index += kWarpSize) {
     const __half2 x = input2[index];
-    const __half2 weight = weight2[index];
+    const __half2 w = weight2[index];
-    sum = fmaf(__half2float(weight.x), __half2float(x.x), sum);
+    sum = fmaf(__half2float(w.x), __half2float(x.x), sum);
-    sum = fmaf(__half2float(weight.y), __half2float(x.y), sum);
+    sum = fmaf(__half2float(w.y), __half2float(x.y), sum);
   }
+
+  // 64-wide warp reduction
   sum = warp_sum(sum);
+
+  // Lane 0 writes the output
   if (lane == 0) {
-    gate_up[warp] = __float2half_rn(sum);
+    gate_up[global_warp] = __float2half_rn(sum);
   }
 }
 
+// =====================================================================
+// W2+reduce kernel: output = sum_k( weights[k] * activated @ W2[eid]^T )
+// =====================================================================
+// Grid maps one warp per output hidden dimension.
+// Each warp loops over kTopK experts, computes dot product, and
+// accumulates the weighted sum.
+//
+// Total warps needed: kHidden = 2048
+// With kWarpsPerBlock=4: 2048/4 = 512 blocks
 __global__ void direct_w2_reduce_kernel(
-    const __half* activated, const __half* w2, const int64_t* expert_ids,
+    const __half* __restrict__ activated,   // (8, 128)
-    const __half* weights, __half* output) {
+    const __half* __restrict__ w2,          // (256, 2048, 128)
-  const int warp =
+    const int64_t* __restrict__ expert_ids, // (8,)
-      (static_cast<int>(blockIdx.x) * blockDim.x + threadIdx.x) / kWarpSize;
+    const __half* __restrict__ weights,     // (8,)
+    __half* __restrict__ output) {          // (1, 2048)
+
+  const int global_warp =
+      static_cast<int>(blockIdx.x) * kWarpsPerBlock +
+      (threadIdx.x / kWarpSize);
   const int lane = threadIdx.x & (kWarpSize - 1);
-  if (warp >= kHidden) {
+
+  if (global_warp >= kHidden)
     return;
-  }
 
   float weighted_sum = 0.0f;
+
 #pragma unroll
   for (int slot = 0; slot < kTopK; ++slot) {
     const int64_t expert = expert_ids[slot];
-    const int64_t weight_row =
+
-        (expert * kHidden + warp) * static_cast<int64_t>(kIntermediate);
+    // Weight row: w2[expert][global_warp][0..kIntermediate)
+    const int64_t weight_offset =
+        (expert * kHidden + global_warp) * static_cast<int64_t>(kIntermediate);
+
     const __half2* activation2 = reinterpret_cast<const __half2*>(
         activated + slot * kIntermediate);
-    const __half2* weight2 =
+    const __half2* weight2 = reinterpret_cast<const __half2*>(
-        reinterpret_cast<const __half2*>(w2 + weight_row);
+        w2 + weight_offset);
+
     float expert_sum = 0.0f;
-    for (int index = lane; index < kIntermediate / 2;
+    for (int index = lane; index < kIntermediate / 2; index += kWarpSize) {
-         index += kWarpSize) {
       const __half2 x = activation2[index];
-      const __half2 weight = weight2[index];
+      const __half2 w = weight2[index];
-      expert_sum = fmaf(
+      expert_sum = fmaf(__half2float(w.x), __half2float(x.x), expert_sum);
-          __half2float(weight.x), __half2float(x.x), expert_sum);
+      expert_sum = fmaf(__half2float(w.y), __half2float(x.y), expert_sum);
-      expert_sum = fmaf(
-          __half2float(weight.y), __half2float(x.y), expert_sum);
     }
+
+    // 64-wide warp reduction
     expert_sum = warp_sum(expert_sum);
+
+    // Lane 0 accumulates weighted result
     if (lane == 0) {
-      const __half expert_half = __float2half_rn(expert_sum);
+      weighted_sum += __half2float(__hmul(
-      const __half product = __hmul(expert_half, weights[slot]);
+          __float2half_rn(expert_sum), weights[slot]));
-      weighted_sum += __half2float(product);
     }
   }
+
   if (lane == 0) {
-    output[warp] = __float2half_rn(weighted_sum);
+    output[global_warp] = __float2half_rn(weighted_sum);
   }
 }
 
+// =====================================================================
+// Input validation helpers
+// =====================================================================
 void check_half_cuda(const torch::Tensor& tensor, const char* name) {
   TORCH_CHECK(tensor.is_cuda(), name, " must be a CUDA tensor");
   TORCH_CHECK(tensor.scalar_type() == torch::kFloat16,
@@ -110,7 +192,11 @@ void check_ids(const torch::Tensor& expert_ids) {
               "expert_ids must have shape (8,)");
 }
 
-}  // namespace
+}  // anonymous namespace
+
+// =====================================================================
+// Python-facing functions
+// =====================================================================
 
 torch::Tensor direct_w13(const torch::Tensor& input,
                          const torch::Tensor& w13,
@@ -120,17 +206,18 @@ torch::Tensor direct_w13(const torch::Tensor& input,
   check_ids(expert_ids);
   TORCH_CHECK(input.dim() == 2 && input.size(0) == 1
                   && input.size(1) == kHidden,
-              "input must have shape (1, 2048)");
+              "input must have shape (1, ", kHidden, ")");
   TORCH_CHECK(w13.dim() == 3 && w13.size(0) == kExperts
                   && w13.size(1) == kW13Rows
                   && w13.size(2) == kHidden,
-              "w13 must have shape (256, 256, 2048)");
+              "w13 must have shape (", kExperts, ", ", kW13Rows, ", ", kHidden, ")");
 
   auto output = torch::empty({kTopK, kW13Rows}, input.options());
-  constexpr int kWarpsPerBlock = kThreads / kWarpSize;
+
-  constexpr int kBlocks =
+  constexpr int total_warps = kTopK * kW13Rows;  // 2048
-      (kTopK * kW13Rows + kWarpsPerBlock - 1) / kWarpsPerBlock;
+  constexpr int blocks = (total_warps + kWarpsPerBlock - 1) / kWarpsPerBlock;
-  direct_w13_kernel<<<kBlocks, kThreads, 0,
+
+  direct_w13_kernel<<<blocks, kThreads, 0,
                       at::cuda::getCurrentCUDAStream()>>>(
       reinterpret_cast<const __half*>(input.data_ptr<at::Half>()),
       reinterpret_cast<const __half*>(w13.data_ptr<at::Half>()),
@@ -150,19 +237,20 @@ torch::Tensor direct_w2_reduce(const torch::Tensor& activated,
   check_ids(expert_ids);
   TORCH_CHECK(activated.dim() == 2 && activated.size(0) == kTopK
                   && activated.size(1) == kIntermediate,
-              "activated must have shape (8, 128)");
+              "activated must have shape (", kTopK, ", ", kIntermediate, ")");
   TORCH_CHECK(w2.dim() == 3 && w2.size(0) == kExperts
                   && w2.size(1) == kHidden
                   && w2.size(2) == kIntermediate,
-              "w2 must have shape (256, 2048, 128)");
+              "w2 must have shape (", kExperts, ", ", kHidden, ", ", kIntermediate, ")");
   TORCH_CHECK(weights.dim() == 1 && weights.numel() == kTopK,
               "weights must have shape (8,)");
 
   auto output = torch::empty({1, kHidden}, activated.options());
-  constexpr int kWarpsPerBlock = kThreads / kWarpSize;
+
-  constexpr int kBlocks =
+  constexpr int total_warps = kHidden;  // 2048
-      (kHidden + kWarpsPerBlock - 1) / kWarpsPerBlock;
+  constexpr int blocks = (total_warps + kWarpsPerBlock - 1) / kWarpsPerBlock;
-  direct_w2_reduce_kernel<<<kBlocks, kThreads, 0,
+
+  direct_w2_reduce_kernel<<<blocks, kThreads, 0,
                             at::cuda::getCurrentCUDAStream()>>>(
       reinterpret_cast<const __half*>(activated.data_ptr<at::Half>()),
       reinterpret_cast<const __half*>(w2.data_ptr<at::Half>()),
@@ -175,7 +263,7 @@ torch::Tensor direct_w2_reduce(const torch::Tensor& activated,
 
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) {
   module.def("w13", &direct_w13,
-             "Direct selected-expert FP16 W13 matvec");
+             "Direct selected-expert FP16 W13 matvec (BI-V100, warp64)");
   module.def("w2_reduce", &direct_w2_reduce,
-             "Direct selected-expert W2 matvec and routed reduction");
+             "Direct selected-expert W2 matvec + routed reduction (BI-V100, warp64)");
 }

qwen3_6_scripts/patch_ops.sh

@@ -199,6 +199,43 @@ if [ -d "$PREBUILT_DIR" ]; then
     done
 fi
 
+# --- Rebuild corex_moe_direct_routed.so for BI-V100 warp_size=64 -----------
+# The prebuilt .so was compiled with kWarpSize=32 which silently corrupts
+# results on BI-V100 (64-wide warps).  Rebuild from the fixed .cu source
+# that uses kWarpSize=64 and 6-step shuffle reductions.
+build_stage "rebuilding corex_moe_direct_routed.so (warp64)"
+COREX_ROOT="${COREX_ROOT:-/usr/local/corex-3.2.3}"
+if [ ! -d "$COREX_ROOT" ]; then
+    COREX_ROOT="/usr/local/corex"
+fi
+TORCH_ROOT="${TORCH_ROOT:-$(python3 -c 'import torch,os;print(os.path.dirname(torch.__file__))' 2>/dev/null || echo "${COREX_ROOT}/lib64/python3/dist-packages/torch")}"
+DIRECT_ROUTED_SRC="./corex_moe_direct_routed.cu"
+DIRECT_ROUTED_DST="${VLLM_ROOT}/corex_moe_direct_routed.so"
+if [ -f "$DIRECT_ROUTED_SRC" ] && [ -x "${COREX_ROOT}/bin/clang++" ]; then
+    "${COREX_ROOT}/bin/clang++" \
+        -std=c++17 -O3 -shared -fPIC \
+        --cuda-path="${COREX_ROOT}" --cuda-gpu-arch=ivcore10 \
+        --no-cuda-version-check -D_GLIBCXX_USE_CXX11_ABI=0 \
+        -DTORCH_EXTENSION_NAME=corex_moe_direct_routed \
+        -DTORCH_API_INCLUDE_EXTENSION_H \
+        -I"${TORCH_ROOT}/include" \
+        -I"${TORCH_ROOT}/include/torch/csrc/api/include" \
+        -I"${TORCH_ROOT}/include/TH" -I"${TORCH_ROOT}/include/THC" \
+        -I/usr/local/include/python3.10 \
+        "$DIRECT_ROUTED_SRC" \
+        -L"${TORCH_ROOT}/lib" -L"${COREX_ROOT}/lib64" \
+        -Wl,-rpath,"${TORCH_ROOT}/lib" -Wl,-rpath,"${COREX_ROOT}/lib64" \
+        -ltorch_python -ltorch_cuda -ltorch_cpu -ltorch \
+        -lc10_cuda -lc10 -lcudart \
+        -o "$DIRECT_ROUTED_DST" 2>&1 && \
+    echo "[patch_ops] REBUILT corex_moe_direct_routed.so (warp64) → ${DIRECT_ROUTED_DST}" || \
+    echo "[patch_ops] WARNING: corex_moe_direct_routed.so rebuild FAILED, using prebuilt"
+elif [ ! -x "${COREX_ROOT}/bin/clang++" ]; then
+    echo "[patch_ops] WARNING: CoreX clang++ not found at ${COREX_ROOT}/bin/clang++, cannot rebuild direct_routed"
+else
+    echo "[patch_ops] WARNING: ${DIRECT_ROUTED_SRC} not found, cannot rebuild direct_routed"
+fi
+
 # --- Deploy ix_bridge Python integration layer --------------------------------
 build_stage "deploying ix_bridge operator replacements"
 EX_ENGINE_DIR="$(cd "$(dirname "$0")/ex_engine" 2>/dev/null && pwd || echo "")"
@@ -458,7 +495,4 @@ python3 ./verify_dlopen_chain.py --vllm-root "${VLLM_ROOT}" || {
     echo "[WARN] dlopen chain verification found issues (non-fatal)"
 }
 
-build_stage "patch script completed"
+build_stage "patch script completed"
-
-
-

qwen3_6_scripts/qwen3_5.py

@@ -1233,11 +1233,12 @@ class GatedDeltaNet(nn.Module):
             # (num_seqs, local_conv_dim, 1)
             mixed_qkv = (mixed_qkv_all
                          .to(weight_2d.dtype)
-                         .unsqueeze(-1))
+                         .unsqueeze(-1)
+                         .contiguous())
 
             if _USE_COREX_GDN_CAUSAL_CONV:
                 mixed_qkv_conv = _corex_gdn_causal_conv.causal_conv_update(
-                    conv_state, mixed_qkv, weight_2d)
+                    conv_state.contiguous(), mixed_qkv, weight_2d)
             else:
                 mixed_qkv_conv = _torch_causal_conv1d_update(
                     mixed_qkv, conv_state, weight_2d,
@@ -1784,6 +1785,9 @@ class Qwen3_5MoeSparseBlock(nn.Module):
             # Total: 3 kernel launches vs previous 16 (top_k*2).
             eids    = topk_ids[0]                              # (K,)
             ws      = topk_weights[0].to(hidden_states.dtype)  # (K,)
+            # --- corex_moe_direct_routed: zero-copy indexed GEMM (warp64) ---
+            # Shape must match the compiled kernel constants:
+            #   kHidden=2048, kExperts=256, kIntermediate=128, kTopK=8
             use_corex_direct = (
                 _USE_COREX_MOE_DIRECT_ROUTED
                 and hidden_states.dtype == torch.float16
@@ -1799,6 +1803,14 @@ class Qwen3_5MoeSparseBlock(nn.Module):
                 and w13.shape == (256, 256, 2048)
                 and w2.shape == (256, 2048, 128)
                 and eids.shape == (8,) and ws.shape == (8,))
+            if not hasattr(self, '_direct_routed_logged'):
+                self._direct_routed_logged = True
+                logger.info(
+                    "MoE T=1 direct_routed check: flag=%s match=%s "
+                    "hs=%s w13=%s w2=%s eids=%s ws=%s",
+                    _USE_COREX_MOE_DIRECT_ROUTED, use_corex_direct,
+                    tuple(hidden_states.shape), tuple(w13.shape),
+                    tuple(w2.shape), tuple(eids.shape), tuple(ws.shape))
             if use_corex_direct:
                 gate_up = _corex_moe_direct_routed.w13(
                     hidden_states, w13, eids)