review cleanups: pp+gpt-oss guard, sparse GEMV asserts, warnings

- --pp with gpt-oss now fails with a clear message instead of a
  cryptic missing-weight panic inside the Qwen3-only PP engine.
- Sparse GEMV wrappers assert K%16==0 (FP8) / K%32==0 (MXFP4) — the
  uint4-vectorized kernels would silently drop a tail otherwise.
- Document the topk_ids buffer holding i32 under an F32 dtype label
  (DType has no I32).
- Drop unused imports/locals and the cuBLASLt scale-mode constants
  orphaned by the strided-batched FP8 rework (e631a71).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
2026-06-12 17:02:59 +08:00
parent 1897b2e17a
commit 5343391dbd
8 changed files with 21 additions and 17 deletions

View File

@@ -1,5 +1,4 @@
use std::ffi::c_void;
use xserv_cuda::GpuBuffer;
use xserv_tensor::{DType, Device, Tensor};
unsafe extern "C" {

View File

@@ -1,5 +1,5 @@
use std::ffi::c_void;
use xserv_tensor::{DType, Device, Tensor};
use xserv_tensor::{DType, Tensor};
use crate::gemm::{cublas_handle, CublasHandle};
@@ -88,6 +88,9 @@ pub fn moe_topk_softmax(
let num_tokens = router_logits.shape()[0];
assert_eq!(router_logits.shape()[1], num_experts);
// NOTE: topk_ids actually holds i32 expert indices; DType has no I32, so
// this is a raw 4-byte buffer mislabeled F32. Never read it as floats —
// all consumers (weighted-sum / sparse GEMV kernels) cast to int*.
let topk_ids = Tensor::empty(&[num_tokens, top_k], DType::F32, router_logits.device());
let topk_weights = Tensor::empty(&[num_tokens, top_k], DType::F32, router_logits.device());
@@ -201,8 +204,12 @@ pub fn moe_sparse_gemv_fp8(
) -> Tensor {
assert_eq!(x.dtype(), DType::BF16);
assert!(x.is_contiguous());
assert_eq!(w_fp8_t.dtype(), DType::FP8E4M3);
let n = w_fp8_t.shape()[1];
let k = w_fp8_t.shape()[2];
// The kernel reads weights as uint4 (16 FP8 values per lane) and would
// silently skip a K%16 tail.
assert_eq!(k % 16, 0, "sparse FP8 GEMV requires K % 16 == 0, got {k}");
assert_eq!(x.shape()[x.ndim() - 1], k);
assert_eq!(x.shape()[0], if x_per_slot { num_tokens * top_k } else { num_tokens });
@@ -233,6 +240,8 @@ pub fn moe_sparse_gemv_mxfp4(
) -> Tensor {
assert_eq!(x.dtype(), DType::BF16);
assert!(x.is_contiguous());
// 32-element MXFP4 blocks, read as uint4 (32 nibbles) per lane.
assert_eq!(k % 32, 0, "sparse MXFP4 GEMV requires K % 32 == 0, got {k}");
assert_eq!(x.shape()[x.ndim() - 1], k);
assert_eq!(x.shape()[0], if x_per_slot { num_tokens * top_k } else { num_tokens });

View File

@@ -107,17 +107,11 @@ const CUDA_R_8F_E4M3: i32 = 28;
// MatmulDesc attributes
const CUBLASLT_MATMUL_DESC_A_SCALE_POINTER: i32 = 17;
const CUBLASLT_MATMUL_DESC_B_SCALE_POINTER: i32 = 18;
const CUBLASLT_MATMUL_DESC_A_SCALE_MODE: i32 = 31;
const CUBLASLT_MATMUL_DESC_B_SCALE_MODE: i32 = 32;
// MatrixLayout attributes
const CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT: i32 = 5;
const CUBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET: i32 = 6;
// Scale modes
const CUBLASLT_MATMUL_MATRIX_SCALE_SCALAR: i32 = 0;
const CUBLASLT_MATMUL_MATRIX_SCALE_OUTER_VEC_32F: i32 = 3;
// MatmulPreference attributes
const CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES: i32 = 1;