fix(ci): resolve kernel failures, expand test coverage, and speed up compilation

.buildkite/pipeline.yml

@@ -1,12 +1,13 @@
 steps:
-  - label: "Humming Tests (B200)"
+  - &b200_tests
+    label: "Base Tests (B200)"
     key: humming-b200-ci
     agents:
       queue: b200-k8s
     timeout_in_minutes: 180
     concurrency: 1
     concurrency_group: humming/b200
-    command: bash .buildkite/run-tests.sh
+    command: HUMMING_TEST_TUNING_SOURCE=heuristic+batch_invariant bash .buildkite/run-tests.sh
     artifact_paths:
       - "test-results/pytest.xml"
     plugins:
@@ -15,7 +16,7 @@ steps:
             runtimeClassName: nvidia
             containers:
               # NGC 26.07 multi-arch: PyTorch 2.13 development snapshot and CUDA 13.3.
-              - image: nvcr.io/nvidia/pytorch:26.07-py3@sha256:2140e699b3beaf7f96a0081fd9c9406bc3832b435cdb60dfa2d261f7d2f34a1c
+              - image: &test_image nvcr.io/nvidia/pytorch:26.07-py3@sha256:2140e699b3beaf7f96a0081fd9c9406bc3832b435cdb60dfa2d261f7d2f34a1c
                 resources:
                   limits:
                     nvidia.com/gpu: 1
@@ -27,14 +28,15 @@ steps:
                 emptyDir:
                   medium: Memory
 
-  - label: "Humming Tests (DGX Spark)"
+  - &dgx_spark_tests
+    label: "Base Tests (DGX Spark)"
     key: humming-dgx-spark-ci
     agents:
       queue: dgx-spark
     timeout_in_minutes: 180
     concurrency: 1
     concurrency_group: humming/dgx-spark
-    command: bash .buildkite/run-tests.sh
+    command: HUMMING_TEST_TUNING_SOURCE=heuristic+batch_invariant bash .buildkite/run-tests.sh
     artifact_paths:
       - "test-results/pytest.xml"
     plugins:
@@ -47,3 +49,115 @@ steps:
           gpus: all
           volumes:
             - /dev/shm:/dev/shm
+
+  - &l4_tests
+    <<: *b200_tests
+    label: "Base Tests (L4)"
+    key: humming-l4-ci
+    agents:
+      queue: l4-k8s
+    concurrency_group: humming/l4
+    plugins:
+      # Match vLLM ci-infra's single-GPU L4 Kubernetes preset.
+      - kubernetes:
+          metadata:
+            annotations:
+              cluster-autoscaler.kubernetes.io/safe-to-evict: "false"
+              karpenter.sh/do-not-disrupt: "true"
+          podSpec:
+            priorityClassName: ci
+            serviceAccountName: buildkite-gpu-jobs
+            tolerations:
+              - key: nvidia.com/gpu
+                operator: Exists
+                effect: NoSchedule
+            nodeSelector:
+              vllm.ci/gpu-pool: l4x4
+            containers:
+              - image: *test_image
+                resources:
+                  limits:
+                    nvidia.com/gpu: 1
+                  requests:
+                    nvidia.com/gpu: 1
+                    cpu: "10"
+                    memory: 40Gi
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+                  sizeLimit: 24Gi
+
+  - &h200_tests
+    label: "Base Tests (H200)"
+    key: humming-h200-ci
+    agents:
+      queue: h200_18gb
+    timeout_in_minutes: 180
+    concurrency: 1
+    concurrency_group: humming/h200
+    env:
+      # The agent environment hook pulls the image before Docker plugin setup.
+      BUILDKITE_PLUGIN_DOCKER_IMAGE: *test_image
+    command: HUMMING_TEST_TUNING_SOURCE=heuristic+batch_invariant bash .buildkite/run-tests.sh
+    artifact_paths:
+      - "test-results/pytest.xml"
+    plugins:
+      # Use the GPU visibility assigned by the MIG agent, as in vLLM ci-infra.
+      - docker#v5.2.0:
+          image: *test_image
+          always-pull: true
+          propagate-environment: true
+          propagate-uid-gid: true
+          environment:
+            - CUDA_VISIBLE_DEVICES
+            - NVIDIA_VISIBLE_DEVICES
+            - PYTORCH_CUDA_ALLOC_CONF=expandable_segments:False
+          volumes:
+            - /dev/shm:/dev/shm
+            - /dev/nvidiactl:/dev/nvidiactl
+
+  - input: "Run Large Sampled Tests (B200)"
+    key: approve-b200-sampled-tests
+
+  - <<: *b200_tests
+    label: "Large Sampled Tests (B200)"
+    key: humming-b200-sampled-ci
+    concurrency_group: humming/b200-sampled
+    depends_on: approve-b200-sampled-tests
+    command: HUMMING_TEST_TUNING_SEED=20261005 HUMMING_TEST_TUNING_SOURCE=sampled32 bash .buildkite/run-tests.sh tests/kernels/humming/
+
+  - input: "Run Large Sampled Tests (DGX Spark)"
+    key: approve-dgx-spark-sampled-tests
+
+  - <<: *dgx_spark_tests
+    label: "Large Sampled Tests (DGX Spark)"
+    key: humming-dgx-spark-sampled-ci
+    concurrency_group: humming/dgx-spark-sampled
+    depends_on: approve-dgx-spark-sampled-tests
+    command: HUMMING_TEST_TUNING_SEED=20261005 HUMMING_TEST_TUNING_SOURCE=sampled32 bash .buildkite/run-tests.sh tests/kernels/humming/
+
+  - input: "Run Large Sampled Tests (L4)"
+    key: approve-l4-sampled-tests
+
+  - <<: *l4_tests
+    label: "Large Sampled Tests (L4)"
+    key: humming-l4-sampled-ci
+    concurrency_group: humming/l4-sampled
+    depends_on: approve-l4-sampled-tests
+    command: HUMMING_TEST_TUNING_SEED=20261005 HUMMING_TEST_TUNING_SOURCE=sampled32 bash .buildkite/run-tests.sh tests/kernels/humming/
+
+  - input: "Run Large Sampled Tests (H200)"
+    key: approve-h200-sampled-tests
+
+  - <<: *h200_tests
+    label: "Large Sampled Tests (H200)"
+    key: humming-h200-sampled-ci
+    agents:
+      queue: h200_35gb
+    concurrency_group: humming/h200-sampled
+    depends_on: approve-h200-sampled-tests
+    command: HUMMING_TEST_TUNING_SEED=20261005 HUMMING_TEST_TUNING_SOURCE=sampled32 bash .buildkite/run-tests.sh tests/kernels/humming/

.buildkite/run-tests.sh

@@ -3,6 +3,9 @@ set -euo pipefail
 
 cd "$(dirname "${BASH_SOURCE[0]}")/.."
 
+echo "--- CPU and compilation limits"
+python3 humming/utils/cpu.py
+
 echo "--- :nvidia: GPU Info"
 nvidia-smi
 
@@ -28,4 +31,7 @@ uv pip install --python .venv/bin/python --no-deps -e .
 
 echo "--- Run tests"
 mkdir -p test-results
-.venv/bin/python -m pytest tests -v --junitxml=test-results/pytest.xml
+if [[ $# -eq 0 ]]; then
+  set -- tests
+fi
+.venv/bin/python -m pytest "$@" -v --junitxml=test-results/pytest.xml

humming/config/config.py

@@ -96,29 +96,15 @@ def is_umma_supported(self):
         )
         has_fp8_output = self.c_dtype in (dtypes.float16, dtypes.bfloat16)
         has_fp8_activation = self.a_dtype in (dtypes.float8e4m3, dtypes.float8e5m2, dtypes.float8e3m4)
-        has_fp8_weights = self.b_dtype in (
-            dtypes.float8e4m3,
-            dtypes.float8e5m2,
-            dtypes.float8e3m4,
-            dtypes.float4e2m1,
-            dtypes.float6e3m2,
-            dtypes.float6e2m3,
-        )
-        has_fp8_operands = has_fp8_activation and (has_fp8_weights or self.b_dtype.is_integer_type)
         has_mx_input_scale = self.input_scale_group_size == 0 or (
             self.input_scale_group_size == 32 and self.as_dtype in (None, dtypes.float8e8m0)
         )
         has_mx_weight_scale = self.weight_scale_group_size == 0 or (
             self.weight_scale_group_size == 32 and self.bs_dtype == dtypes.float8e8m0
         )
-        has_mx_scales = (
-            has_mx_input_scale
-            and has_mx_weight_scale
-            and not self.has_zero_point
-            and not self.is_block_weight_scale
-        )
+        has_mx_scales = has_mx_input_scale and has_mx_weight_scale and not self.is_block_weight_scale
         has_supported_fp8_scales = has_fp8_epilogue_scales or has_mx_scales
-        use_fp8_umma = has_fp8_output and has_fp8_operands and has_supported_fp8_scales
+        use_fp8_umma = has_fp8_output and has_fp8_activation and has_supported_fp8_scales
         fp4_dtypes = (dtypes.float4e2m1, dtypes.float4e0m3)
         has_fp4_weights = self.b_dtype in fp4_dtypes or self.b_dtype.is_integer_type
         has_fp4_operands = self.a_dtype in fp4_dtypes and has_fp4_weights
@@ -134,12 +120,7 @@ def is_umma_supported(self):
             self.input_scale_group_size == self.weight_scale_group_size
             and self.as_dtype in (None, self.bs_dtype)
         )
-        has_fp4_scales = (
-            has_fp4_scale_format
-            and has_matching_group_scales
-            and not self.has_zero_point
-            and not self.is_block_weight_scale
-        )
+        has_fp4_scales = has_fp4_scale_format and has_matching_group_scales and not self.is_block_weight_scale
         use_fp4_umma = has_fp8_output and has_fp4_operands and has_fp4_scales
         use_f16_umma = self.a_dtype in (dtypes.float16, dtypes.bfloat16) and has_fp8_output
         has_int8_operands = self.a_dtype == dtypes.int8 and self.b_dtype.is_integer_type
@@ -609,6 +590,11 @@ def __post_init__(self):
         if self.use_cp_async is None:
             self.use_cp_async = current_device.sm_major >= 8
 
+        if self.use_warp_spec:
+            assert self.use_cp_async, "warp specialization requires cp.async"
+        if self.mma_type in (MmaType.WGMMA, MmaType.UMMA):
+            assert self.use_cp_async, f"{self.mma_type.value} requires cp.async"
+
         self.num_math_threads = math.prod(self.block_shape) // math.prod(self.warp_shape) * 32
         if self.use_warp_spec:
             self.num_load_threads = 128

humming/config/mma.py

@@ -89,34 +89,54 @@ def get_mxmma_compiler_error(layer_config, compiler_version):
 
 
 def get_register_budget_error(layer_config, tuning_config, use_f16_accum=False, registers_per_sm=None):
-    """Return a diagnostic when accumulators and one operand buffer exhaust the budget."""
     mma_type = tuning_config.mma_type or get_default_mma_type(layer_config)
-    if mma_type == MmaType.UMMA:
-        return None  # UMMA accumulators use TMEM and have separate resource checks.
     if registers_per_sm is None:
         registers_per_sm = current_device.max_registers_per_sm
     warp_m, warp_n, warp_k = tuning_config.warp_shape
     num_math_threads = math.prod(tuning_config.block_shape) // math.prod(tuning_config.warp_shape) * 32
     num_threads = num_math_threads + (128 if tuning_config.use_warp_spec else 0)
+    if mma_type == MmaType.UMMA:
+        if layer_config.use_raw_weight:
+            dequant_threads = 0
+        else:
+            dequant_threads = 128 * tuning_config.umma_num_dequant_warpgroups
+        num_threads = 256 + dequant_threads
     launch_budget = registers_per_sm // (num_threads * tuning_config.num_ctas_per_sm) // 8 * 8
-    accumulator_registers = warp_m * warp_n / (64 if use_f16_accum else 32)
-    has_group_accumulator = (
-        mma_type != MmaType.MXMMA
-        and layer_config.a_dtype.num_bits < 16
-        and (
-            layer_config.input_scale_group_size > 0
-            or (
-                not layer_config.use_fused_e8m0_scale
-                and (layer_config.is_group_weight_scale or layer_config.is_block_weight_scale)
-            )
+    if launch_budget < 56:
+        return (
+            "register budget exceeded: requires at least 56 registers per thread; "
+            f"launch budget {launch_budget:g}"
         )
-    )
+    if mma_type == MmaType.UMMA:
+        return None  # UMMA accumulators use TMEM and have separate resource checks.
+
+    if mma_type == MmaType.WGMMA:
+        # One instruction spans project N=64 across four warps. ptxas needs
+        # 26 extra registers for SS, or 30 for RS, beyond its accumulator tuple.
+        # setmaxnreg can raise the runtime math budget, but not this compile limit.
+        instruction_accumulators = warp_m // (4 if use_f16_accum else 2)
+        instruction_overhead = 26 if layer_config.use_raw_weight else 30
+        instruction_registers = instruction_accumulators + instruction_overhead
+        if instruction_registers > launch_budget:
+            return (
+                "register budget exceeded: WGMMA instruction requires at least "
+                f"{instruction_registers} registers per thread; launch budget {launch_budget:g}"
+            )
+
+    accumulator_registers = warp_m * warp_n / (64 if use_f16_accum else 32)
+    has_group_accumulator = False
+    if mma_type != MmaType.MXMMA and layer_config.a_dtype.num_bits < 16:
+        has_group_scale = layer_config.input_scale_group_size > 0
+        if not has_group_scale and not layer_config.use_fused_e8m0_scale:
+            has_group_scale = layer_config.is_group_weight_scale or layer_config.is_block_weight_scale
+        has_group_accumulator = has_group_scale
     math_budget = min(255, launch_budget)
     if tuning_config.use_warp_spec:
         # Match humming_ws.cuh's allocation, using the full physical accumulator
         # count for its preferred budget, independently of our spill estimate.
         physical_accumulators = accumulator_registers * (2 if has_group_accumulator else 1)
         preferred_budget = min(232, max(128, physical_accumulators * 2 + 96))
+        preferred_budget = math.ceil(preferred_budget / 8) * 8
         needs_more_load_registers = num_math_threads > 256 or (
             tuning_config.num_ctas_per_sm == 1 and layer_config.a_dtype.num_bits != 16
         )
@@ -126,6 +146,12 @@ def get_register_budget_error(layer_config, tuning_config, use_f16_accum=False,
             preferred_budget = 96
         math_budget = min(preferred_budget, max(24, available_registers // num_math_threads // 8 * 8))
 
+    if math_budget < 56:
+        return (
+            "register budget exceeded: requires at least 56 registers per math thread; "
+            f"launch budget {launch_budget:g}, math-thread budget {math_budget:g}"
+        )
+
     # Each ordinary MMA buffer spans K=256/activation_bits, hence M/4 and N/4
     # 32-bit registers per thread. Count dequantized operands, not both copies.
     buffer_registers = (warp_m + warp_n) / 4
@@ -137,11 +163,15 @@ def get_register_budget_error(layer_config, tuning_config, use_f16_accum=False,
             buffer_registers = warp_k * layer_config.a_dtype.num_bits / 64
     if has_group_accumulator:
         accumulator_registers *= 1.25
-    demand = accumulator_registers + buffer_registers
+    input_scale_registers = 0
+    if mma_type == MmaType.WGMMA and layer_config.input_scale_group_size > 0:
+        input_scale_registers = math.ceil(warp_m / 16) * 4
+    demand = accumulator_registers + buffer_registers + input_scale_registers
     if demand >= math_budget - 8:
         return (
             f"register budget exceeded: accumulator {accumulator_registers:g} + "
-            f"single-buffer {buffer_registers:g} = {demand:g} must be < "
+            f"single-buffer {buffer_registers:g} + input scales {input_scale_registers:g} "
+            f"= {demand:g} must be < "
             f"math-thread budget {math_budget:g} - 8 ({math_budget - 8:g})"
         )
     return None

humming/csrc/launcher/tensor.h

@@ -216,7 +216,7 @@ inline void check_tensor_bzp(std::optional<Tensor> &tensor, KernelData &kernel_d
     expected_dtype = dtype_id_to_tensor_dtype(kernel_data.c_dtype_id);
   } else {
     uint32_t scale_vec = 1;
-    if (kernel_data.mma_type == MmaType::MXMMA && get_dtype_num_bits(kernel_data.a_dtype_id) == 4 && group_size > 0) {
+    if (kernel_data.use_block_scaled_mma && get_dtype_num_bits(kernel_data.a_dtype_id) == 4 && group_size > 0) {
       scale_vec = 256 / 4 / group_size;
     }
     expected_shape.push_back(num_groups / scale_vec);
@@ -488,8 +488,7 @@ inline CUtensorMap make_tma_desc_bzp(std::optional<Tensor> &tensor_, KernelData
   }
 
   uint32_t num_bits = get_dtype_num_bits(kernel_data.b_dtype_id) <= 4 ? 4 : 8;
-  if (kernel_data.mma_type == MmaType::MXMMA &&
-      get_dtype_num_bits(kernel_data.a_dtype_id) == 4 && group_size > 0) {
+  if (kernel_data.use_block_scaled_mma && get_dtype_num_bits(kernel_data.a_dtype_id) == 4 && group_size > 0) {
     uint32_t scale_vec = 256 / 4 / group_size;
     if (scale_vec > 1) return make_tma_desc(tensor, {block_shape_n * num_bits * scale_vec / 32, num_groups / scale_vec}, 0, "bzp");
   }

humming/include/humming/arith/epilogue_arith.cuh

@@ -84,7 +84,6 @@ public:
     if constexpr (kUseNativeChannelScale) {
       // Apply channel scaling before narrowing: an FP8 dot product can exceed
       // FP16's range even when the scaled output is representable.
-      if (row_group == 0 && column_group == 0) may_process_on_smem_write(0, 0);
       const uint32_t *scales = kIsChannelWeightScale2 ? bs2 : (ElementBS::kBits == 8 ? dq_bs : bs);
       uint32_t channel = (threadIdx.x % 32) / 4;
       uint32_t packed_scale = __shfl_sync(0xffffffff, scales[column_group], channel / 2);