add kvbm flag

Author: nv-anants

container/Dockerfile.trtllm

@@ -4,6 +4,7 @@
 ARG BASE_IMAGE="nvcr.io/nvidia/pytorch"
 ARG BASE_IMAGE_TAG="25.06-py3"
 ARG RELEASE_BUILD
+ARG ENABLE_KVBM=false
 ARG RUNTIME_IMAGE="nvcr.io/nvidia/cuda"
 ARG RUNTIME_IMAGE_TAG="12.9.1-runtime-ubuntu24.04"
 
@@ -240,6 +241,8 @@ ARG ARCH_ALT
 
 FROM quay.io/pypa/manylinux_2_28_${ARCH_ALT} AS wheel_builder
 ARG RELEASE_BUILD
+# Use arg ENABLE_KVBM = true to turn on the block-manager feature
+ARG ENABLE_KVBM
 ARG CARGO_BUILD_JOBS
 # Set CARGO_BUILD_JOBS to 16 if not provided
 # This is to prevent cargo from building $(nproc) jobs in parallel,
@@ -292,7 +295,11 @@ RUN cargo build \
 # Build dynamo wheels
 RUN uv build --wheel --out-dir /workspace/dist && \
     cd /workspace/lib/bindings/python && \
-    uv build --wheel --out-dir /workspace/dist --python 3.12 && \
+    if [ "$ENABLE_KVBM" = "true" ]; then \
+        maturin build --release --features block-manager --out /workspace/dist; \
+    else \
+        maturin build --release --out /workspace/dist; \
+    fi && \
     if [ "$RELEASE_BUILD" = "true" ]; then \
         uv build --wheel --out-dir /workspace/dist --python 3.11 && \
         uv build --wheel --out-dir /workspace/dist --python 3.10; \