upd

Author: cyx-6

.github/workflows/nightly-release.yml

@@ -98,7 +98,7 @@ jobs:
         run: |
           python -m pip install --upgrade pip
           pip install build twine wheel
-          pip install setuptools>=61.0 requests filelock torch tqdm numpy apache-tvm-ffi==0.1.0b19
+          pip install setuptools>=61.0 requests filelock torch tqdm numpy "apache-tvm-ffi>=0.1,<0.2"
 
       - name: Build flashinfer-cubin wheel
         env:

.github/workflows/release.yml

@@ -136,7 +136,7 @@ jobs:
         run: |
           python -m pip install --upgrade pip
           pip install build twine wheel
-          pip install setuptools>=61.0 requests filelock torch tqdm numpy apache-tvm-ffi==0.1.0b19
+          pip install setuptools>=61.0 requests filelock torch tqdm numpy "apache-tvm-ffi>=0.1,<0.2"
 
       - name: Build flashinfer-cubin wheel
         run: |

csrc/batch_mla_config.jinja

@@ -13,7 +13,7 @@ using namespace flashinfer;
 #ifdef FLASHINFER_ENABLE_PROFILER
 #define ADDITIONAL_FUNC_PARAMS , Tensor profiler_buffer
 #define ADDITIONAL_PARAMS_SETTER \
-  params.profiler_buffer = static_cast<uint64_t*>(profiler_buffer->data);
+  params.profiler_buffer = static_cast<uint64_t*>(profiler_buffer.data_ptr());
 #else
 #define ADDITIONAL_FUNC_PARAMS
 #define ADDITIONAL_PARAMS_SETTER

csrc/nv_internal/tensorrt_llm/thop/fp4Op.cpp

@@ -149,7 +149,7 @@ void BlockScaleInterleave(TensorView blockScale, TensorView interleavedBlockScal
   }
   CHECK_CONTIGUOUS(blockScale);
   CHECK_INPUT_TYPE(blockScale, dl_uint8);
-  auto blockScaleShape = blockScale.shape();
+  auto blockScaleShape = blockScale.sizes();
   TVM_FFI_ICHECK(blockScaleShape.size() == 2 || blockScaleShape.size() == 3)
       << "Block Scale should be 2D or 3D tensor.";
   auto num_experts = blockScaleShape.size() == 3 ? blockScaleShape[0] : 1;
@@ -204,7 +204,7 @@ void BlockScaleInterleaveReverse(TensorView const& blockScale, TensorView revers
   }
   CHECK_CONTIGUOUS(blockScale);
   CHECK_INPUT_TYPE(blockScale, dl_uint8);
-  auto blockScaleShape = blockScale.shape();
+  auto blockScaleShape = blockScale.sizes();
   TVM_FFI_ICHECK(blockScaleShape.size() == 2 || blockScaleShape.size() == 3)
       << "Block Scale should be 2D or 3D tensor.";
   auto num_experts = blockScaleShape.size() == 3 ? blockScaleShape[0] : 1;
@@ -251,8 +251,8 @@ void E2M1AndUFP8SFScaleToFloatV2(TensorView valueE2M1, TensorView scaleFP8SF,
                                  bool isSfSwizzledLayout = true) {
   CHECK_CPU_INPUT(valueE2M1, dl_uint8);
   CHECK_CPU_INPUT(scaleFP8SF, dl_uint8);
-  auto packedShape = valueE2M1.shape();
-  auto scaleShape = scaleFP8SF.shape();
+  auto packedShape = valueE2M1.sizes();
+  auto scaleShape = scaleFP8SF.sizes();
   TVM_FFI_ICHECK_EQ(packedShape.size(), 2) << "valueE2M1 should be 2D tensor.";
   TVM_FFI_ICHECK_EQ(scaleShape.size(), 1) << "scaleFP8SF should be 1D tensor.";
 

csrc/nv_internal/tensorrt_llm/thop/fp4Quantize.cpp

@@ -50,7 +50,7 @@ void fp4_quantize(TensorView self, Optional<TensorView> const& globalScale, Tens
     globalScalePtr = static_cast<float*>(globalScale.value().data_ptr());
   }
 
-  auto const& inputShape = self.shape();
+  auto const& inputShape = self.sizes();
   auto const& rank = inputShape.size();
 
   TVM_FFI_ICHECK_GE(rank, 2) << "Input should be >=2D tensor.";
@@ -140,7 +140,7 @@ void fp4_batched_quantize(TensorView self, Optional<TensorView> const& mask, Ten
   CHECK_INPUT_TYPE(globalScale, fp32_dtype);
   TVM_FFI_ICHECK_EQ(sfVecSize, 16) << "sfVecSize can only be 16";
 
-  auto const& inputShape = self.shape();
+  auto const& inputShape = self.sizes();
   auto const& rank = inputShape.size();
 
   TVM_FFI_ICHECK_EQ(rank, 3) << "Input should be 3D tensor.";
@@ -205,7 +205,7 @@ void silu_and_mul_nvfp4_batched_quantize(TensorView const& self, TensorView cons
   CHECK_INPUT_TYPE(globalScale, fp32_dtype);
   TVM_FFI_ICHECK_EQ(sfVecSize, 16) << "sfVecSize can only be 16";
 
-  auto const& inputShape = self.shape();
+  auto const& inputShape = self.sizes();
   auto const& rank = inputShape.size();
   auto const& mask_rank = mask.ndim();
 

csrc/nv_internal/tensorrt_llm/thop/fp8Quantize.cpp

@@ -36,7 +36,7 @@ void mxfp8_quantize(TensorView input, TensorView valMxFP8, TensorView scaleFP8SF
   TVM_FFI_ICHECK_EQ(alignment % SF_VEC_SIZE, 0)
       << "alignment must be divisible by SF_VEC_SIZE = 32";
 
-  auto const& inputShape = input.shape();
+  auto const& inputShape = input.sizes();
   auto const& rank = inputShape.size();
 
   TVM_FFI_ICHECK_GE(rank, 2) << "Input should be >=2D tensor.";
@@ -98,7 +98,7 @@ void mxfp8_quantize_host(TensorView x_fp32, TensorView fp8_tensor, TensorView sc
   int32_t const sf_vec_size = 32;
   auto fp32_dtype = DLDataType{kDLFloat, 32, 1};
   CHECK_INPUT_TYPE(x_fp32, fp32_dtype);
-  auto data_shape = x_fp32.shape();
+  auto data_shape = x_fp32.sizes();
   TVM_FFI_ICHECK_EQ(data_shape.size(), 2) << "x_fp32 should be 2D tensor.";
   int num_tokens = data_shape[0];
   int hidden_dim = data_shape[1];
@@ -145,8 +145,8 @@ void mxfp8_dequantize_host(TensorView value_e4m3, TensorView scale_ue8m08sf,
   int32_t const sf_vec_size = 32;
   CHECK_INPUT_TYPE(value_e4m3, dl_uint8);
   CHECK_INPUT_TYPE(scale_ue8m08sf, dl_uint8);
-  auto data_shape = value_e4m3.shape();
-  auto scale_shape = scale_ue8m08sf.shape();
+  auto data_shape = value_e4m3.sizes();
+  auto scale_shape = scale_ue8m08sf.sizes();
   TVM_FFI_ICHECK_EQ(data_shape.size(), 2) << "value_e4m3 should be 2D tensor.";
   TVM_FFI_ICHECK_EQ(scale_shape.size(), 1) << "scale_ue8m08sf should be 1D tensor.";
 

csrc/rope.cu

@@ -285,29 +285,29 @@ void rope_quantize(TensorView q_rope_in, TensorView k_rope_in, TensorView q_nope
   CHECK_INPUT(pos_ids);
 
   // Extract dimensions from tensor shapes (flexible)
-  uint32_t rope_dim = q_rope_in->shape[q_rope_in->ndim - 1];
-  uint32_t no_rope_dim = q_nope_in->shape[q_nope_in->ndim - 1];
+  uint32_t rope_dim = q_rope_in.size(-1);
+  uint32_t no_rope_dim = q_nope_in.size(-1);
 
   // Validate rope and no_rope dimensions are consistent
-  TVM_FFI_ICHECK_EQ(k_rope_in->shape[k_rope_in->ndim - 1], rope_dim);
-  TVM_FFI_ICHECK_EQ(k_nope_in->shape[k_nope_in->ndim - 1], no_rope_dim);
-  TVM_FFI_ICHECK_EQ(q_rope_out->shape[q_rope_out->ndim - 1], rope_dim);
-  TVM_FFI_ICHECK_EQ(k_rope_out->shape[k_rope_out->ndim - 1], rope_dim);
-  TVM_FFI_ICHECK_EQ(q_nope_out->shape[q_nope_out->ndim - 1], no_rope_dim);
-  TVM_FFI_ICHECK_EQ(k_nope_out->shape[k_nope_out->ndim - 1], no_rope_dim);
-  TVM_FFI_ICHECK_EQ(q_rope_in->dtype, k_rope_in->dtype);
-  TVM_FFI_ICHECK_EQ(q_rope_in->dtype, q_nope_in->dtype);
-  TVM_FFI_ICHECK_EQ(q_rope_in->dtype, k_nope_in->dtype);
-  TVM_FFI_ICHECK_EQ(q_rope_out->dtype, k_rope_out->dtype);
-  TVM_FFI_ICHECK_EQ(q_rope_out->dtype, q_nope_out->dtype);
-  TVM_FFI_ICHECK_EQ(q_rope_out->dtype, k_nope_out->dtype);
+  TVM_FFI_ICHECK_EQ(k_rope_in.size(-1), rope_dim);
+  TVM_FFI_ICHECK_EQ(k_nope_in.size(-1), no_rope_dim);
+  TVM_FFI_ICHECK_EQ(q_rope_out.size(-1), rope_dim);
+  TVM_FFI_ICHECK_EQ(k_rope_out.size(-1), rope_dim);
+  TVM_FFI_ICHECK_EQ(q_nope_out.size(-1), no_rope_dim);
+  TVM_FFI_ICHECK_EQ(k_nope_out.size(-1), no_rope_dim);
+  TVM_FFI_ICHECK_EQ(q_rope_in.dtype(), k_rope_in.dtype());
+  TVM_FFI_ICHECK_EQ(q_rope_in.dtype(), q_nope_in.dtype());
+  TVM_FFI_ICHECK_EQ(q_rope_in.dtype(), k_nope_in.dtype());
+  TVM_FFI_ICHECK_EQ(q_rope_out.dtype(), k_rope_out.dtype());
+  TVM_FFI_ICHECK_EQ(q_rope_out.dtype(), q_nope_out.dtype());
+  TVM_FFI_ICHECK_EQ(q_rope_out.dtype(), k_nope_out.dtype());
 
   // Validate supported input data types (float16 or bfloat16)
-  TVM_FFI_ICHECK(q_rope_in->dtype == dl_float16 || q_rope_in->dtype == dl_bfloat16)
+  TVM_FFI_ICHECK(q_rope_in.dtype() == dl_float16 || q_rope_in.dtype() == dl_bfloat16)
       << "Input dtype must be float16 or bfloat16";
 
   // Validate supported output quantization data types (float8_e4m3fn or float8_e5m2)
-  TVM_FFI_ICHECK(q_rope_out->dtype == dl_float8_e4m3fn || q_rope_out->dtype == dl_float8_e5m2)
+  TVM_FFI_ICHECK(q_rope_out.dtype() == dl_float8_e4m3fn || q_rope_out.dtype() == dl_float8_e5m2)
       << "Output dtype must be float8_e4m3fn or float8_e5m2";
 
   // Q tensors are always 3D: (nnz, num_qo_heads, rope_dim/no_rope_dim)
@@ -318,7 +318,7 @@ void rope_quantize(TensorView q_rope_in, TensorView k_rope_in, TensorView q_nope
 
   // K tensors can be 2D (MLA) or 3D (GQA/MHA)
   uint32_t num_kv_heads;
-  if (k_rope_in->ndim == 2) {
+  if (k_rope_in.ndim() == 2) {
     // MLA case: k_rope_in: (nnz, rope_dim), k_nope_in: (nnz, no_rope_dim)
     CHECK_DIM(2, k_rope_in);
     CHECK_DIM(2, k_nope_in);
@@ -331,81 +331,82 @@ void rope_quantize(TensorView q_rope_in, TensorView k_rope_in, TensorView q_nope
     CHECK_DIM(3, k_nope_in);
     CHECK_DIM(3, k_rope_out);
     CHECK_DIM(3, k_nope_out);
-    num_kv_heads = k_rope_in->shape[1];
+    num_kv_heads = k_rope_in.size(1);
   }
-  uint32_t nnz = q_rope_in->shape[0];
-  uint32_t num_qo_heads = q_rope_in->shape[1];
+  uint32_t nnz = q_rope_in.size(0);
+  uint32_t num_qo_heads = q_rope_in.size(1);
 
   // Validate consistent dimensions across all tensors
-  TVM_FFI_ICHECK_EQ(q_nope_in->shape[0], nnz);
-  TVM_FFI_ICHECK_EQ(k_rope_in->shape[0], nnz);
-  TVM_FFI_ICHECK_EQ(k_nope_in->shape[0], nnz);
-  TVM_FFI_ICHECK_EQ(q_rope_out->shape[0], nnz);
-  TVM_FFI_ICHECK_EQ(k_rope_out->shape[0], nnz);
-  TVM_FFI_ICHECK_EQ(q_nope_out->shape[0], nnz);
-  TVM_FFI_ICHECK_EQ(k_nope_out->shape[0], nnz);
+  TVM_FFI_ICHECK_EQ(q_nope_in.size(0), nnz);
+  TVM_FFI_ICHECK_EQ(k_rope_in.size(0), nnz);
+  TVM_FFI_ICHECK_EQ(k_nope_in.size(0), nnz);
+  TVM_FFI_ICHECK_EQ(q_rope_out.size(0), nnz);
+  TVM_FFI_ICHECK_EQ(k_rope_out.size(0), nnz);
+  TVM_FFI_ICHECK_EQ(q_nope_out.size(0), nnz);
+  TVM_FFI_ICHECK_EQ(k_nope_out.size(0), nnz);
 
   // Validate Q tensor head dimensions are consistent
-  TVM_FFI_ICHECK_EQ(q_nope_in->shape[1], num_qo_heads);
-  TVM_FFI_ICHECK_EQ(q_rope_out->shape[1], num_qo_heads);
-  TVM_FFI_ICHECK_EQ(q_nope_out->shape[1], num_qo_heads);
+  TVM_FFI_ICHECK_EQ(q_nope_in.size(1), num_qo_heads);
+  TVM_FFI_ICHECK_EQ(q_rope_out.size(1), num_qo_heads);
+  TVM_FFI_ICHECK_EQ(q_nope_out.size(1), num_qo_heads);
 
   // Validate K tensor head dimensions (if 3D)
-  if (k_rope_in->ndim == 3) {
-    TVM_FFI_ICHECK_EQ(k_nope_in->shape[1], num_kv_heads);
-    TVM_FFI_ICHECK_EQ(k_rope_out->shape[1], num_kv_heads);
-    TVM_FFI_ICHECK_EQ(k_nope_out->shape[1], num_kv_heads);
+  if (k_rope_in.ndim() == 3) {
+    TVM_FFI_ICHECK_EQ(k_nope_in.size(1), num_kv_heads);
+    TVM_FFI_ICHECK_EQ(k_rope_out.size(1), num_kv_heads);
+    TVM_FFI_ICHECK_EQ(k_nope_out.size(1), num_kv_heads);
   }
 
-  const uint32_t q_rope_in_stride_n = q_rope_in->strides[0];
-  const uint32_t q_rope_in_stride_h = q_rope_in->strides[1];
-  const uint32_t q_nope_in_stride_n = q_nope_in->strides[0];
-  const uint32_t q_nope_in_stride_h = q_nope_in->strides[1];
-  const uint32_t q_rope_out_stride_n = q_rope_out->strides[0];
-  const uint32_t q_rope_out_stride_h = q_rope_out->strides[1];
-  const uint32_t q_nope_out_stride_n = q_nope_out->strides[0];
-  const uint32_t q_nope_out_stride_h = q_nope_out->strides[1];
+  const uint32_t q_rope_in_stride_n = q_rope_in.stride(0);
+  const uint32_t q_rope_in_stride_h = q_rope_in.stride(1);
+  const uint32_t q_nope_in_stride_n = q_nope_in.stride(0);
+  const uint32_t q_nope_in_stride_h = q_nope_in.stride(1);
+  const uint32_t q_rope_out_stride_n = q_rope_out.stride(0);
+  const uint32_t q_rope_out_stride_h = q_rope_out.stride(1);
+  const uint32_t q_nope_out_stride_n = q_nope_out.stride(0);
+  const uint32_t q_nope_out_stride_h = q_nope_out.stride(1);
 
   // K tensor strides depend on dimensionality
   uint32_t k_rope_in_stride, k_nope_in_stride, k_rope_out_stride, k_nope_out_stride;
   uint32_t k_rope_in_stride_h, k_nope_in_stride_h, k_rope_out_stride_h, k_nope_out_stride_h;
 
-  if (k_rope_in->ndim == 2) {
+  if (k_rope_in.ndim() == 2) {
     // 2D K tensors (MLA): only have batch stride
-    k_rope_in_stride = k_rope_in->strides[0];
-    k_nope_in_stride = k_nope_in->strides[0];
-    k_rope_out_stride = k_rope_out->strides[0];
-    k_nope_out_stride = k_nope_out->strides[0];
+    k_rope_in_stride = k_rope_in.stride(0);
+    k_nope_in_stride = k_nope_in.stride(0);
+    k_rope_out_stride = k_rope_out.stride(0);
+    k_nope_out_stride = k_nope_out.stride(0);
     // For 2D tensors, head stride is the same as batch stride (shared K/V)
     k_rope_in_stride_h = k_rope_in_stride;
     k_nope_in_stride_h = k_nope_in_stride;
     k_rope_out_stride_h = k_rope_out_stride;
     k_nope_out_stride_h = k_nope_out_stride;
   } else {
     // 3D K tensors (GQA/MHA): have both batch and head strides
-    k_rope_in_stride = k_rope_in->strides[0];
-    k_rope_in_stride_h = k_rope_in->strides[1];
-    k_nope_in_stride = k_nope_in->strides[0];
-    k_nope_in_stride_h = k_nope_in->strides[1];
-    k_rope_out_stride = k_rope_out->strides[0];
-    k_rope_out_stride_h = k_rope_out->strides[1];
-    k_nope_out_stride = k_nope_out->strides[0];
-    k_nope_out_stride_h = k_nope_out->strides[1];
+    k_rope_in_stride = k_rope_in.stride(0);
+    k_rope_in_stride_h = k_rope_in.stride(1);
+    k_nope_in_stride = k_nope_in.stride(0);
+    k_nope_in_stride_h = k_nope_in.stride(1);
+    k_rope_out_stride = k_rope_out.stride(0);
+    k_rope_out_stride_h = k_rope_out.stride(1);
+    k_nope_out_stride = k_nope_out.stride(0);
+    k_nope_out_stride_h = k_nope_out.stride(1);
   }
 
-  cudaSetDevice(q_rope_in->device.device_id);
-  const cudaStream_t stream = get_stream(q_rope_in->device);
-  DISPATCH_DLPACK_DTYPE_TO_CTYPE_FP16(q_rope_in->dtype, c_type, [&] {
-    return DISPATCH_DLPACK_DTYPE_TO_CTYPE_FP8(q_rope_out->dtype, c_quant_type, [&] {
-      return DISPATCH_DLPACK_IDTYPE_TO_CTYPE(pos_ids->dtype, c_idtype, [&] {
+  cudaSetDevice(q_rope_in.device().device_id);
+  const cudaStream_t stream = get_stream(q_rope_in.device());
+  DISPATCH_DLPACK_DTYPE_TO_CTYPE_FP16(q_rope_in.dtype(), c_type, [&] {
+    return DISPATCH_DLPACK_DTYPE_TO_CTYPE_FP8(q_rope_out.dtype(), c_quant_type, [&] {
+      return DISPATCH_DLPACK_IDTYPE_TO_CTYPE(pos_ids.dtype(), c_idtype, [&] {
         cudaError_t status = RopeQuantize(
-            static_cast<c_type*>(q_rope_in->data), static_cast<c_type*>(k_rope_in->data),
-            static_cast<c_type*>(q_nope_in->data), static_cast<c_type*>(k_nope_in->data),
-            static_cast<c_quant_type*>(q_rope_out->data),
-            static_cast<c_quant_type*>(k_rope_out->data),
-            static_cast<c_quant_type*>(q_nope_out->data),
-            static_cast<c_quant_type*>(k_nope_out->data), static_cast<float*>(cos_sin_cache->data),
-            static_cast<c_idtype*>(pos_ids->data), nnz, num_qo_heads, num_kv_heads, rope_dim,
+            static_cast<c_type*>(q_rope_in.data_ptr()), static_cast<c_type*>(k_rope_in.data_ptr()),
+            static_cast<c_type*>(q_nope_in.data_ptr()), static_cast<c_type*>(k_nope_in.data_ptr()),
+            static_cast<c_quant_type*>(q_rope_out.data_ptr()),
+            static_cast<c_quant_type*>(k_rope_out.data_ptr()),
+            static_cast<c_quant_type*>(q_nope_out.data_ptr()),
+            static_cast<c_quant_type*>(k_nope_out.data_ptr()),
+            static_cast<float*>(cos_sin_cache.data_ptr()),
+            static_cast<c_idtype*>(pos_ids.data_ptr()), nnz, num_qo_heads, num_kv_heads, rope_dim,
             no_rope_dim, q_rope_in_stride_n, q_rope_in_stride_h, q_nope_in_stride_n,
             q_nope_in_stride_h, q_rope_out_stride_n, q_rope_out_stride_h, q_nope_out_stride_n,
             q_nope_out_stride_h, k_rope_in_stride, k_rope_in_stride_h, k_nope_in_stride,