[None][fix] fix FA4 install in devel docker (#14706)

Signed-off-by: Olivia Stoner <245287810+o-stoner@users.noreply.github.com>

Author: o-stoner

docker/Dockerfile.multi

@@ -71,6 +71,7 @@ RUN --mount=type=bind,source=docker/common,target=/opt/docker/common \
     GITHUB_MIRROR=${GITHUB_MIRROR} bash /opt/docker/common/install_ucx.sh && \
     GITHUB_MIRROR=${GITHUB_MIRROR} bash /opt/docker/common/install_nixl.sh && \
     bash /opt/docker/common/install_etcd.sh && \
+    GITHUB_MIRROR=${GITHUB_MIRROR} bash /opt/docker/common/install_fa4.sh && \
     rm -rf /root/.cache/uv/archive-v0 && \
     # WAR against https://github.com/advisories/GHSA-58pv-8j8x-9vj2
     rm -rf /usr/local/lib/python3.12/dist-packages/setuptools/_vendor/jaraco.context-5.3.0.dist-info && \

docker/common/install_fa4.sh

@@ -0,0 +1,10 @@
+#!/bin/bash
+
+set -ex
+
+FLASH_ATTN_4_VERSION="4.0.0b11"
+
+if [ -n "${GITHUB_MIRROR}" ]; then
+  export PIP_INDEX_URL="https://urm.nvidia.com/artifactory/api/pypi/pypi-remote/simple"
+fi
+pip3 install "flash-attn-4==${FLASH_ATTN_4_VERSION}"

jenkins/current_image_tags.properties

@@ -13,8 +13,8 @@
 #     images are adopted from PostMerge pipelines, the abbreviated commit hash is used instead.
 IMAGE_NAME=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm
 
-LLM_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:pytorch-26.02-py3-x86_64-ubuntu24.04-trt10.15.1.29-skip-tritondevel-202606051544-14972
-LLM_SBSA_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:pytorch-26.02-py3-sbsa-ubuntu24.04-trt10.15.1.29-skip-tritondevel-202606051544-14972
-LLM_ROCKYLINUX8_PY310_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:cuda-13.1.0-devel-rocky8-x86_64-rocky8-py310-trt10.15.1.29-skip-tritondevel-202606051544-14972
-LLM_ROCKYLINUX8_PY312_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:cuda-13.1.0-devel-rocky8-x86_64-rocky8-py312-trt10.15.1.29-skip-tritondevel-202606051544-14972
-LLM_SBSA_WHEEL_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:cuda-13.1.0-devel-ubuntu24.04-sbsa-ubuntu24.04-py312-trt10.15.1.29-skip-tritondevel-202606051544-14972
+LLM_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:pytorch-26.02-py3-x86_64-ubuntu24.04-trt10.15.1.29-skip-tritondevel-202606091844-14706
+LLM_SBSA_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:pytorch-26.02-py3-sbsa-ubuntu24.04-trt10.15.1.29-skip-tritondevel-202606091844-14706
+LLM_ROCKYLINUX8_PY310_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:cuda-13.1.0-devel-rocky8-x86_64-rocky8-py310-trt10.15.1.29-skip-tritondevel-202606091844-14706
+LLM_ROCKYLINUX8_PY312_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:cuda-13.1.0-devel-rocky8-x86_64-rocky8-py312-trt10.15.1.29-skip-tritondevel-202606091844-14706
+LLM_SBSA_WHEEL_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:cuda-13.1.0-devel-ubuntu24.04-sbsa-ubuntu24.04-py312-trt10.15.1.29-skip-tritondevel-202606091844-14706

tests/unittest/_torch/sampler/test_torch_sampler.py

@@ -2173,11 +2173,11 @@ def _validate_token_frequencies(
             test_expected_counts != 0, 0, test_token_counts
         )
         assert (test_token_counts_for_zero_prob == 0).all()
-        test_expected_counts_ma = np.ma.masked_array(
+        test_expected_counts_ma = np.ma.MaskedArray(
             test_expected_counts.numpy(),
             mask=(test_expected_counts.numpy() == 0),
         )
-        test_token_counts_ma = np.ma.masked_array(
+        test_token_counts_ma = np.ma.MaskedArray(
             test_token_counts.numpy(),
             mask=test_expected_counts_ma.mask,
         )
@@ -2227,7 +2227,7 @@ def _validate_token_frequencies(
             prob_delta = np.where(prob_delta > 5e-2, prob_delta, 0)  # NB: this is rather liberal
             # bound relative differences on remaining probs
             prob_delta_rel = (
-                np.ma.masked_array(num_samples * prob_delta, mask=test_expected_counts_ma.mask)
+                np.ma.MaskedArray(num_samples * prob_delta, mask=test_expected_counts_ma.mask)
                 / test_expected_counts_ma.data
             )
             assert prob_delta_rel.max() < 0.05