diff --git a/Dockerfile.spark b/Dockerfile.spark
index 4a101c1..9345e90 100644
--- a/Dockerfile.spark
+++ b/Dockerfile.spark
@@ -9,8 +9,10 @@ COPY wheels/sageattn3-*.whl /tmp/wheels/
 RUN python -m pip install --no-cache-dir --no-deps /tmp/wheels/sageattn3-*.whl \
     && rm -rf /tmp/wheels
 
+RUN python -m pip install --no-cache-dir --no-deps comfy-kitchen==0.2.28
+
 RUN python -m pip install --no-cache-dir --no-deps -e . \
-    && python -c "import comfy_kitchen, torch; from sageattn3 import sageattn3_blackwell; print(torch.__version__, torch.version.cuda)"
+    && python -c "import comfy_kitchen, torch; from sageattn3 import sageattn3_blackwell; assert hasattr(torch.ops.comfy_kitchen, 'rms_rope_split_half_'); print(torch.__version__, torch.version.cuda)"
 
 ENV H3_MODEL_PATH=/models/minimax_h3_ref2va_pruned_nvfp4.safetensors
 ENV TORCH_COMPILE_DISABLE=0 TORCHDYNAMO_DISABLE=0