diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 8f463a26..b2b54ebc 100644 --- a/.buildkite/pipeline.yml +++ b/.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/ diff --git a/.buildkite/run-tests.sh b/.buildkite/run-tests.sh index b8d6898f..d509ae6d 100644 --- a/.buildkite/run-tests.sh +++ b/.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 diff --git a/humming/config/config.py b/humming/config/config.py index b8302a0a..c512ce41 100644 --- a/humming/config/config.py +++ b/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 diff --git a/humming/config/mma.py b/humming/config/mma.py index c8af71f7..c0c489d7 100644 --- a/humming/config/mma.py +++ b/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 diff --git a/humming/csrc/launcher/tensor.h b/humming/csrc/launcher/tensor.h index c86d6b7a..bbbd76ac 100644 --- a/humming/csrc/launcher/tensor.h +++ b/humming/csrc/launcher/tensor.h @@ -216,7 +216,7 @@ inline void check_tensor_bzp(std::optional &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_, 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"); } diff --git a/humming/include/humming/arith/epilogue_arith.cuh b/humming/include/humming/arith/epilogue_arith.cuh index ef2350ae..81d77e05 100644 --- a/humming/include/humming/arith/epilogue_arith.cuh +++ b/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); diff --git a/humming/include/humming/arith/mainloop_arith.cuh b/humming/include/humming/arith/mainloop_arith.cuh index dd28c22b..7ec8a297 100644 --- a/humming/include/humming/arith/mainloop_arith.cuh +++ b/humming/include/humming/arith/mainloop_arith.cuh @@ -47,7 +47,7 @@ private: static constexpr uint32_t kInputScaleGroupSize = kIsGroupInputScale ? Ctx::kInputScaleGroupSize : 1; static constexpr uint32_t kWeightScaleGroupSize = kIsGroupOrBlockWeightScale ? Ctx::kWeightScaleGroupSize : 1; - static constexpr uint32_t kNumZPGroupsPerMma = kUseMxmma && ElementA::kBits == 4 && kIsGroupWeightScale ? kPartMmaShapeK / kWeightScaleGroupSize : 1; + static constexpr uint32_t kNumZPGroupsPerMma = Ctx::kUseBlockScaledMma && ElementA::kBits == 4 && kIsGroupWeightScale ? kPartMmaShapeK / kWeightScaleGroupSize : 1; static constexpr bool kUsePackedKLayout = Ctx::kUsePackedKLayout; static constexpr uint32_t kPackedKFactor = Ctx::kPackedKFactor; @@ -247,7 +247,7 @@ public: is_last_iter = (k + 1 == kNumKSlabs); } else { k_index = iter_id * kPartMmaShapeK + (k + 1) * MmaShape::K; - is_last_iter = iter_id == (Ctx::kWarpIters - 1); + is_last_iter = k_index == WarpShape::K; } uint32_t is_as_group_end = kIsGroupInputScale && k_index % kInputScaleGroupSize == 0; constexpr bool kProcessGroupWeightScale = kIsGroupWeightScale && !kUseFusedE8m0Scale; @@ -347,9 +347,11 @@ public: constexpr uint32_t kFullPackets = kNumBSPerGroup / 4; constexpr uint32_t kTailScales = kNumBSPerGroup % 4; - PRAGMA_UNROLL - for (uint32_t i = 0; i < kFullPackets; i++) - dq_bs_vals[i] = F8Conversion::num42float4(bs_vals[i]); + if constexpr (kFullPackets > 0) { + PRAGMA_UNROLL + for (uint32_t i = 0; i < kFullPackets; i++) + dq_bs_vals[i] = F8Conversion::num42float4(bs_vals[i]); + } if constexpr (kTailScales != 0) { uint32_t packed = 0; @@ -388,7 +390,7 @@ public: uint32_t buffer_id = iter_id % 2; uint32_t k_index = iter_id * kPartMmaShapeK + (k + 1) * MmaShape::K; - uint32_t is_last_iter = iter_id == (Ctx::kWarpIters - 1); + uint32_t is_last_iter = k_index == WarpShape::K; uint32_t is_as_group_end = kApplyGroupInputScaleOnC && k_index % kInputScaleGroupSize == 0; uint32_t is_bs_group_end = (kApplyGroupWeightScaleOnC || kApplyBlockWeightScaleOnC) && k_index % kWeightScaleGroupSize == 0; @@ -517,7 +519,7 @@ public: is_last_iter = (k + 1 == kNumKSlabs); } else { k_index = iter_id * kPartMmaShapeK + (k + 1) * MmaShape::K; - is_last_iter = iter_id == (Ctx::kWarpIters - 1); + is_last_iter = k_index == WarpShape::K; } uint32_t is_as_group_end = kApplyGroupInputScaleOnC && k_index % kInputScaleGroupSize == 0; uint32_t is_bs_group_end = (kApplyGroupWeightScaleOnC || kApplyBlockWeightScaleOnC) && k_index % kWeightScaleGroupSize == 0; @@ -635,6 +637,6 @@ public: template CUDA_INLINE T *regs_zp_as_ptr(uint32_t buffer_id) { - return reinterpret_cast(zp[buffer_id]); + return reinterpret_cast(zp[kIsChannelWeightScale ? 0 : buffer_id]); }; }; diff --git a/humming/include/humming/epilogue/smem_writer.cuh b/humming/include/humming/epilogue/smem_writer.cuh index 05a23e8e..aafb142b 100644 --- a/humming/include/humming/epilogue/smem_writer.cuh +++ b/humming/include/humming/epilogue/smem_writer.cuh @@ -220,6 +220,9 @@ public: template CUDA_INLINE void write_umma(MMA &mma, uint32_t slice_id, uint32_t slice_count, WriteChunk write_chunk) { + // Overlapped accumulators can emit other rows before row zero. Prepare + // channel scales before either output order reads them. + arith.may_process_on_smem_write(0, 0); // Specialize both output orders so scale-array indices stay compile-time constants. if constexpr (MMA::kAccumulatorStride != 0 && MMA::kAccumulatorStride < WarpShape::M) { if (mma.output_chunk_index(0) != 0) write_umma_order(mma, slice_id, slice_count, write_chunk); diff --git a/humming/include/humming/kernel/humming_umma.cuh b/humming/include/humming/kernel/humming_umma.cuh index 92b23f5e..30ec982b 100644 --- a/humming/include/humming/kernel/humming_umma.cuh +++ b/humming/include/humming/kernel/humming_umma.cuh @@ -394,6 +394,8 @@ __global__ __launch_bounds__(TuningConfig::kNumThreads, TuningConfig::kNumCtasPe s2r.load_channel(scheduler.slice_id); if constexpr (Consumer::kHasChannelData && kSeparateOutputStorage) { if (group + 1 == kOutputGroups) { + // Order channel reads before the producer's next async overwrite. + tma_fence_async_shared(); ctx.sync_math_threads(); consumer.arrive(kNumStages); } diff --git a/humming/include/humming/kernel/humming_ws.cuh b/humming/include/humming/kernel/humming_ws.cuh index e7a0b62d..c9cb2663 100644 --- a/humming/include/humming/kernel/humming_ws.cuh +++ b/humming/include/humming/kernel/humming_ws.cuh @@ -31,7 +31,8 @@ struct WarpSpecializationRegisterAllocation { static constexpr uint32_t kLoadThreadRegisters = TuningConfig::kNumMathThreads > 256 || (TuningConfig::kNumCtasPerSm == 1 && ElementA::kBits != 16) ? 40 : 24; static constexpr uint32_t kAccumulatorRegistersPerThread = sizeof(typename MMA::CRegistersArrayType) / sizeof(uint32_t) * (MMA::final_regs_c_index() + 1); static constexpr uint32_t kEstimatedMathThreadRegisters = MIN(232, MAX(128, kAccumulatorRegistersPerThread * 2 + 96)); - static constexpr uint32_t kPreferredMathThreadRegisters = TuningConfig::kNumMathThreads > 256 ? 96 : kEstimatedMathThreadRegisters; + // setmaxnreg requires multiples of eight; the physical budget below is rounded down. + static constexpr uint32_t kPreferredMathThreadRegisters = TuningConfig::kNumMathThreads > 256 ? 96 : CEIL_DIV(kEstimatedMathThreadRegisters, 8) * 8; static constexpr uint32_t kNumWarps = TuningConfig::kNumThreads / 32; static constexpr uint32_t kRegisterAllocationGranularityPerWarp = 256; static constexpr uint32_t kRegisterBudgetPerWarp = diff --git a/humming/include/humming/memory/g2s_loader/loader_bzp.cuh b/humming/include/humming/memory/g2s_loader/loader_bzp.cuh index 61940d1e..116dd218 100644 --- a/humming/include/humming/memory/g2s_loader/loader_bzp.cuh +++ b/humming/include/humming/memory/g2s_loader/loader_bzp.cuh @@ -21,11 +21,10 @@ private: static constexpr bool kIsChannel = Ctx::kIsChannelWeightScale; static constexpr bool kIsGroup = Ctx::kIsGroupWeightScale; static constexpr bool kEvictWeightsFirst = Ctx::kUseUmma && Ctx::kIsDenseGemm && kIsGroup && Ctx::kRasterGroupM > 1; - static constexpr bool kUseMxmma = Ctx::kUseMxmma; static constexpr uint32_t kGroupSize = kIsGroup ? Ctx::kWeightScaleGroupSize : ProblemShape::K; static constexpr uint32_t kPartMmaShapeK = Ctx::kPartMmaShapeK; - static constexpr uint32_t kNumGroupsPerMma = kUseMxmma && ElementA::kBits == 4 && kIsGroup ? kPartMmaShapeK / kGroupSize : 1; + static constexpr uint32_t kNumGroupsPerMma = Ctx::kUseBlockScaledMma && ElementA::kBits == 4 && kIsGroup ? kPartMmaShapeK / kGroupSize : 1; static constexpr uint32_t kNumZPBits = kIsFpZeroPoint ? 16 : MAX(4, static_next_power_of_2(ElementB::kBits)); static constexpr uint32_t kSmemStride = BlockShape::N * kNumZPBits / 32 / 4 * kNumGroupsPerMma; static constexpr uint32_t kGmemStride = ProblemShape::N * kNumZPBits / 32 / 4 * kNumGroupsPerMma; diff --git a/humming/include/humming/memory/s2r_loader/loader_bzp.cuh b/humming/include/humming/memory/s2r_loader/loader_bzp.cuh index 0cf7e6f1..0daf0a9e 100644 --- a/humming/include/humming/memory/s2r_loader/loader_bzp.cuh +++ b/humming/include/humming/memory/s2r_loader/loader_bzp.cuh @@ -14,7 +14,6 @@ private: static constexpr bool kIsFpZeroPoint = Ctx::kIsFpZeroPoint; static constexpr bool kIsChannelScale = Ctx::kIsChannelWeightScale; static constexpr bool kIsGroupScale = Ctx::kIsGroupWeightScale; - static constexpr bool kUseMxmma = Ctx::kUseMxmma; static constexpr uint32_t kGroupSize = kIsChannelScale ? BlockShape::K : Ctx::kWeightScaleGroupSize; static constexpr uint32_t kPartMmaShapeK = Ctx::kPartMmaShapeK; @@ -23,7 +22,7 @@ private: static constexpr uint32_t K_WARPS = Ctx::K_WARPS; static constexpr uint32_t kNumZPBits = kIsFpZeroPoint ? 16 : MAX(4, static_next_power_of_2(ElementB::kBits)); - static constexpr uint32_t kNumGroupsPerMma = kUseMxmma && ElementA::kBits == 4 && kIsGroupScale ? kPartMmaShapeK / kGroupSize : 1; + static constexpr uint32_t kNumGroupsPerMma = Ctx::kUseBlockScaledMma && ElementA::kBits == 4 && kIsGroupScale ? kPartMmaShapeK / kGroupSize : 1; static constexpr uint32_t kLoadBytes = WarpShape::N / 8 * kNumZPBits / 8 * kNumGroupsPerMma; using LoadType = typename LoadTypeChooser::Type; static constexpr uint32_t kSmemStride = BlockShape::N * kNumZPBits / 32 / 4 * kNumGroupsPerMma; diff --git a/humming/include/humming/mma/umma.cuh b/humming/include/humming/mma/umma.cuh index ea054a3b..e6478b6f 100644 --- a/humming/include/humming/mma/umma.cuh +++ b/humming/include/humming/mma/umma.cuh @@ -326,7 +326,8 @@ struct UMMA : WMMA { return base + buffer * kOperandBufferColumns + k * 8; } }(); - using ElementB = std::conditional_t; + // The TMEM path has already dequantized weights to the activation format. + using ElementB = std::conditional_t; // f8f6f4 descriptor format 2 selects the undocumented E3M4 format. constexpr uint32_t kWeightFormat = std::is_same::value ? 5 : std::is_same::value ? 4 diff --git a/humming/jit/runtime.py b/humming/jit/runtime.py index 66f69528..2c0dd159 100644 --- a/humming/jit/runtime.py +++ b/humming/jit/runtime.py @@ -10,6 +10,7 @@ from humming import dtypes from humming.device import current_device from humming.jit.compiler import NVCCCompiler, NVRTCCompiler +from humming.utils.cpu import get_parallel_build_workers from humming.utils.cubin import get_cubin_kernel_names @@ -111,10 +112,9 @@ def compile_kernel(spec): return kernel_type(**kernel_args) kernel_specs = list(kernel_specs) - parallel = len(kernel_specs) > 1 - parallel &= os.environ.get("HUMMING_DISABLE_PARALLEL_BUILD", "0") != "1" + workers = min(len(kernel_specs), get_parallel_build_workers()) if kernel_specs else 0 + parallel = workers > 1 if parallel: - workers = min(16, len(kernel_specs)) with ThreadPoolExecutor( max_workers=workers, ) as executor: diff --git a/humming/kernel/humming.py b/humming/kernel/humming.py index 59d0b67a..926549d9 100644 --- a/humming/kernel/humming.py +++ b/humming/kernel/humming.py @@ -303,12 +303,7 @@ def select_mma_op_class(self): and self.a_dtype in (dtypes.float8e4m3, dtypes.float8e5m2, dtypes.float8e3m4) and self.b_dtype in (dtypes.float4e2m1, dtypes.float6e3m2, dtypes.float6e2m3) ) - umma_native_mixed = ( - self.mma_type == MmaType.UMMA - and self.a_dtype.num_bits <= 8 - and self.b_dtype.is_floating_point_type - ) - use_native_weight = self.use_raw_weight or mma_native_mixed or umma_native_mixed + use_native_weight = self.use_raw_weight or mma_native_mixed self.mma_b_dtype = self.b_dtype if use_native_weight else self.a_dtype scale_dtype = dtypes.float8e8m0 @@ -423,7 +418,9 @@ def check_dtype(self): assert self.b_dtype.num_bits <= self.a_dtype.mantissa_bits + 2 elif self.b_dtype.is_floating_point_type and self.a_dtype.is_floating_point_type: assert self.b_dtype.is_signed - uses_native_umma = self.mma_type == MmaType.UMMA and self.a_dtype.num_bits <= 8 + uses_native_umma = ( + self.mma_type == MmaType.UMMA and self.a_dtype.num_bits <= 8 and self.use_raw_weight + ) if not self.use_block_scaled_mma and not uses_native_umma: assert self.b_dtype.exponent_bits <= self.a_dtype.exponent_bits assert self.b_dtype.mantissa_bits <= self.a_dtype.mantissa_bits @@ -497,15 +494,7 @@ def check_config(self): block_m = self.block_shape[0] assert block_m <= 32 or block_m % 16 == 0, "INT8 UMMA requires M divisible by 16 above M=32" elif self.a_dtype.num_bits == 8: - assert self.b_dtype.is_integer_type or self.b_dtype in ( - dtypes.float8e4m3, - dtypes.float8e5m2, - dtypes.float8e3m4, - dtypes.float4e2m1, - dtypes.float6e3m2, - dtypes.float6e2m3, - ) - assert not self.is_block_weight_scale and not self.has_zero_point + assert not self.is_block_weight_scale if self.use_block_scaled_mma: for group_size, scale_dtype in ( (self.input_scale_group_size, self.as_dtype), @@ -516,7 +505,7 @@ def check_config(self): "mxf8f6f4 requires E8M0 scales with group size 32" ) if self.a_dtype.num_bits == 4: - assert not self.is_block_weight_scale and not self.has_zero_point + assert not self.is_block_weight_scale assert self.b_dtype.is_integer_type or self.b_dtype in (dtypes.float4e2m1, dtypes.float4e0m3) assert self.use_block_scaled_mma for group_size, scale_dtype in ( diff --git a/humming/testing/runner.py b/humming/testing/runner.py index b817abfb..10ab8145 100644 --- a/humming/testing/runner.py +++ b/humming/testing/runner.py @@ -1,6 +1,7 @@ import dataclasses import json import os +import re from collections import Counter from datetime import datetime, timezone from pathlib import Path @@ -22,28 +23,48 @@ generate_random_topk_ids, ) from humming.testing.tuning import ( + NUM_SAMPLED_TUNING_CONFIGS, create_tuning_config, generate_heuristics_configs, sample_test_tuning_configs, ) from humming.transform import transform_humming_tensors -_DEFAULT_SHAPE_MS = (1, 17, 64, 257, 1024, 4096) +_DEFAULT_SHAPE_MS = (1, 30, 257, 1024, 4090) TEST_TUNING_SOURCE_ENV = "HUMMING_TEST_TUNING_SOURCE" NUMERICAL_ERROR_LOG_ENV = "HUMMING_TEST_NUMERICAL_ERROR_LOG" +def _get_sample_size(tuning_source: str) -> int | None: + match = re.fullmatch(r"sampled([1-9][0-9]*)?", tuning_source) + if match is None: + return None + return int(match[1]) if match[1] else NUM_SAMPLED_TUNING_CONFIGS + + +def _get_test_tuning_sources() -> tuple[str, ...]: + value = os.environ.get(TEST_TUNING_SOURCE_ENV, "heuristic+batch_invariant") + sources = tuple(dict.fromkeys(source.strip() for source in value.split("+"))) + for source in sources: + if source not in ("heuristic", "batch_invariant") and _get_sample_size(source) is None: + raise ValueError(f"invalid tuning source: {value}") + return sources + + def assert_kernel_test_shape_coverage( results: list["KernelTestResult"], shape_ms: list[int] | tuple[int, ...] | None = None, ) -> None: if shape_ms is None: shape_ms = _DEFAULT_SHAPE_MS - counts = Counter(result.shape_m for result in results) - assert set(counts) == set(shape_ms) - assert len(set(counts.values())) == 1 - expected_minimum = 100 if os.environ.get(TEST_TUNING_SOURCE_ENV) == "sampled" else 1 - assert next(iter(counts.values())) >= expected_minimum + tuning_sources = _get_test_tuning_sources() + assert {result.tuning_source for result in results} == set(tuning_sources) + for tuning_source in tuning_sources: + counts = Counter(result.shape_m for result in results if result.tuning_source == tuning_source) + assert set(counts) == set(shape_ms) + assert len(set(counts.values())) == 1 + expected_minimum = _get_sample_size(tuning_source) or 1 + assert next(iter(counts.values())) >= expected_minimum @dataclasses.dataclass(frozen=True, kw_only=True) @@ -83,6 +104,7 @@ def uses_m_major_input_scale(self) -> bool: @dataclasses.dataclass(frozen=True, kw_only=True) class KernelTestResult: + tuning_source: str shape_m: int tuning_config: TuningConfig tuning_values: dict @@ -95,13 +117,16 @@ def __init__(self, test_case: KernelTestCase): self.test_case = test_case self.layer_config = test_case.layer_config self.compute_config = test_case.compute_config + self.tuning_source = "heuristic" self.device = torch.cuda.current_device() self.weight_ref: torch.Tensor self.bias_ref: torch.Tensor | None = None self.kernel_tensors: dict[str, torch.Tensor] self.prepare_weight() - def prepare_kernels(self, shape_ms: tuple[int, ...]) -> dict[int, list[tuple[torch.Tensor, dict, int]]]: + def prepare_kernels( + self, shape_ms: tuple[int, ...] + ) -> list[tuple[str, ComputeConfig, dict[int, list[tuple[torch.Tensor, dict, int]]]]]: if self.layer_config.mxmma_supported: compiler_version = _cuda_compiler_version(HummingKernel._get_compiler()) compiler_error = get_mxmma_compiler_error(self.layer_config, compiler_version) @@ -110,49 +135,68 @@ def prepare_kernels(self, shape_ms: tuple[int, ...]) -> dict[int, list[tuple[tor pytest.skip(compiler_error) - tuning_source = os.environ.get(TEST_TUNING_SOURCE_ENV, "heuristic") - if tuning_source == "batch_invariant": - self.compute_config = dataclasses.replace(self.compute_config, use_batch_invariant=True) - - if tuning_source != "sampled" and get_default_mma_type(self.layer_config) == MmaType.WGMMA: - min_warp_shape_n = 32 if self.layer_config.a_dtype.num_bits == 16 else 16 - if self.layer_config.shape_n % (min_warp_shape_n * 4): + tuning_sources = _get_test_tuning_sources() + if any(_get_sample_size(source) is None for source in tuning_sources): + if get_default_mma_type(self.layer_config) == MmaType.WGMMA: + min_warp_shape_n = 32 if self.layer_config.a_dtype.num_bits == 16 else 16 + if self.layer_config.shape_n % (min_warp_shape_n * 4): + import pytest + + pytest.skip("shape_n cannot form four WGMMA warp tiles") + + groups = [] + compile_configs = [] + effective_shape_ms = [self.test_case.effective_shape_m(shape_m) for shape_m in shape_ms] + for tuning_source in tuning_sources: + compute_config = self.test_case.compute_config + if tuning_source == "batch_invariant": + compute_config = dataclasses.replace(compute_config, use_batch_invariant=True) + sample_size = _get_sample_size(tuning_source) + if sample_size is not None: + tuning_configs = sample_test_tuning_configs( + self.layer_config, compute_config, sample_size=sample_size + ) + else: + tuning_configs = generate_heuristics_configs( + self.layer_config, compute_config, effective_shape_ms + ) + if not tuning_configs: import pytest - pytest.skip("shape_n cannot form four WGMMA warp tiles") - - if tuning_source == "sampled": - tuning_configs = sample_test_tuning_configs(self.layer_config, self.compute_config) - elif tuning_source in ("heuristic", "batch_invariant"): - effective_shape_ms = [self.test_case.effective_shape_m(shape_m) for shape_m in shape_ms] - tuning_configs = generate_heuristics_configs( - self.layer_config, - self.compute_config, - effective_shape_ms, + pytest.skip(f"no legal tuning configs for {tuning_source}") + groups.append((tuning_source, compute_config, tuning_configs)) + compile_configs.extend( + (0, 1 << 30, values | {"use_batch_invariant": compute_config.use_batch_invariant}) + for values in tuning_configs ) - else: - raise ValueError(f"invalid tuning source: {tuning_source}") - - if not tuning_configs: - import pytest - - pytest.skip("no legal tuning configs for this layer and compute config") kernel_configs = HummingKernel.prepare_kernels( self.layer_config.to_str(), self.compute_config.to_str(), - [(0, 1 << 30, values) for values in tuning_configs], + compile_configs, device=self.device, ).reshape(-1, 4) - assert kernel_configs.shape[0] == len(tuning_configs) + assert kernel_configs.shape[0] == len(compile_configs) for kernel_id in set(kernel_configs[:, 2].tolist()): HummingKernel._id2kernel[kernel_id].assert_smem_size_matches_estimate() - enum_iter_objs = enumerate(zip(kernel_configs, tuning_configs, strict=True)) - kernels = [(*values, index) for index, values in enum_iter_objs] - if tuning_source == "sampled": - return dict.fromkeys(shape_ms, kernels) - return {shape_m: [kernel] for shape_m, kernel in zip(shape_ms, kernels, strict=True)} + prepared_groups = [] + offset = 0 + for tuning_source, compute_config, tuning_configs in groups: + group_configs = kernel_configs[offset : offset + len(tuning_configs)] + kernels = [ + (kernel_config, values, index) + for index, (kernel_config, values) in enumerate( + zip(group_configs, tuning_configs, strict=True) + ) + ] + if _get_sample_size(tuning_source) is not None: + shape_kernels = dict.fromkeys(shape_ms, kernels) + else: + shape_kernels = {shape_m: [kernel] for shape_m, kernel in zip(shape_ms, kernels, strict=True)} + prepared_groups.append((tuning_source, compute_config, shape_kernels)) + offset += len(tuning_configs) + return prepared_groups def prepare_weight(self) -> None: torch.manual_seed(self.test_case.seed + 123) @@ -436,6 +480,23 @@ def _launch_kernel( **self.kernel_tensors, ) + def _assert_close(self, outputs: torch.Tensor, outputs_ref: torch.Tensor, use_stream_k: bool) -> None: + rtol, atol = self.test_case.rtol, self.test_case.atol + try: + torch.testing.assert_close(outputs, outputs_ref, rtol=rtol, atol=atol) + except AssertionError: + if outputs.dtype != torch.bfloat16 or not use_stream_k: + raise + # Keep rtol unchanged and validate tensor metadata and non-finite values. + torch.testing.assert_close(outputs, outputs_ref, rtol=rtol, atol=2 * atol) + outliers = ~torch.isclose(outputs, outputs_ref, rtol=rtol, atol=atol).reshape(-1) + # Allow at most 0.01% isolated outliers, with a minimum allowance of two. + max_outliers = max(2, outputs.numel() // 10000) + if outliers.sum().item() > max_outliers: + raise + if torch.any(outliers[:-1] & outliers[1:]): + raise + def _run_kernel( self, shape_m: int, @@ -453,12 +514,8 @@ def _run_kernel( f"tuning_index={tuning_index}, tuning_config={tuning_values}" ) from error try: - torch.testing.assert_close( - outputs, - outputs_ref, - rtol=self.test_case.rtol, - atol=self.test_case.atol, - ) + compiled_kernel = HummingKernel._id2kernel[int(kernel_config[2])] + self._assert_close(outputs, outputs_ref, compiled_kernel.use_stream_k) except AssertionError as error: self._record_numerical_error(error, shape_m, tuning_values, tuning_index) raise AssertionError( @@ -472,6 +529,7 @@ def _run_kernel( ) from error return KernelTestResult( + tuning_source=self.tuning_source, shape_m=shape_m, tuning_config=create_tuning_config(tuning_values), tuning_values=tuning_values, @@ -489,7 +547,25 @@ def run( if not shape_ms: return [] - kernels = self.prepare_kernels(shape_ms) + prepared_groups = self.prepare_kernels(shape_ms) + results = [] + original_compute_config = self.compute_config + original_tuning_source = self.tuning_source + try: + for tuning_source, compute_config, kernels in prepared_groups: + self.tuning_source = tuning_source + self.compute_config = compute_config + results.extend(self._run_prepared_kernels(shape_ms, kernels)) + finally: + self.compute_config = original_compute_config + self.tuning_source = original_tuning_source + return results + + def _run_prepared_kernels( + self, + shape_ms: tuple[int, ...], + kernels: dict[int, list[tuple[torch.Tensor, dict, int]]], + ) -> list[KernelTestResult]: max_shape_m = max(shape_ms) shape_k = self.layer_config.shape_k - self.layer_config.pad_shape_k output_top_k = 1 if self.compute_config.gemm_type == GemmType.DENSE else self.test_case.top_k @@ -584,7 +660,7 @@ def _record_numerical_error( "pytest_test": os.environ.get("PYTEST_CURRENT_TEST"), "test_case": dataclasses.asdict(self.test_case), "shape_m": shape_m, - "tuning_source": os.environ.get(TEST_TUNING_SOURCE_ENV, "heuristic"), + "tuning_source": self.tuning_source, "tuning_index": tuning_index, "tuning_config": tuning_values, "device": current_device.name, diff --git a/humming/testing/tuning.py b/humming/testing/tuning.py index ed18f395..8f5c9b4f 100644 --- a/humming/testing/tuning.py +++ b/humming/testing/tuning.py @@ -17,7 +17,7 @@ TEST_TUNING_SEED_ENV = "HUMMING_TEST_TUNING_SEED" SAMPLED_TUNING_VALUES = { "mma_type": tuple(mma_type.value for mma_type in MmaType), - "num_stages": (2, 3, 4, 6, 8), + "num_stages": (2, 3, 4, 5, 6, 8), "use_tma": (True, False, 123, 456, 789), "use_warp_spec": (True, False), "use_mbarrier": (True, False), @@ -66,9 +66,7 @@ def _is_legal_mma_type(layer_config, compute_config, mma_type): if has_mixed_raw_weights and mma_type != MmaType.UMMA: return False if mma_type == MmaType.UMMA: - if not layer_config.is_umma_supported or compute_config.use_f16_accum: - return False - return layer_config.a_dtype.num_bits == 16 or not layer_config.has_zero_point + return layer_config.is_umma_supported and not compute_config.use_f16_accum if mma_type == MmaType.MXMMA: return sm_version // 10 == 12 and layer_config.use_block_scaled_mma if mma_type == MmaType.WGMMA: @@ -271,6 +269,9 @@ def _generate_transfer_candidates( continue if (use_tma or signature["use_warp_spec"]) and not signature["use_mbarrier"]: continue + requires_cp_async = signature["use_warp_spec"] or mma_type in (MmaType.WGMMA, MmaType.UMMA) + if requires_cp_async and not signature["use_cp_async"]: + continue if compute_config.gemm_type == GemmType.INDEXED: tma_values.update(use_tma_a=False, use_tma_as=False, use_tma_as2=False, use_tma_c=False) if not ( @@ -352,13 +353,17 @@ def _get_seed(layer_config: LayerConfig, compute_config: ComputeConfig) -> int: return int.from_bytes(hashlib.sha256(content).digest()[:8], "little") -def _select_pairwise(candidates: list[tuple[dict, dict]], rng: random.Random) -> list[tuple[dict, dict]]: +def _select_pairwise( + candidates: list[tuple[dict, dict]], rng: random.Random, max_count: int | None = None +) -> list[tuple[dict, dict]]: pair_sets = [frozenset(_get_covered_pairs(candidate)) for candidate in candidates] uncovered = set().union(*pair_sets) remaining = list(range(len(candidates))) rng.shuffle(remaining) selected = [] while uncovered and remaining: + if max_count is not None and len(selected) >= max_count: + break candidate_index = max(remaining, key=lambda index: len(pair_sets[index] & uncovered)) covered = pair_sets[candidate_index] & uncovered if not covered: @@ -391,19 +396,6 @@ def _fits_device_resources( registers_per_sm = current_device.max_registers_per_sm if num_threads * num_ctas_per_sm > max_threads: return False - launch_bound_registers = registers_per_sm // (num_threads * num_ctas_per_sm) // 8 * 8 - if mma_type == MmaType.UMMA and launch_bound_registers < 40: - # TMEM transfers and their address operands cannot compile at a 32-register limit. - return False - - if config["use_warp_spec"] and mma_type != MmaType.UMMA: - needs_more_load_registers = num_math_threads > 256 or ( - num_ctas_per_sm == 1 and layer_config.a_dtype.num_bits != 16 - ) - load_thread_registers = 40 if needs_more_load_registers else 24 - if launch_bound_registers < load_thread_registers: - return False - tuning_config = create_tuning_config(config) if ( get_register_budget_error(layer_config, tuning_config, compute_config.use_f16_accum, registers_per_sm) @@ -476,6 +468,7 @@ def _enumerate_backend_candidates( layer_config: LayerConfig, compute_config: ComputeConfig, mma_type: MmaType, + sample_size: int, ) -> list[tuple[dict, dict]]: rng = random.Random(_get_seed(layer_config, compute_config)) groups = ( @@ -499,7 +492,7 @@ def add(items): for items in itertools.product(*reduced_groups): add(items) - target_pool_size = NUM_SAMPLED_TUNING_CONFIGS * 5 + target_pool_size = sample_size * 5 product_size = math.prod(len(group) for group in groups) trial_count = min(product_size, target_pool_size * 500) for flat_index in rng.sample(range(product_size), trial_count): @@ -517,12 +510,15 @@ def add(items): def enumerate_test_tuning_configs( layer_config: LayerConfig, compute_config: ComputeConfig, + sample_size: int = NUM_SAMPLED_TUNING_CONFIGS, ) -> list[tuple[dict, dict]]: candidates = [] for value in SAMPLED_TUNING_VALUES["mma_type"]: mma_type = MmaType(value) if _is_legal_mma_type(layer_config, compute_config, mma_type): - candidates.extend(_enumerate_backend_candidates(layer_config, compute_config, mma_type)) + candidates.extend( + _enumerate_backend_candidates(layer_config, compute_config, mma_type, sample_size) + ) return candidates @@ -531,13 +527,13 @@ def sample_test_tuning_configs( compute_config: ComputeConfig, sample_size: int = NUM_SAMPLED_TUNING_CONFIGS, ) -> list[dict]: - candidates = enumerate_test_tuning_configs(layer_config, compute_config) + candidates = enumerate_test_tuning_configs(layer_config, compute_config, sample_size) rng = random.Random(_get_seed(layer_config, compute_config)) - selected = _select_pairwise(candidates, rng) + selected = _select_pairwise(candidates, rng, max_count=sample_size) selected_ids = {id(candidate) for candidate in selected} remaining = [candidate for candidate in candidates if id(candidate) not in selected_ids] rng.shuffle(remaining) - target_size = min(max(sample_size, len(selected)), len(candidates)) + target_size = min(sample_size, len(candidates)) selected.extend(remaining[: target_size - len(selected)]) rng.shuffle(selected) return [config for config, _ in selected] diff --git a/humming/transform.py b/humming/transform.py index 42f91918..365ee1eb 100644 --- a/humming/transform.py +++ b/humming/transform.py @@ -525,14 +525,13 @@ def transform_humming_tensors( ) if zero_point is not None: - use_mxmma_zp_layout = ( + use_grouped_fp4_zp_layout = ( config.use_block_scaled_mma - and config.sm_version // 10 == 12 and config.a_dtype.num_bits == 4 and config.weight_scale_group_size > 0 ) num_groups_per_mma = 1 - if use_mxmma_zp_layout: + if use_grouped_fp4_zp_layout: num_groups_per_mma = 256 // config.a_dtype.num_bits // config.weight_scale_group_size zero_point = transform_humming_zero_point( zero_point, diff --git a/humming/tune/__init__.py b/humming/tune/__init__.py index e5fdcf2f..72c22126 100644 --- a/humming/tune/__init__.py +++ b/humming/tune/__init__.py @@ -157,7 +157,8 @@ def get_heuristics_config( device_index = get_device_index(device) # Backend heuristics inspect the test environment; keep their cache entries separate. default_test_source = "heuristic" if "PYTEST_CURRENT_TEST" in os.environ else "" - is_heuristic_test = os.environ.get("HUMMING_TEST_TUNING_SOURCE", default_test_source) == "heuristic" + test_sources = os.environ.get("HUMMING_TEST_TUNING_SOURCE", default_test_source).split("+") + is_heuristic_test = "heuristic" in (source.strip() for source in test_sources) with torch.cuda.device(device_index): if isinstance(layer_config, dict): layer_config = LayerConfig(**layer_config) diff --git a/humming/tune/sm100.py b/humming/tune/sm100.py index 04dc5082..1e2b07ad 100644 --- a/humming/tune/sm100.py +++ b/humming/tune/sm100.py @@ -31,7 +31,8 @@ def get_config( config = super().get_config(layer_config, shape_m, use_f16_accum, use_batch_invariant, gemm_type) block_m, block_n, _ = config["block_shape"] warp_m = config["warp_shape"][0] - if use_batch_invariant: + # This four-warp schedule targets small M tiles. + if use_batch_invariant or block_m > 32: return config warp_k = 1024 // layer_config.a_dtype.num_bits @@ -591,11 +592,11 @@ def get_config( if shape_m <= 0: raise ValueError("shape_m must be positive") default_test_source = "heuristic" if "PYTEST_CURRENT_TEST" in os.environ else "" - is_heuristic_test = os.environ.get("HUMMING_TEST_TUNING_SOURCE", default_test_source) == "heuristic" + test_sources = os.environ.get("HUMMING_TEST_TUNING_SOURCE", default_test_source).split("+") + is_heuristic_test = "heuristic" in (source.strip() for source in test_sources) prefer_umma = is_heuristic_test and layer_config.is_umma_supported and not use_f16_accum prefer_umma &= layer_config.shape_n % 128 == 0 prefer_umma &= layer_config.shape_k % (512 // layer_config.a_dtype.num_bits) == 0 - prefer_umma &= layer_config.a_dtype.num_bits == 16 or not layer_config.has_zero_point if prefer_umma or get_default_mma_type(layer_config) == MmaType.UMMA: has_native_mixed_operands = layer_config.a_dtype.num_bits == 8 has_native_mixed_operands &= layer_config.b_dtype.is_floating_point_type diff --git a/humming/tune/sm120.py b/humming/tune/sm120.py index 9f56fe80..53cbe33d 100644 --- a/humming/tune/sm120.py +++ b/humming/tune/sm120.py @@ -20,6 +20,11 @@ class Sm120Heuristics(Sm89Heuristics): ] b4_allowed_dtypes: list[dtypes.DataType] = [dtypes.float4e2m1, dtypes.float4e0m3] + @staticmethod + def _has_complete_warpgroups(config) -> bool: + num_math_warps = math.prod(config["block_shape"]) // math.prod(config["warp_shape"]) + return num_math_warps % 4 == 0 + @classmethod def should_use_pdl_for_input(cls, layer_config, shape_m: int) -> bool: return layer_config.shape_n >= 4096 and shape_m <= 32 @@ -127,16 +132,10 @@ def get_config( is_wna16 = a.is_floating_point_type and a.num_bits == 16 and not use_f16_accum if a.is_floating_point_type and a.num_bits <= 8 and not layer_config.use_fused_e8m0_scale: config["use_tma"] = True - config["use_warp_spec"] = True + config["use_warp_spec"] = cls._has_complete_warpgroups(config) elif is_wna16: config["use_tma"] = True - block_shape = config["block_shape"] - warp_shape = config["warp_shape"] - m_warps = block_shape[0] // warp_shape[0] - n_warps = block_shape[1] // warp_shape[1] - k_warps = block_shape[2] // warp_shape[2] - num_math_threads = m_warps * n_warps * k_warps * 32 - config["use_warp_spec"] = num_math_threads % 128 == 0 + config["use_warp_spec"] = cls._has_complete_warpgroups(config) config["num_stages"] = cls._fit_num_stages( layer_config, config, @@ -209,6 +208,8 @@ def get_config( mn_warps = block_m // warp_m * (block_n // warp_n) target_k_warps = max(1, 4 // mn_warps) config["warp_shape"] = (warp_m, warp_n, max(warp_k, block_k // target_k_warps)) + if not cls._has_complete_warpgroups(config): + config["use_warp_spec"] = False cls._rebalance_dense_warps(layer_config, config, shape_m) if is_mxmma: @@ -278,7 +279,9 @@ def _rebalance_dense_warps(cls, layer_config, config, shape_m: int) -> None: block_m, block_n, _ = config["block_shape"] _, warp_n, _ = config["warp_shape"] mma_k = 1024 // layer_config.a_dtype.num_bits - block_k = min(layer_config.shape_k, mma_k * 2) + block_k = mma_k * 2 + while block_k > layer_config.shape_k: + block_k //= 2 if layer_config.shape_k % block_k: return @@ -350,6 +353,7 @@ def _use_m_tile(cls, layer_config, config, gemm_type, block_m: int) -> None: _, warp_n, warp_k = config["warp_shape"] config["block_shape"] = (block_m, block_n, block_k) config["warp_shape"] = (block_m, warp_n, warp_k) + config["use_warp_spec"] = config.get("use_warp_spec", False) and cls._has_complete_warpgroups(config) config["num_stages"] = cls._fit_num_stages( layer_config, config, diff --git a/humming/tune/sm121.py b/humming/tune/sm121.py index fca48b9a..dc443d0b 100644 --- a/humming/tune/sm121.py +++ b/humming/tune/sm121.py @@ -46,6 +46,7 @@ def _tune_small_m_dense(cls, layer_config, config, shape_m: int, is_wna16: bool) shape_k = layer_config.shape_k if num_a_bits >= 8 and use_tma and block_k >= warp_k * 2: config["block_shape"] = (block_m, block_n, block_k // 2) + config["use_warp_spec"] &= cls._has_complete_warpgroups(config) if is_wna16: cls._use_tma_b_only(config) elif num_b_bits <= 4 and not use_stream_k and block_k == 128 and shape_k % (block_k * 2) == 0: diff --git a/humming/tune/sm90.py b/humming/tune/sm90.py index 4e48793d..7a7b3fec 100644 --- a/humming/tune/sm90.py +++ b/humming/tune/sm90.py @@ -1,12 +1,12 @@ from humming import dtypes from humming.config import GemmType, LayerConfig from humming.device import current_device -from humming.tune.base import DeviceHeuristics from humming.tune.candidate import ( DeviceProfile, TuningDecision, TuningProblem, ) +from humming.tune.sm8x import Sm80Heuristics from humming.tune.sm90_policies import ( Sm90CandidatePolicy, build_sm90_seed_config, @@ -17,7 +17,7 @@ from humming.utils.smem import estimate_smem_size_layer -class Sm90Heuristics(DeviceHeuristics): +class Sm90Heuristics(Sm80Heuristics): max_smem_size: int = 227 * 1024 candidate_policy = Sm90CandidatePolicy() b16_allowed_dtypes: list[dtypes.DataType] = [dtypes.float16, dtypes.bfloat16] @@ -167,6 +167,9 @@ def get_config( use_batch_invariant: bool = False, gemm_type: GemmType = GemmType.DENSE, ): + if layer_config.a_dtype == dtypes.int4: + return super().get_config(layer_config, shape_m, use_f16_accum, use_batch_invariant, gemm_type) + use_grouped_scale_candidates = cls._uses_grouped_scale_candidates(layer_config) use_indexed_a16_policy = cls._uses_indexed_a16_policy(layer_config, use_batch_invariant, gemm_type) use_candidates = use_grouped_scale_candidates or use_indexed_a16_policy diff --git a/humming/tune/sm90_policies.py b/humming/tune/sm90_policies.py index 80872f47..bb7deb13 100644 --- a/humming/tune/sm90_policies.py +++ b/humming/tune/sm90_policies.py @@ -176,6 +176,12 @@ def build_sm90_seed_config(problem: TuningProblem) -> dict: block_shape_k = 64 warp_shape_n = 32 warp_shape_k = 64 + if problem.use_batch_invariant: + # Keep one K partition and the same reduction tile for every batch size. + batch_invariant_k = 128 if layer_config.use_packed_k_layout else 1024 // layer_config.a_dtype.num_bits + block_shape_k = min(batch_invariant_k, layer_config.shape_k & -layer_config.shape_k) + warp_shape_k = block_shape_k + config = { "block_shape": (block_shape_m, block_shape_n, block_shape_k), "warp_shape": (block_shape_m, warp_shape_n, warp_shape_k), @@ -218,7 +224,10 @@ def select_grouped_scale( problem.shape_m, max_block_m, ) - block_ks = (256, 128, 64) if block_shape_m <= 32 else (128, 64) + if problem.use_batch_invariant: + block_ks = (min(128, layer_config.shape_k & -layer_config.shape_k),) + else: + block_ks = (256, 128, 64) if block_shape_m <= 32 else (128, 64) use_multicast = problem.gemm_type == GemmType.DENSE and problem.shape_m / block_shape_m >= 4 candidates = [] diff --git a/humming/utils/cpu.py b/humming/utils/cpu.py new file mode 100644 index 00000000..074333c4 --- /dev/null +++ b/humming/utils/cpu.py @@ -0,0 +1,100 @@ +"""CPU limits shared by JIT compilation and the CI startup report.""" + +import os +import re +from pathlib import Path + + +def _read_text(path: Path) -> str: + try: + return path.read_text().strip() + except OSError: + return "" + + +def _decode_mount_path(value: str) -> Path: + return Path(re.sub(r"\\([0-7]{3})", lambda match: chr(int(match[1], 8)), value)) + + +def get_cgroup_cpu_count() -> int | None: + memberships = {} + for line in _read_text(Path("/proc/self/cgroup")).splitlines(): + _, controllers, path = line.split(":", 2) + if not controllers or "cpu" in controllers.split(","): + memberships["cgroup2" if not controllers else "cgroup"] = Path(path) + + limits = [] + for line in _read_text(Path("/proc/self/mountinfo")).splitlines(): + mount, separator, filesystem = line.partition(" - ") + if not separator: + continue + fields = filesystem.split() + filesystem_type = fields[0] + if filesystem_type not in memberships: + continue + if filesystem_type == "cgroup" and "cpu" not in fields[2].split(","): + continue + mount_fields = mount.split() + root = _decode_mount_path(mount_fields[3]) + mount_point = _decode_mount_path(mount_fields[4]) + membership = memberships[filesystem_type] + try: + relative_path = membership.relative_to(root) + except ValueError: + # A cgroup namespace can expose membership relative to its own root. + relative_path = membership.relative_to("/") + if ".." in relative_path.parts: + relative_path = Path(".") + directory = mount_point / relative_path + while True: + try: + if filesystem_type == "cgroup2": + quota, period = _read_text(directory / "cpu.max").split() + else: + quota = _read_text(directory / "cpu.cfs_quota_us") + period = _read_text(directory / "cpu.cfs_period_us") + if int(quota) > 0 and int(period) > 0: + limits.append(max(1, int(quota) // int(period))) + except ValueError: + pass # Missing or unlimited quota. + if directory == mount_point: + break + directory = directory.parent + return min(limits) if limits else None + + +def get_cpu_affinity_count() -> int: + try: + return len(os.sched_getaffinity(0)) + except (AttributeError, OSError): + return os.cpu_count() or 1 + + +def get_available_cpu_count() -> int: + available = min(os.cpu_count() or 1, get_cpu_affinity_count()) + quota = get_cgroup_cpu_count() + return max(1, min(available, quota)) if quota is not None else max(1, available) + + +def get_parallel_build_workers() -> int: + available = get_available_cpu_count() + if os.environ.get("HUMMING_DISABLE_PARALLEL_BUILD", "0") == "1": + return 1 + value = os.environ.get("HUMMING_MAX_PARALLEL_BUILD_WORKERS") + if value is None: + return available + try: + requested = int(value) + if requested > 0: + return min(available, requested) + except ValueError: + pass + raise ValueError("HUMMING_MAX_PARALLEL_BUILD_WORKERS must be a positive integer") + + +if __name__ == "__main__": + print(f"System CPU threads: {os.cpu_count() or 1}") + print(f"CPU affinity threads: {get_cpu_affinity_count()}") + print(f"Cgroup CPU quota threads (rounded down, minimum 1): {get_cgroup_cpu_count() or 'unlimited'}") + print(f"Available CPU threads: {get_available_cpu_count()}") + print(f"Parallel compilation worker limit: {get_parallel_build_workers()}") diff --git a/tests/kernels/humming/test_accumulation.py b/tests/kernels/humming/test_accumulation.py index 68a4472f..22d87e7f 100644 --- a/tests/kernels/humming/test_accumulation.py +++ b/tests/kernels/humming/test_accumulation.py @@ -120,6 +120,8 @@ def test_f16_accumulation(test_case): assert config.c_dtype == dtypes.float16 if config.use_block_scaled_mma: pytest.skip("block-scaled layers require FP32 accumulation") + if config.use_raw_weight and config.a_dtype.num_bits != config.b_dtype.num_bits: + pytest.skip("native mixed weight layout requires UMMA with FP32 accumulation") skip_if_unsupported(a_dtype=config.a_dtype, mma_type=get_default_mma_type(config).value) runner = KernelTestRunner(test_case) diff --git a/tests/kernels/humming/test_mxmma.py b/tests/kernels/humming/test_block_scaled.py similarity index 55% rename from tests/kernels/humming/test_mxmma.py rename to tests/kernels/humming/test_block_scaled.py index 9060a060..604c2a47 100644 --- a/tests/kernels/humming/test_mxmma.py +++ b/tests/kernels/humming/test_block_scaled.py @@ -1,9 +1,9 @@ import pytest -import torch from humming import dtypes from humming.config import ComputeConfig, GemmType, LayerConfig, MmaType from humming.config.mma import get_default_mma_type +from humming.device import current_device from humming.testing import ( KernelTestCase, KernelTestRunner, @@ -29,10 +29,18 @@ def _case( input_group_size: int | None = None, weight_group_size: int | None = None, gemm_type: GemmType = GemmType.DENSE, + use_m_major_input_scale: bool = False, + **layer_values, ) -> KernelTestCase: input_group_size = group_size if input_group_size is None else input_group_size weight_group_size = group_size if weight_group_size is None else weight_group_size + if input_group_size and not weight_group_size: + scale_dtype = dtypes.float8e4m3 if input_group_size == 16 else dtypes.float8e8m0 + layer_values.setdefault("as_dtype", scale_dtype) is_dense = gemm_type == GemmType.DENSE + sm_version = current_device.sm_version + if sm_version // 10 not in (10, 11, 12): + sm_version = 120 return KernelTestCase( name=name, layer_config=LayerConfig( @@ -46,15 +54,16 @@ def _case( input_scale_group_size=input_group_size, weight_scale_group_size=weight_group_size, has_zero_point=has_zero_point, - sm_version=120, + sm_version=sm_version, + **layer_values, ), - compute_config=ComputeConfig(gemm_type=gemm_type), + compute_config=ComputeConfig(gemm_type=gemm_type, use_m_major_input_scale=use_m_major_input_scale), top_k=1 if is_dense else TOP_K, seed=2026, ) -MXMMA_FORMAT_CASES = ( +FORMAT_CASES = ( _case( "e3m4-fp4-e8m0-g32", a_dtype=dtypes.float8e3m4, @@ -161,7 +170,16 @@ def _case( ), ) -MXMMA_ZERO_POINT_CASES = ( +ZERO_POINT_CASES = ( + _case( + "e4m3-uint4-channel-zp", + a_dtype=dtypes.float8e4m3, + b_dtype=dtypes.uint4, + bs_dtype=dtypes.bfloat16, + group_size=0, + has_zero_point=True, + input_quant_mode="static_tensor", + ), _case( "e3m4-uint5-e8m0-g32-zp", a_dtype=dtypes.float8e3m4, @@ -211,55 +229,122 @@ def _case( group_size=16, has_zero_point=True, ), + _case( + "e2m1-uint2-e8m0-g16-zp-indexed", + a_dtype=dtypes.float4e2m1, + b_dtype=dtypes.uint2, + bs_dtype=dtypes.float8e8m0, + group_size=16, + has_zero_point=True, + gemm_type=GemmType.INDEXED, + ), ) -MXMMA_CASES = MXMMA_FORMAT_CASES + MXMMA_ZERO_POINT_CASES +NATIVE_QUANTIZATION_CASES = tuple( + _case( + f"{a_dtype}-{b_dtype}-{scale_dtype}-g{group_size}-{quant_mode}-{gemm_type.value}", + a_dtype=a_dtype, + b_dtype=b_dtype, + bs_dtype=scale_dtype, + as_dtype=scale_dtype, + group_size=group_size, + input_quant_mode=quant_mode, + weight_scale_2_type="tensor", + has_bias=True, + gemm_type=gemm_type, + use_m_major_input_scale=gemm_type != GemmType.INDEXED, + ) + for a_dtype, b_dtype, group_size, scale_dtype, quant_mode, gemm_type in ( + ("float4e2m1", "float4e2m1", 32, "float8e8m0", "dynamic_group", GemmType.DENSE), + ("float4e2m1", "float4e2m1", 16, "float8e4m3", "dynamic_group_token", GemmType.INDEXED), + ("float4e2m1", "float4e2m1", 16, "float8e8m0", "static_tensor_dynamic_group", GemmType.DENSE), + ("float4e0m3", "float4e0m3", 16, "float8e4m3", "dynamic_group_token", GemmType.DENSE), + ("float4e0m3", "float4e0m3", 16, "float8e4m3", "dynamic_group_token", GemmType.GROUPED_MASKED), + ("float4e0m3", "float4e0m3", 16, "float8e8m0", "static_tensor_dynamic_group", GemmType.DENSE), + ("float4e0m3", "float4e2m1", 16, "float8e4m3", "dynamic_group_token", GemmType.INDEXED), + ("float4e0m3", "float4e2m1", 16, "float8e8m0", "dynamic_group", GemmType.DENSE), + ("float4e2m1", "float4e0m3", 16, "float8e8m0", "dynamic_group", GemmType.GROUPED_CONTIGUOUS), + ("float4e2m1", "float4e0m3", 16, "float8e4m3", "dynamic_group_token", GemmType.INDEXED), + ("float8e4m3", "float4e2m1", 32, "float8e8m0", "dynamic_group", GemmType.DENSE), + ("float8e5m2", "float6e3m2", 32, "float8e8m0", "dynamic_group", GemmType.INDEXED), + ("float8e3m4", "float8e3m4", 32, "float8e8m0", "dynamic_group", GemmType.DENSE), + ("float8e4m3", "float6e2m3", 32, "float8e8m0", "dynamic_group", GemmType.GROUPED_MASKED), + ) +) +OPTIONAL_SCALE_CASES = tuple( + _case( + f"{a_dtype}-{b_dtype}-as{input_group}-bs{weight_group}-{quant_mode}", + a_dtype=a_dtype, + b_dtype=b_dtype, + bs_dtype=dtypes.float8e8m0 if weight_group else dtypes.bfloat16, + as_dtype=dtypes.float8e8m0 if input_group else None, + group_size=32, + input_group_size=input_group, + weight_group_size=weight_group, + input_quant_mode=quant_mode, + ) + for a_dtype, b_dtype, input_group, weight_group, quant_mode in ( + ("float8e4m3", "float8e4m3", 0, 32, "dynamic_token"), + ("float8e4m3", "float4e2m1", 32, 0, "dynamic_group"), + ("float4e2m1", "float4e2m1", 0, 0, "dynamic_token"), + ("float8e4m3", "float8e4m3", 0, 0, "static_tensor"), + ("float8e4m3", "float4e2m1", 0, 32, "dynamic_token"), + ("float4e2m1", "float4e2m1", 32, 0, "dynamic_group"), + ) +) + +BLOCK_SCALED_CASES = FORMAT_CASES + ZERO_POINT_CASES + NATIVE_QUANTIZATION_CASES + OPTIONAL_SCALE_CASES -@pytest.mark.parametrize("test_case", MXMMA_CASES, ids=str) -def test_mxmma(test_case): + +@pytest.mark.parametrize("test_case", BLOCK_SCALED_CASES, ids=str) +def test_block_scaled(test_case): config = test_case.layer_config - assert get_default_mma_type(config) == MmaType.MXMMA + assert get_default_mma_type(config) in (MmaType.UMMA, MmaType.MXMMA) skip_if_unsupported(a_dtype=config.a_dtype, mma_type=get_default_mma_type(config).value) results = KernelTestRunner(test_case).run() - for result in results: - torch.testing.assert_close( - result.outputs, - result.outputs_ref, - rtol=test_case.rtol, - atol=test_case.atol, - ) assert_kernel_test_shape_coverage(results) -def test_mxmma_case_coverage(): - assert all(get_default_mma_type(case.layer_config) == MmaType.MXMMA for case in MXMMA_CASES) - assert {case.layer_config.a_dtype for case in MXMMA_CASES} == { +def test_block_scaled_case_coverage(): + assert all( + get_default_mma_type(case.layer_config) in (MmaType.UMMA, MmaType.MXMMA) + for case in BLOCK_SCALED_CASES + ) + assert {case.layer_config.a_dtype for case in BLOCK_SCALED_CASES} == { dtypes.float4e0m3, dtypes.float4e2m1, dtypes.float8e3m4, dtypes.float8e4m3, dtypes.float8e5m2, } - assert {case.layer_config.bs_dtype for case in MXMMA_CASES} == { + assert {case.layer_config.bs_dtype for case in BLOCK_SCALED_CASES} == { dtypes.bfloat16, dtypes.float8e4m3, dtypes.float8e8m0, } - assert {case.layer_config.weight_scale_group_size for case in MXMMA_CASES} == { + assert {case.layer_config.weight_scale_group_size for case in BLOCK_SCALED_CASES} == { 0, 16, 32, } + assert {case.compute_config.gemm_type for case in BLOCK_SCALED_CASES} == { + GemmType.DENSE, + GemmType.INDEXED, + GemmType.GROUPED_CONTIGUOUS, + GemmType.GROUPED_MASKED, + } + assert any(case.layer_config.has_bias for case in BLOCK_SCALED_CASES) + assert any(case.compute_config.use_m_major_input_scale for case in BLOCK_SCALED_CASES) assert any( case.layer_config.a_dtype in (dtypes.float8e4m3, dtypes.float8e5m2) and case.layer_config.b_dtype in (dtypes.float4e2m1, dtypes.float6e3m2, dtypes.float6e2m3) - for case in MXMMA_CASES + for case in BLOCK_SCALED_CASES ) - assert len(MXMMA_ZERO_POINT_CASES) == 6 - assert all(case.layer_config.has_zero_point for case in MXMMA_ZERO_POINT_CASES) - assert {case.layer_config.a_dtype for case in MXMMA_ZERO_POINT_CASES} == { + assert len(ZERO_POINT_CASES) == 8 + assert all(case.layer_config.has_zero_point for case in ZERO_POINT_CASES) + assert {case.layer_config.a_dtype for case in ZERO_POINT_CASES} == { dtypes.float4e0m3, dtypes.float4e2m1, dtypes.float8e3m4, diff --git a/tests/kernels/humming/test_moe.py b/tests/kernels/humming/test_moe.py index cb3aaf43..2ef44a85 100644 --- a/tests/kernels/humming/test_moe.py +++ b/tests/kernels/humming/test_moe.py @@ -4,6 +4,7 @@ from humming import dtypes, ops from humming.config import ComputeConfig, GemmType, LayerConfig from humming.config.mma import get_default_mma_type +from humming.device import current_device from humming.kernel.humming import HummingKernel from humming.testing import ( KernelTestCase, @@ -17,6 +18,9 @@ SHAPE_K = 1024 NUM_EXPERTS = 8 TOP_K = 2 +BLOCK_SCALED_SM_VERSION = ( + current_device.sm_version if current_device.sm_version // 10 in (10, 11, 12) else 120 +) def _case( @@ -104,7 +108,7 @@ def _case( has_zero_point=True, is_fp_zero_point=True, ), - _case("indexed", GemmType.INDEXED), + _case("indexed", GemmType.INDEXED, weight_scale_group_size=128), _case( "indexed-static-input", GemmType.INDEXED, @@ -116,30 +120,40 @@ def _case( GemmType.INDEXED, a_dtype=dtypes.float4e2m1, b_dtype=dtypes.float4e2m1, + as_dtype=dtypes.float8e4m3, input_scale_group_size=16, input_quant_mode="static_tensor_dynamic_group", - sm_version=120, + sm_version=BLOCK_SCALED_SM_VERSION, ), _case( "indexed-dynamic-group-token", GemmType.INDEXED, a_dtype=dtypes.float4e2m1, b_dtype=dtypes.float4e2m1, + as_dtype=dtypes.float8e4m3, input_scale_group_size=16, input_quant_mode="dynamic_group_token", - sm_version=120, + sm_version=BLOCK_SCALED_SM_VERSION, ), _case("indexed-partial-k-tile", GemmType.INDEXED, shape_k=96), _case("grouped-contiguous", GemmType.GROUPED_CONTIGUOUS), + _case( + "grouped-contiguous-short-k-zp", + GemmType.GROUPED_CONTIGUOUS, + shape_k=192, + weight_scale_group_size=64, + has_zero_point=True, + ), _case( "grouped-contiguous-dynamic-group-token", GemmType.GROUPED_CONTIGUOUS, use_m_major_input_scale=True, a_dtype=dtypes.float4e2m1, b_dtype=dtypes.float4e2m1, + as_dtype=dtypes.float8e4m3, input_scale_group_size=16, input_quant_mode="dynamic_group_token", - sm_version=120, + sm_version=BLOCK_SCALED_SM_VERSION, ), _case("grouped-masked", GemmType.GROUPED_MASKED), _case( @@ -148,9 +162,10 @@ def _case( use_m_major_input_scale=True, a_dtype=dtypes.float4e2m1, b_dtype=dtypes.float4e2m1, + as_dtype=dtypes.float8e4m3, input_scale_group_size=16, input_quant_mode="dynamic_group_token", - sm_version=120, + sm_version=BLOCK_SCALED_SM_VERSION, ), _case( "indexed-bias-pad-k", @@ -219,34 +234,6 @@ def test_moe_case_coverage(): } -def test_indexed_mma_tma_scale_pipeline(monkeypatch): - skip_if_unsupported(a_dtype=dtypes.bfloat16, mma_type="mma", use_tma=True) - case = next(case for case in MOE_CASES if case.name == "indexed-partial-k-tile") - config = dict( - mma_type="mma", - block_shape=(32, 128, 32), - warp_shape=(16, 64, 32), - num_stages=4, - num_ctas_per_sm=1, - use_tma=True, - use_tma_a=False, - use_tma_b=False, - use_tma_bs=True, - use_tma_c=False, - use_warp_spec=False, - use_mbarrier=True, - use_cp_async=False, - use_stream_k=False, - output_chunk_rows=32, - smem_reuse_mode="none", - ) - monkeypatch.setenv("HUMMING_TEST_TUNING_SOURCE", "heuristic") - monkeypatch.setattr("humming.testing.tuning.get_heuristics_config", lambda *args, **kwargs: dict(config)) - shape_ms = (1, 17, 257) - results = KernelTestRunner(case).run(shape_ms) - assert_kernel_test_shape_coverage(results, shape_ms) - - @pytest.mark.parametrize("raster_group_m", [1, 8]) @pytest.mark.parametrize("use_stream_k", [False, True]) @pytest.mark.parametrize("offset_dtype", [torch.int32, torch.int64]) @@ -312,68 +299,3 @@ def launch(): torch.testing.assert_close(output[start : start + count], reference, atol=0.01, rtol=0.01) start += count assert (output[start:] == -123).all() - - -@pytest.mark.parametrize( - "gemm_type", - (GemmType.INDEXED, GemmType.GROUPED_CONTIGUOUS, GemmType.GROUPED_MASKED), -) -@pytest.mark.parametrize("multicast", (1, 2)) -@pytest.mark.parametrize("use_stream_k", (False, True)) -def test_raw_wgmma_tma_b_general_moe(gemm_type, multicast, use_stream_k, monkeypatch): - skip_if_unsupported(a_dtype=dtypes.float8e4m3, mma_type="wgmma") - case = _case( - "raw-wgmma-tma-b-general", - gemm_type, - a_dtype=dtypes.float8e4m3, - b_dtype=dtypes.float8e4m3, - has_bias=True, - ) - tuning = dict( - mma_type="wgmma", - block_shape=(16, 256, 128), - warp_shape=(16, 32, 128), - num_stages=3, - use_tma=True, - use_warp_spec=True, - use_stream_k=use_stream_k, - multi_cast_size_b=multicast, - use_tma_a=gemm_type != GemmType.INDEXED, - use_tma_c=gemm_type != GemmType.INDEXED, - ) - monkeypatch.setattr("humming.testing.tuning.get_heuristics_config", lambda *args, **kwargs: dict(tuning)) - if multicast > 1: - with pytest.raises(AssertionError, match="B multicast requires dense GEMM"): - KernelTestRunner(case).run((1,)) - return - results = KernelTestRunner(case).run((1, 17, 129)) - assert_kernel_test_shape_coverage(results, (1, 17, 129)) - - -@pytest.mark.parametrize("gemm_type", (GemmType.GROUPED_CONTIGUOUS, GemmType.GROUPED_MASKED)) -@pytest.mark.parametrize("a_dtype,b_dtype", (("float8e4m3", "float8e4m3"), ("bfloat16", "uint4"))) -@pytest.mark.parametrize("pad_k", (0, 128, 32)) -@pytest.mark.parametrize("use_stream_k", (False, True)) -def test_wgmma_tma_a_general_moe(gemm_type, a_dtype, b_dtype, pad_k, use_stream_k, monkeypatch): - skip_if_unsupported(a_dtype=a_dtype, mma_type="wgmma") - case = _case( - "wgmma-tma-a-general-moe", - gemm_type, - a_dtype=a_dtype, - b_dtype=b_dtype, - shape_k=1024, - pad_shape_k=pad_k, - has_bias=True, - ) - tuning = dict( - mma_type="wgmma", - block_shape=(16, 128, 256), - warp_shape=(16, 32, 64 if a_dtype == "bfloat16" else 128), - num_stages=3, - use_tma=True, - use_warp_spec=True, - use_stream_k=use_stream_k, - ) - monkeypatch.setattr("humming.testing.tuning.get_heuristics_config", lambda *args, **kwargs: dict(tuning)) - results = KernelTestRunner(case).run((1, 17, 129)) - assert_kernel_test_shape_coverage(results, (1, 17, 129)) diff --git a/tests/kernels/humming/test_mxfp4.py b/tests/kernels/humming/test_mxfp4.py index d47ed1a9..a76592cb 100644 --- a/tests/kernels/humming/test_mxfp4.py +++ b/tests/kernels/humming/test_mxfp4.py @@ -6,7 +6,6 @@ from humming import dtypes from humming.config import ComputeConfig, GemmType, LayerConfig from humming.config.mma import get_default_mma_type -from humming.device import current_device from humming.schema.compressed_tensors import CompressedTensorsInputSchema from humming.schema.humming import HummingWeightSchema from humming.testing import ( @@ -206,56 +205,3 @@ def test_mxfp4_input_schema_compatibility(checkpoint_format, group_size): ).to_humming_schema(torch.bfloat16) assert inputs.input_scale_dtype is None assert inputs.is_compatible_with(weight, torch.bfloat16) == (group_size == INPUT_GROUP_SIZE) - - -@pytest.mark.parametrize( - "a_dtype,b_dtype,group_size,scale_dtype,quant_mode,gemm_type", - ( - ("float4e2m1", "float4e2m1", 32, "float8e8m0", "dynamic_group", GemmType.DENSE), - ("float4e2m1", "float4e2m1", 16, "float8e4m3", "dynamic_group_token", GemmType.INDEXED), - ("float4e2m1", "float4e2m1", 16, "float8e8m0", "static_tensor_dynamic_group", GemmType.DENSE), - ("float4e0m3", "float4e0m3", 16, "float8e4m3", "dynamic_group_token", GemmType.DENSE), - ("float4e0m3", "float4e2m1", 16, "float8e4m3", "dynamic_group_token", GemmType.INDEXED), - ("float4e0m3", "float4e2m1", 16, "float8e8m0", "dynamic_group", GemmType.DENSE), - ("float4e2m1", "float4e0m3", 16, "float8e8m0", "dynamic_group", GemmType.GROUPED_CONTIGUOUS), - ("float4e2m1", "float4e0m3", 16, "float8e4m3", "dynamic_group_token", GemmType.INDEXED), - ("float8e4m3", "float4e2m1", 32, "float8e8m0", "dynamic_group", GemmType.DENSE), - ("float8e5m2", "float6e3m2", 32, "float8e8m0", "dynamic_group", GemmType.INDEXED), - ("float8e3m4", "float8e3m4", 32, "float8e8m0", "dynamic_group", GemmType.DENSE), - ("float8e4m3", "float6e2m3", 32, "float8e8m0", "dynamic_group", GemmType.GROUPED_MASKED), - ), -) -def test_native_block_scaled(a_dtype, b_dtype, group_size, scale_dtype, quant_mode, gemm_type): - """Numerical contracts shared by native block-scaled backends and sampled tuning.""" - if current_device.sm_version < 100: - pytest.skip("native block-scaled MMA requires SM100 or newer") - layer = LayerConfig( - shape_n=512, - shape_k=1024, - a_dtype=a_dtype, - b_dtype=b_dtype, - c_dtype=dtypes.bfloat16, - as_dtype=scale_dtype, - bs_dtype=scale_dtype, - input_scale_group_size=group_size, - weight_scale_group_size=group_size, - input_quant_mode=quant_mode, - weight_scale_2_type="tensor", - has_bias=True, - num_experts=0 if gemm_type == GemmType.DENSE else 4, - ) - backend = get_default_mma_type(layer) - skip_if_unsupported(a_dtype=layer.a_dtype, mma_type=backend.value) - if not layer.use_block_scaled_mma: - pytest.skip("the device does not support this native block-scaled format") - case = KernelTestCase( - name="native-block-scaled", - layer_config=layer, - compute_config=ComputeConfig( - gemm_type=gemm_type, - use_m_major_input_scale=gemm_type != GemmType.INDEXED, - ), - seed=2026, - ) - results = KernelTestRunner(case).run((1, 17, 129)) - assert_kernel_test_shape_coverage(results, (1, 17, 129)) diff --git a/tests/kernels/humming/test_scale.py b/tests/kernels/humming/test_scale.py index dbbcc00b..54087fd3 100644 --- a/tests/kernels/humming/test_scale.py +++ b/tests/kernels/humming/test_scale.py @@ -236,99 +236,3 @@ def test_scale_config_case_coverage(): for bs_dtype in GROUP_SCALE_DTYPES: for scale_type in SECONDARY_SCALE_TYPES: assert (bs_dtype, scale_type) in secondary_scale_pairs - - -@pytest.mark.parametrize("warp_n", (16, 32, 64)) -@pytest.mark.parametrize("use_tma", (False, True)) -@pytest.mark.parametrize("bs_dtype", (dtypes.bfloat16, dtypes.float8e4m3)) -def test_raw_wgmma_ss_group_scale_layout(warp_n, use_tma, bs_dtype, monkeypatch): - """Group scales follow SS's N16 fragments, including wide per-warp tiles.""" - skip_if_unsupported(a_dtype=dtypes.float8e4m3, mma_type="wgmma") - case = _case( - "raw-wgmma-ss-group64", - a_dtype=dtypes.float8e4m3, - b_dtype=dtypes.float8e4m3, - bs_dtype=bs_dtype, - input_scale_group_size=64, - weight_scale_group_size=64, - use_int_weight_scale=False, - use_fused_e8m0_scale=False, - ) - tuning = dict( - mma_type="wgmma", - block_shape=(16, warp_n * 4, 256), - warp_shape=(16, warp_n, 128), - num_stages=3, - use_tma=use_tma, - use_warp_spec=use_tma, - use_stream_k=False, - ) - monkeypatch.setattr("humming.testing.tuning.get_heuristics_config", lambda *args, **kwargs: dict(tuning)) - results = KernelTestRunner(case).run((1, 17, 129)) - assert_kernel_test_shape_coverage(results, (1, 17, 129)) - - -@pytest.mark.parametrize("warp_n", (16, 32, 64)) -@pytest.mark.parametrize("input_group", (0, 64)) -def test_raw_wgmma_ss_block_scale_layout(warp_n, input_group, monkeypatch): - skip_if_unsupported(a_dtype=dtypes.float8e4m3, mma_type="wgmma") - case = _case( - "raw-wgmma-ss-block64", - a_dtype=dtypes.float8e4m3, - b_dtype=dtypes.float8e4m3, - bs_dtype=dtypes.float32, - weight_scale_type=WeightScaleType.BLOCK, - weight_scale_group_size=64, - weight_scale_group_size_n=64, - input_scale_group_size=input_group, - ) - tuning = dict( - mma_type="wgmma", - block_shape=(16, warp_n * 4, 128), - warp_shape=(16, warp_n, 128), - num_stages=3, - use_tma=True, - use_warp_spec=True, - use_stream_k=False, - ) - monkeypatch.setattr("humming.testing.tuning.get_heuristics_config", lambda *args, **kwargs: dict(tuning)) - results = KernelTestRunner(case).run((1, 17, 129)) - assert_kernel_test_shape_coverage(results, (1, 17, 129)) - - -@pytest.mark.parametrize( - "a_dtype,b_dtype,input_group,weight_group,quant_mode", - ( - ("float8e4m3", "float8e4m3", 0, 32, "dynamic_token"), - ("float8e4m3", "float4e2m1", 32, 0, "dynamic_group"), - ("float4e2m1", "float4e2m1", 0, 0, "dynamic_token"), - ("float8e4m3", "float8e4m3", 0, 0, "static_tensor"), - ("float8e4m3", "float4e2m1", 0, 32, "dynamic_token"), - ("float4e2m1", "float4e2m1", 32, 0, "dynamic_group"), - ), -) -def test_optional_native_group_scales(a_dtype, b_dtype, input_group, weight_group, quant_mode): - if current_device.sm_version < 100: - pytest.skip("native block-scaled MMA requires SM100 or newer") - layer = LayerConfig( - shape_n=256, - shape_k=512, - a_dtype=a_dtype, - b_dtype=b_dtype, - c_dtype=dtypes.bfloat16, - as_dtype=dtypes.float8e8m0 if input_group else None, - bs_dtype=dtypes.float8e8m0 if weight_group else dtypes.bfloat16, - input_quant_mode=quant_mode, - input_scale_group_size=input_group, - weight_scale_group_size=weight_group, - weight_scale_type="group" if weight_group else "channel", - ) - skip_if_unsupported(a_dtype=layer.a_dtype, mma_type=get_default_mma_type(layer).value) - case = KernelTestCase( - name="optional-native-group-scales", - layer_config=layer, - compute_config=ComputeConfig(gemm_type=GemmType.DENSE), - seed=2026, - ) - results = KernelTestRunner(case).run((1, 17, 129)) - assert_kernel_test_shape_coverage(results, (1, 17, 129)) diff --git a/tests/kernels/humming/test_shape.py b/tests/kernels/humming/test_shape.py index 93d9ad2f..5fdcc927 100644 --- a/tests/kernels/humming/test_shape.py +++ b/tests/kernels/humming/test_shape.py @@ -3,6 +3,7 @@ from humming import dtypes from humming.config import ComputeConfig, GemmType, LayerConfig, MmaType from humming.config.mma import get_default_mma_type +from humming.device import current_device from humming.testing import ( KernelTestCase, KernelTestRunner, @@ -12,6 +13,9 @@ MIN_SHAPE_N = 64 MIN_SHAPE_K = 32 +BLOCK_SCALED_SM_VERSION = ( + current_device.sm_version if current_device.sm_version // 10 in (10, 11, 12) else 120 +) def _case( @@ -27,9 +31,11 @@ def _case( as_dtype=None, input_scale_group_size: int = 0, input_quant_mode: str | None = None, + weight_scale_type: str | None = None, weight_scale_group_size: int = 0, weight_scale_group_size_n: int = 0, mma_type: MmaType | None = None, + sm_version: int | None = None, ) -> KernelTestCase: return KernelTestCase( name=name, @@ -45,9 +51,10 @@ def _case( as_dtype=as_dtype, input_scale_group_size=input_scale_group_size, input_quant_mode=input_quant_mode, + weight_scale_type=weight_scale_type, weight_scale_group_size=weight_scale_group_size, weight_scale_group_size_n=weight_scale_group_size_n, - sm_version=90 if mma_type == MmaType.WGMMA else None, + sm_version=90 if mma_type == MmaType.WGMMA else sm_version, ), compute_config=ComputeConfig(gemm_type=GemmType.DENSE), seed=2026, @@ -55,6 +62,12 @@ def _case( PROBLEM_SHAPE_CASES = ( + _case( + "short-k-group-scale", + shape_n=256, + shape_k=64, + weight_scale_group_size=64, + ), _case( "minimum", shape_n=MIN_SHAPE_N, @@ -164,6 +177,7 @@ def _case( NATIVE_SHAPE_CASES = tuple( _case( f"mxfp8-mxfp4-n{shape_n}-k{shape_k}", + sm_version=BLOCK_SCALED_SM_VERSION, shape_n=shape_n, shape_k=shape_k, a_dtype=dtypes.float8e4m3, @@ -176,6 +190,15 @@ def _case( ) for shape_n, shape_k in ((128, 64), (384, 192), (256, 2816), (256, 2880), (256, 2944), (512, 4160)) ) + ( + _case( + "fp8-large-rectangular", + shape_n=4096, + shape_k=8192, + a_dtype=dtypes.float8e4m3, + b_dtype=dtypes.float8e4m3, + input_quant_mode="static_tensor", + weight_scale_type="tensor", + ), _case( "fp8-uint4-n192-k2880", shape_n=192, diff --git a/tests/kernels/humming/test_special_paths.py b/tests/kernels/humming/test_special_paths.py index 7d47d313..a5d41bfd 100644 --- a/tests/kernels/humming/test_special_paths.py +++ b/tests/kernels/humming/test_special_paths.py @@ -1,14 +1,10 @@ -import dataclasses - import pytest import torch from torch._dynamo.testing import CompileCounterWithBackend -import humming.testing.runner as runner_module from humming import dtypes from humming.config import ComputeConfig, GemmType, LayerConfig, MmaType, WeightScale2Type from humming.config.mma import get_default_mma_type -from humming.device import current_device from humming.forward import humming_forward from humming.testing import ( KernelTestCase, @@ -23,49 +19,6 @@ GROUPED_INPUT_SIZE = 128 -@pytest.mark.parametrize("use_cp_async", (False, True)) -@pytest.mark.parametrize("block_k", (64, 128, 1024)) -@pytest.mark.parametrize("gemm_type", (GemmType.DENSE, GemmType.INDEXED)) -def test_ppu_raw_weight_loading(use_cp_async, block_k, gemm_type, monkeypatch): - skip_if_unsupported(a_dtype=dtypes.int8, use_cp_async=use_cp_async) - if not current_device.is_ppu: - pytest.skip("PPU operand layout and AIU regression") - # K=1024 needs eight slabs with four load warps, exercising the AIU slab loops. - config = dict( - mma_type="mma", - block_shape=(32, 64, block_k), - warp_shape=(32, 16, block_k), - num_stages=2, - num_ctas_per_sm=1, - use_stream_k=False, - use_cp_async=use_cp_async, - use_tma=False, - use_mbarrier=False, - ) - monkeypatch.setenv("HUMMING_TEST_TUNING_SOURCE", "heuristic") - monkeypatch.setattr("humming.testing.tuning.get_heuristics_config", lambda *args, **kwargs: config) - case = KernelTestCase( - name="ppu-raw-weight-loading", - layer_config=LayerConfig( - shape_n=256, - shape_k=2048, - pad_shape_n=24, - pad_shape_k=32, - a_dtype=dtypes.int8, - b_dtype=dtypes.int8, - c_dtype=dtypes.float16, - num_experts=0 if gemm_type == GemmType.DENSE else 4, - has_bias=True, - ), - compute_config=ComputeConfig(gemm_type=gemm_type), - top_k=1 if gemm_type == GemmType.DENSE else 2, - seed=2026, - ) - shape_ms = (1, 17, 64) - results = KernelTestRunner(case).run(shape_ms) - assert_kernel_test_shape_coverage(results, shape_ms) - - SPECIAL_FEATURES = { "use_int_weight_scale", "use_fused_e8m0_scale", @@ -243,6 +196,7 @@ def _kernel_case( ) +@pytest.mark.filterwarnings(r"ignore:`torch\.jit\.script_method` is deprecated\. ") def test_forward_fullgraph(): """Catch graph breaks in the forward path and reuse the graph across token counts.""" torch._dynamo.reset() @@ -326,60 +280,6 @@ def test_special_weight_path_coverage(): assert odd_bit_fallback.use_packed_k_layout is False -@pytest.mark.parametrize( - "a_dtype,use_fused,has_zero_point", - [ - (dtypes.float8e4m3, False, False), - (dtypes.int8, False, False), - (dtypes.float8e4m3, True, False), - (dtypes.int8, True, False), - (dtypes.float8e4m3, False, True), - ], - ids=["fp8", "int8", "fp8-fused", "int8-fused", "fp8-zp"], -) -@pytest.mark.parametrize( - "warp_m,warp_k,warp_n,k_warps", - [(64, 128, 16, 1), (64, 128, 16, 2), (64, 128, 32, 1), (64, 128, 64, 2), (176, 128, 16, 2)], -) -@pytest.mark.parametrize("use_warp_spec", [False, True]) -def test_packed_k_geometry( - monkeypatch, warp_m, warp_k, warp_n, k_warps, use_warp_spec, a_dtype, use_fused, has_zero_point -): - skip_if_unsupported(a_dtype=a_dtype, mma_type="wgmma") - case = next( - case - for _, case in SPECIAL_WEIGHT_CASES - if case.layer_config.use_packed_k_layout - and case.layer_config.use_fused_e8m0_scale == use_fused - and case.layer_config.a_dtype == a_dtype - and case.layer_config.has_zero_point == has_zero_point - and (not use_fused or case.layer_config.input_scale_group_size == 128) - and case.layer_config.weight_scale_group_size == (32 if use_fused else 128) - ) - layer = dataclasses.replace(case.layer_config, num_experts=33) - compute = dataclasses.replace( - case.compute_config, gemm_type=GemmType.GROUPED_CONTIGUOUS, use_m_major_input_scale=True - ) - case = dataclasses.replace(case, layer_config=layer, compute_config=compute) - tuning = dict( - block_shape=(warp_m, warp_n * 4, warp_k * k_warps), - warp_shape=(warp_m, warp_n, warp_k), - num_stages=3, - use_warp_spec=use_warp_spec, - use_stream_k=k_warps > 1, - raster_group_m=8, - multi_cast_size_a=1, - multi_cast_size_b=1, - ) - monkeypatch.setattr( - runner_module, - "generate_heuristics_configs", - lambda layer, compute, shape_ms: [dict(tuning) for _ in shape_ms], - ) - results = KernelTestRunner(case).run(shape_ms=[17, 257]) - assert_kernel_test_shape_coverage(results, [17, 257]) - - @pytest.mark.parametrize( "mma_type,block_m,chunk_rows,shape_n,gemm_type,use_tma,cta_group_size", ( diff --git a/tests/kernels/humming/test_umma.py b/tests/kernels/humming/test_umma.py deleted file mode 100644 index 0eced495..00000000 --- a/tests/kernels/humming/test_umma.py +++ /dev/null @@ -1,772 +0,0 @@ -import dataclasses - -import pytest -import torch - -from humming import dtypes -from humming.config import ComputeConfig, GemmType, LayerConfig, MmaType -from humming.config.config import _cuda_compiler_version -from humming.config.mma import get_default_mma_type -from humming.jit.runtime import KernelRuntime -from humming.kernel.humming import HummingKernel -from humming.layer import HummingLayer -from humming.schema import HummingWeightSchema -from humming.testing import KernelTestCase, KernelTestRunner -from humming.testing.data import generate_moe_tensors, generate_random_tensor -from humming.tune.sm100 import Sm100Heuristics, Sm100MmaHeuristics, Sm100UmmaHeuristics - -WEIGHT_CONFIGS = { - "uint4": dict(b_dtype="uint4", weight_scale_group_size=128), - "uint4-zp": dict(b_dtype="uint4", weight_scale_group_size=128, has_zero_point=True), - "uint4-fp-zp": dict( - b_dtype="uint4", - weight_scale_group_size=128, - has_zero_point=True, - is_fp_zero_point=True, - ), - "nvfp4": dict( - b_dtype="float4e2m1", - bs_dtype="float8e4m3", - weight_scale_group_size=16, - weight_scale_2_type="tensor", - ), - "mxfp4": dict( - b_dtype="float4e2m1", - bs_dtype="float8e8m0", - weight_scale_group_size=32, - ), - "fp8": dict(b_dtype="float8e4m3"), -} - - -@pytest.fixture(autouse=True) -def require_umma_device(monkeypatch): - if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] not in (10, 11): - pytest.skip("UMMA requires an SM10x or SM11x GPU") - if _cuda_compiler_version(KernelRuntime._get_compiler()) < (12, 9): - pytest.skip("UMMA requires CUDA 12.9 or newer") - monkeypatch.setattr(torch.backends.cuda.matmul, "allow_tf32", False) - # These regressions require the exact backend and geometry selected below. - monkeypatch.setenv("HUMMING_TEST_TUNING_SOURCE", "heuristic") - - def force_umma(layer_config, shape_m, gemm_type, **kwargs): - return Sm100Heuristics.get_umma_config(layer_config, shape_m, gemm_type) | {"mma_type": "umma"} - - monkeypatch.setattr("humming.testing.tuning.get_heuristics_config", force_umma) - - -def _case(name, gemm_type, **weight_values): - return KernelTestCase( - name=name, - layer_config=LayerConfig( - shape_n=256, - shape_k=256, - num_experts=0 if gemm_type == GemmType.DENSE else 4, - a_dtype=dtypes.bfloat16, - c_dtype=dtypes.bfloat16, - **(dict(bs_dtype="bfloat16") | weight_values), - ), - compute_config=ComputeConfig(gemm_type=gemm_type), - top_k=2, - seed=2026, - ) - - -def _assert_results(case, shape_ms): - runner = KernelTestRunner(case) - kernels = runner.prepare_kernels(shape_ms) - for variants in kernels.values(): - for kernel in variants: - compiled = HummingKernel._id2kernel[int(kernel[0][2])] - assert compiled.mma_type == MmaType.UMMA - dequant_threads = 0 if compiled.use_umma_ss else 128 * compiled.umma_num_dequant_warpgroups - expected_threads = 256 + dequant_threads - assert compiled.num_threads == expected_threads - assert compiled.num_math_threads == 128 - assert compiled.num_load_threads in (64, 96) - compiled.assert_smem_size_matches_estimate() - results = runner.run(shape_ms) - assert {result.shape_m for result in results} == set(shape_ms) - for result in results: - torch.testing.assert_close(result.outputs, result.outputs_ref, rtol=case.rtol, atol=case.atol) - - -@pytest.mark.parametrize("block_n,block_k", ((256, 64), (128, 128))) -def test_umma_operand_wait_once_per_iteration(block_n, block_k, monkeypatch): - # Odd stages must not wait again for a second output group or fragment. - def select_config(layer_config, shape_m, gemm_type, **kwargs): - return Sm100Heuristics.get_umma_config(layer_config, shape_m, gemm_type) | { - "block_shape": (24, block_n, block_k), - "warp_shape": (24, 32, block_k), - "num_stages": 3, - "num_ctas_per_sm": 1, - "use_stream_k": False, - "use_tma": True, - "use_tma_a": True, - "use_tma_c": True, - } - - monkeypatch.setattr("humming.testing.tuning.get_heuristics_config", select_config) - case = _case("operand-wait", GemmType.DENSE, **WEIGHT_CONFIGS["nvfp4"]) - case = dataclasses.replace(case, layer_config=dataclasses.replace(case.layer_config, shape_k=1024)) - _assert_results(case, (64, 257)) - - -@pytest.mark.parametrize("shape_n,shape_k,num_sms", ((256, 256, 3), (128, 8192, 64))) -@pytest.mark.parametrize("use_tma_c", (False, True)) -@pytest.mark.parametrize( - "weight_values", - ( - dict(b_dtype="uint4", weight_scale_group_size=128, has_bias=True), - dict(b_dtype="uint4", weight_scale_type="channel", has_bias=True), - ), -) -def test_umma_native_output_stream_k_bias(weight_values, use_tma_c, shape_n, shape_k, num_sms, monkeypatch): - """Only the first K slice contributes bias to native output.""" - case = _case("native-output-stream-k", GemmType.DENSE, **weight_values) - - case = dataclasses.replace( - case, layer_config=dataclasses.replace(case.layer_config, shape_n=shape_n, shape_k=shape_k) - ) - - def select_stream_k(layer_config, shape_m, gemm_type, **kwargs): - return Sm100Heuristics.get_umma_config(layer_config, shape_m, gemm_type) | { - "block_shape": (128, 128, 64), - "warp_shape": (128, 32, 64), - "num_stages": 2, - "num_ctas_per_sm": 1, - "num_sms": num_sms, - "use_tma": True, - "use_tma_c": use_tma_c, - "use_stream_k": True, - } - - monkeypatch.setattr("humming.testing.tuning.get_heuristics_config", select_stream_k) - _assert_results(case, (17, 257)) - - -@pytest.mark.parametrize( - "gemm_type,block_m,shape_k,block_n,weight_name", - ( - (GemmType.DENSE, 64, 64, 128, "uint4"), - (GemmType.GROUPED_CONTIGUOUS, 64, 192, 256, "uint4-zp"), - (GemmType.INDEXED, 176, 256, 128, "uint4"), - ), -) -def test_umma_pipeline_stage_reuse(gemm_type, block_m, shape_k, block_n, weight_name, monkeypatch): - """Retire async reads before reuse across persistent tiles and experts.""" - weights = WEIGHT_CONFIGS[weight_name] | {"weight_scale_group_size": 64} - case = _case("pipeline-stage-reuse", gemm_type, **weights) - case = dataclasses.replace(case, layer_config=dataclasses.replace(case.layer_config, shape_k=shape_k)) - - def minimum_stages(layer_config, shape_m, gemm_type, **kwargs): - return Sm100Heuristics.get_umma_config(layer_config, shape_m, gemm_type) | { - "block_shape": (block_m, block_n, 64), - "warp_shape": (block_m, 32, 64), - "num_stages": 3, - "num_sms": 2, - } - - monkeypatch.setattr("humming.testing.tuning.get_heuristics_config", minimum_stages) - _assert_results(case, (17, 257)) - - -def _public_problem(weight_ref, shape_m, gemm_type, block_m): - device = weight_ref.device - shape_k = weight_ref.shape[-1] - if gemm_type == GemmType.DENSE: - inputs = generate_random_tensor((shape_m, shape_k), torch.bfloat16, device=device) - return dict(inputs=inputs), slice(None), inputs.float() @ weight_ref.T - - # Leave experts 1 and 3 empty; experts 0 and 2 have tail tiles. - topk_ids = torch.tensor([0, 2], device=device, dtype=torch.int32) - topk_ids = topk_ids.expand(shape_m, -1).contiguous() - expert_max_tokens = shape_m + 3 - _, layout, sorted_ids, expert_ids, padded = generate_moe_tensors( - topk_ids, - 4, - gemm_type, - block_size_config=block_m, - expert_max_tokens=expert_max_tokens, - ) - if gemm_type == GemmType.INDEXED: - inputs = generate_random_tensor((shape_m, shape_k), torch.bfloat16, device=device) - reference = torch.stack([inputs.float() @ weight_ref[e].T for e in (0, 2)], dim=1).flatten(0, 1) - return ( - dict( - inputs=inputs, - sorted_ids=sorted_ids, - expert_ids=expert_ids, - num_tokens_padded=padded, - top_k=2, - ), - slice(None), - reference, - ) - - total_m = shape_m * 2 if gemm_type == GemmType.GROUPED_CONTIGUOUS else 4 * expert_max_tokens - inputs = generate_random_tensor((total_m, shape_k), torch.bfloat16, device=device) - output_ids, references = [], [] - for expert in (0, 2): - offset = ( - int(layout[expert]) if gemm_type == GemmType.GROUPED_CONTIGUOUS else expert * expert_max_tokens - ) - ids = torch.arange(offset, offset + shape_m, device=device) - output_ids.append(ids) - references.append(inputs[ids].float() @ weight_ref[expert].T) - return ( - dict(inputs=inputs, expert_layout=layout, valid_shape_m=shape_m * 2), - torch.cat(output_ids), - torch.cat(references), - ) - - -def _public_layer(weight_name, gemm_type, shape_n=256, shape_k=256): - torch.manual_seed(2026) - schema = HummingWeightSchema(**WEIGHT_CONFIGS[weight_name]) - num_experts = 0 if gemm_type == GemmType.DENSE else 4 - weight_shape = (4, shape_n, shape_k) if num_experts else (shape_n, shape_k) - weight = generate_random_tensor(weight_shape, torch.bfloat16, device="cuda") - tensors = schema.quant_tensor(weight, schema, torch.bfloat16) - if num_experts and "weight_scale_2" in tensors: - tensors["weight_scale_2"] *= torch.arange(1, num_experts + 1, device=weight.device).reshape_as( - tensors["weight_scale_2"] - ) - weight_ref = schema.dequant_tensors(tensors) - layer = HummingLayer( - shape_n=shape_n, - shape_k=shape_k, - num_experts=num_experts, - weight_config=schema, - input_config={"dtype": "bfloat16"}, - torch_dtype=torch.bfloat16, - ).cuda() - layer.load_state_dict(tensors, strict=False) - layer.transform() - return layer, weight_ref - - -def _selected_backend(layer, gemm_type, kwargs, tuning=None): - prepared = HummingKernel.prepare_kernels( - layer.humming_config.to_str(), - {"gemm_type": gemm_type.value}, - tuning, - ).reshape(-1, 4) - dispatch_m = kwargs.get("valid_shape_m", 0) or kwargs["inputs"].shape[0] * ( - kwargs.get("top_k", 1) if gemm_type == GemmType.INDEXED else 1 - ) - selected = [row for row in prepared if int(row[0]) < dispatch_m <= int(row[1])] - assert len(selected) == 1 - return HummingKernel._id2kernel[int(selected[0][2])].mma_type - - -@pytest.mark.parametrize( - "gemm_type,weight_name", - ( - (GemmType.DENSE, "uint4"), - (GemmType.INDEXED, "uint4-zp"), - (GemmType.GROUPED_CONTIGUOUS, "uint4-fp-zp"), - (GemmType.GROUPED_MASKED, "nvfp4"), - (GemmType.DENSE, "mxfp4"), - (GemmType.INDEXED, "fp8"), - ), -) -def test_umma_public_layer_switches_without_repacking(weight_name, gemm_type): - """MMA and UMMA consume one transformed layer, including routed calls.""" - layer, weight_ref = _public_layer(weight_name, gemm_type) - num_experts = layer.num_experts - packed = { - name: (value.data_ptr(), value.detach().view(torch.uint8).clone()) - for name, value in layer.named_parameters() - } - for shape_m, mma_type in ( - (17, MmaType.MMA), - (257, MmaType.UMMA), - (257, MmaType.MMA), - (17, MmaType.MMA), - ): - torch.manual_seed(2026 + shape_m) - config = layer.humming_config - get_config = Sm100MmaHeuristics.get_config - if mma_type == MmaType.UMMA: - get_config = Sm100Heuristics.get_umma_config - tuning = get_config( - config, - shape_m=shape_m * (2 if num_experts else 1), - gemm_type=gemm_type, - ) - tuning |= {"mma_type": mma_type.value} - kwargs, output_ids, reference = _public_problem( - weight_ref, shape_m, gemm_type, tuning["block_shape"][0] - ) - outputs = layer( - **kwargs, - compute_config={"gemm_type": gemm_type.value}, - tuning_config=tuning, - ) - actual_mma = _selected_backend(layer, gemm_type, kwargs, tuning) - assert actual_mma == mma_type - torch.testing.assert_close(outputs[output_ids], reference.to(torch.bfloat16), rtol=0.01, atol=0.05) - for name, value in layer.named_parameters(): - pointer, original = packed[name] - assert pointer == value.data_ptr() - assert torch.equal(value.detach().view(torch.uint8), original) - - -@pytest.mark.parametrize( - "gemm_type,use_stream_k", - ( - (GemmType.DENSE, False), - (GemmType.INDEXED, True), - ), -) -def test_umma_k32_tile_reuse(gemm_type, use_stream_k, monkeypatch): - def select_k32(layer_config, shape_m, gemm_type, **kwargs): - return Sm100Heuristics.get_umma_config(layer_config, shape_m, gemm_type) | { - "block_shape": (8, 256, 32), - "warp_shape": (8, 32, 32), - "num_stages": 5, - "num_sms": 2, - "num_ctas_per_sm": 1, - "use_stream_k": use_stream_k, - } - - monkeypatch.setattr("humming.testing.tuning.get_heuristics_config", select_k32) - case = _case("uint4", gemm_type, **WEIGHT_CONFIGS["uint4"]) - case = dataclasses.replace(case, layer_config=dataclasses.replace(case.layer_config, shape_k=1024)) - _assert_results(case, (1, 17, 257)) - - -def test_umma_fp8_persistent_weight_reuse(monkeypatch): - def select_config(layer_config, shape_m, gemm_type, **kwargs): - return Sm100Heuristics.get_umma_config(layer_config, shape_m, gemm_type) | { - "block_shape": (256, 128, 128), - "warp_shape": (256, 32, 128), - "num_stages": 3, - "num_ctas_per_sm": 1, - "num_sms": 2, - "use_stream_k": False, - } - - monkeypatch.setattr("humming.testing.tuning.get_heuristics_config", select_config) - case = _case("fp8-persistent", GemmType.DENSE, b_dtype=dtypes.float8e4m3, weight_scale_type="tensor") - layer_config = dataclasses.replace( - case.layer_config, - a_dtype=dtypes.float8e4m3, - shape_k=8192, - input_quant_mode="static_tensor", - ) - _assert_results(dataclasses.replace(case, layer_config=layer_config), (257, 4096)) - - -@pytest.mark.parametrize( - "shape_k,activation_dtype,weight_dtype", - ( - (64, "float8e4m3", "float8e4m3"), - (256, "float8e3m4", "float8e3m4"), - (64, "float8e4m3", "float4e2m1"), - (256, "float8e3m4", "float6e2m3"), - ), -) -def test_umma_fp8_public_dispatch(shape_k, activation_dtype, weight_dtype, monkeypatch): - # This regression checks production small-M dispatch, not test backend preference. - monkeypatch.delenv("HUMMING_TEST_TUNING_SOURCE", raising=False) - monkeypatch.delenv("PYTEST_CURRENT_TEST", raising=False) - schema = HummingWeightSchema(b_dtype=weight_dtype, weight_scale_type="tensor") - weight = generate_random_tensor((256, shape_k), torch.bfloat16, device="cuda") - tensors = schema.quant_tensor(weight, schema, torch.bfloat16) - weight_ref = schema.dequant_tensors(tensors).float() - tensors["input_scale"] = torch.tensor([0.25], device="cuda") - layer = HummingLayer( - shape_n=256, - shape_k=shape_k, - weight_config=schema, - input_config={"dtype": activation_dtype, "quant_mode": "static_tensor"}, - torch_dtype=torch.bfloat16, - ).cuda() - layer.load_state_dict(tensors, strict=False) - layer.transform() - assert get_default_mma_type(layer.humming_config) == MmaType.UMMA - assert layer.humming_config.use_raw_weight - runner = KernelTestRunner( - KernelTestCase( - name="fp8-public", - layer_config=layer.humming_config, - compute_config=ComputeConfig(gemm_type=GemmType.DENSE), - ) - ) - for shape_m in (17, 257): - inputs = torch.randn((shape_m, shape_k), device="cuda", dtype=torch.bfloat16) - inputs_ref, _, _, _ = runner.prepare_inputs(inputs, tensors["input_scale"]) - outputs = layer(inputs) - expected = inputs_ref @ weight_ref.T - torch.testing.assert_close(outputs, expected.to(torch.bfloat16), rtol=0.01, atol=0.05) - backend = _selected_backend(layer, GemmType.DENSE, {"inputs": inputs}) - expect_umma = shape_m == 257 or weight_dtype != "float8e4m3" or activation_dtype == "float8e3m4" - assert backend == (MmaType.UMMA if expect_umma else MmaType.MMA) - - -@pytest.mark.parametrize( - "block_n,block_k,weight_name,gemm_type,num_stages,use_tma,num_ctas", - ( - (128, 128, "fp8", GemmType.DENSE, 3, True, 1), - (512, 64, "uint4", GemmType.DENSE, 4, True, 1), - ), -) -def test_umma_cooperative_dequant( - block_n, block_k, weight_name, gemm_type, num_stages, use_tma, num_ctas, monkeypatch -): - def select_config(layer_config, shape_m, gemm_type, **kwargs): - return Sm100Heuristics.get_umma_config(layer_config, shape_m, gemm_type) | { - "block_shape": (32, block_n, block_k), - "warp_shape": (32, 32, block_k), - "num_stages": num_stages, - "num_ctas_per_sm": num_ctas, - "umma_num_dequant_warpgroups": 2, - "num_sms": 2, - "use_stream_k": True, - "use_tma": use_tma, - "use_tma_a": use_tma, - "use_tma_c": use_tma, - } - - monkeypatch.setattr("humming.testing.tuning.get_heuristics_config", select_config) - case = _case("cooperative-dequant", gemm_type, **WEIGHT_CONFIGS[weight_name]) - layer_config = dataclasses.replace(case.layer_config, shape_n=1024, shape_k=1024) - if weight_name == "fp8": - layer_config = dataclasses.replace( - layer_config, a_dtype=dtypes.float8e4m3, input_quant_mode="dynamic_token" - ) - _assert_results(dataclasses.replace(case, layer_config=layer_config), (17, 257)) - - -@pytest.mark.parametrize( - "weight_dtype,activation_dtype,group_size,scale_dtype,quant_mode", - ( - ("float4e2m1", "float8e3m4", 32, "float8e8m0", "dynamic_group"), - ("uint2", "float4e0m3", 16, "float8e4m3", "dynamic_group_token"), - ("float4e2m1", "float8e4m3", 32, "float8e8m0", "dynamic_group"), - ("uint2", "float4e2m1", 32, "float8e8m0", "dynamic_group"), - ("float4e2m1", "float4e2m1", 16, "float8e4m3", "dynamic_group_token"), - ("uint2", "float4e2m1", 16, "float8e4m3", "static_tensor_dynamic_group"), - ("float4e2m1", "float8e4m3", 32, "float8e8m0", "static_tensor_dynamic_group"), - ), -) -def test_umma_mxf8_mxf4_public_dispatch(activation_dtype, group_size, scale_dtype, quant_mode, weight_dtype): - schema = HummingWeightSchema( - b_dtype=weight_dtype, bs_dtype=scale_dtype, weight_scale_group_size=group_size - ) - weight = generate_random_tensor((256, 256), torch.bfloat16, device="cuda") - tensors = schema.quant_tensor(weight, schema, torch.bfloat16, allow_negative_scale=False) - weight_ref = schema.dequant_tensors(tensors).float() - static_scale = None - if quant_mode == "static_tensor_dynamic_group": - static_scale = torch.tensor([0.25], device="cuda") - tensors["input_scale_2"] = static_scale - layer = HummingLayer( - shape_n=256, - shape_k=256, - weight_config=schema, - input_config={ - "dtype": activation_dtype, - "group_size": group_size, - "scale_dtype": scale_dtype, - "quant_mode": quant_mode, - }, - torch_dtype=torch.bfloat16, - ).cuda() - layer.load_state_dict(tensors, strict=False) - layer.transform() - assert layer.humming_config.use_block_scaled_mma - assert get_default_mma_type(layer.humming_config) == MmaType.UMMA - runner = KernelTestRunner( - KernelTestCase( - name="mx-public", - layer_config=layer.humming_config, - compute_config=ComputeConfig(gemm_type=GemmType.DENSE), - ) - ) - for shape_m in (1, 17, 257): - inputs = torch.randn((shape_m, 256), device="cuda", dtype=torch.bfloat16) - inputs_ref, _, _, _ = runner.prepare_inputs(inputs, static_scale) - outputs = layer(inputs) - expected = inputs_ref @ weight_ref.T - torch.testing.assert_close(outputs, expected.to(torch.bfloat16), rtol=0.01, atol=0.05) - - -@pytest.mark.parametrize( - "dtype,weight_dtype,block_m,gemm_type,shape_k", - ( - (dtypes.float8e3m4, dtypes.float8e3m4, 64, GemmType.DENSE, 512), - (dtypes.float4e2m1, dtypes.float4e2m1, 16, GemmType.DENSE, 512), - (dtypes.float4e0m3, dtypes.float4e0m3, None, GemmType.GROUPED_MASKED, 1024), - ), -) -def test_umma_ss_small_tile(dtype, weight_dtype, block_m, gemm_type, shape_k, monkeypatch): - block_k = 128 if dtype.num_bits == 4 else 64 - - def select_config(layer_config, shape_m, gemm_type, **kwargs): - if block_m is None: - return Sm100UmmaHeuristics.get_config(layer_config, shape_m, gemm_type=gemm_type) - return dict( - mma_type="umma", - block_shape=(block_m, 128, block_k), - warp_shape=(block_m, 32, block_k), - num_stages=3, - num_sms=4, - use_tma=True, - use_tma_a=gemm_type != GemmType.INDEXED, - use_tma_c=gemm_type != GemmType.INDEXED, - use_stream_k=True, - smem_reuse_mode="none", - ) - - monkeypatch.setattr("humming.testing.tuning.get_heuristics_config", select_config) - is_fp4 = dtype.num_bits == 4 - config = LayerConfig( - shape_n=256, - shape_k=shape_k, - num_experts=0 if gemm_type == GemmType.DENSE else 4, - a_dtype=dtype, - b_dtype=weight_dtype, - c_dtype=dtypes.bfloat16, - as_dtype=dtypes.float8e4m3 if is_fp4 else None, - bs_dtype=dtypes.float8e4m3 if is_fp4 else dtypes.bfloat16, - input_quant_mode="dynamic_group_token" if is_fp4 else None, - input_scale_group_size=16 if is_fp4 else 0, - weight_scale_group_size=16 if is_fp4 else 0, - weight_scale_type="group" if is_fp4 else "channel", - ) - case = KernelTestCase( - name="ss-small-tile", - layer_config=config, - compute_config=ComputeConfig( - gemm_type=gemm_type, - use_m_major_input_scale=block_m is None and gemm_type != GemmType.INDEXED, - ), - top_k=2, - seed=2026, - ) - _assert_results(case, (1, 137, 833) if block_m is None else (1, 137)) - - -@pytest.mark.parametrize( - "gemm_type,block_m,block_n,block_k,use_tma_c,reuse_mode,cta_group_size,use_tma,use_stream_k", - ( - (GemmType.DENSE, 40, 128, 64, False, "all_stages", 1, True, True), - (GemmType.DENSE, 48, 128, 64, True, "last_stage", 2, True, True), - (GemmType.GROUPED_MASKED, 40, 128, 64, True, "all_stages", 1, True, True), - ), -) -def test_umma_chunked_output_layout( - gemm_type, - block_m, - block_n, - block_k, - use_tma_c, - reuse_mode, - cta_group_size, - use_tma, - use_stream_k, - monkeypatch, -): - - def select_output(layer_config, shape_m, gemm_type, **kwargs): - return { - "mma_type": "umma", - "block_shape": (block_m, block_n, block_k), - "warp_shape": (block_m, 32, block_k), - "num_stages": 3, - "num_sms": 6, - "num_ctas_per_sm": 1, - "use_tma": use_tma or use_tma_c, - "use_tma_b": use_tma, - "use_tma_bs": use_tma, - "use_tma_bzp": use_tma, - "use_tma_a": use_tma and gemm_type != GemmType.INDEXED, - "use_tma_c": use_tma_c, - "use_stream_k": use_stream_k, - "smem_reuse_mode": reuse_mode, - "umma_cta_group_size": cta_group_size, - "output_chunk_rows": 32, - } - - monkeypatch.setattr("humming.testing.tuning.get_heuristics_config", select_output) - case = _case( - "chunked-layout", - gemm_type, - b_dtype="uint4", - weight_scale_group_size=0, - has_zero_point=True, - has_bias=True, - ) - layer_config = dataclasses.replace(case.layer_config, shape_n=1024, shape_k=1024) - _assert_results(dataclasses.replace(case, layer_config=layer_config), (17, 13 * block_m + 1)) - - -@pytest.mark.parametrize( - "a_dtype,b_dtype,microscale,gemm_type,block_m,block_n,block_k,stream_k", - ( - (dtypes.float8e4m3, dtypes.float4e2m1, True, GemmType.DENSE, 160, 128, 256, False), - (dtypes.float8e4m3, dtypes.float4e2m1, True, GemmType.DENSE, 32, 512, 128, True), - (dtypes.float8e4m3, dtypes.float4e2m1, True, GemmType.GROUPED_MASKED, 48, 128, 128, True), - ), -) -def test_umma_cooperative_fp8( - a_dtype, - b_dtype, - microscale, - gemm_type, - block_m, - block_n, - block_k, - stream_k, - monkeypatch, -): - - def select_config(layer_config, shape_m, gemm_type, **kwargs): - return { - "mma_type": "umma", - "block_shape": (block_m, block_n, block_k), - "warp_shape": (block_m, 32, block_k), - "num_stages": 2 if block_n == 512 else 3, - "num_ctas_per_sm": 1, - "num_sms": 4, - "use_warp_spec": True, - "use_tma": True, - "use_tma_a": gemm_type != GemmType.INDEXED, - "use_tma_c": gemm_type != GemmType.INDEXED, - "use_tma_b": True, - "use_tma_as": microscale and gemm_type != GemmType.INDEXED, - "use_stream_k": stream_k, - "smem_reuse_mode": "none", - "umma_cta_group_size": 2, - "output_chunk_rows": 32, - } - - monkeypatch.setattr("humming.testing.tuning.get_heuristics_config", select_config) - output_dtype = dtypes.float16 if a_dtype == dtypes.float8e5m2 else dtypes.bfloat16 - layer_config = LayerConfig( - shape_n=2 * block_n, - shape_k=1024, - num_experts=0 if gemm_type == GemmType.DENSE else 4, - a_dtype=a_dtype, - b_dtype=b_dtype, - c_dtype=output_dtype, - bs_dtype=dtypes.float8e8m0 if microscale else output_dtype, - as_dtype=dtypes.float8e8m0 if microscale else None, - input_scale_group_size=32 if microscale else 0, - weight_scale_group_size=32 if microscale else 0, - input_quant_mode="dynamic_group" if microscale else "dynamic_token", - weight_scale_type="group" if microscale else "channel", - has_bias=not microscale, - ) - case = KernelTestCase( - name="cooperative-fp8", - layer_config=layer_config, - compute_config=ComputeConfig( - gemm_type=gemm_type, - use_m_major_input_scale=microscale and gemm_type != GemmType.INDEXED, - ), - top_k=2, - seed=2026, - ) - _assert_results(case, (17, 13 * block_m + 1)) - - -@pytest.mark.parametrize( - "a_dtype,b_dtype,group_size,scale_dtype,quant_mode,cta_group_size,gemm_type,use_tma,block_shape", - ( - ( - dtypes.float4e2m1, - dtypes.float4e2m1, - 32, - dtypes.float8e8m0, - "dynamic_group", - 1, - GemmType.DENSE, - True, - (48, 256, 256), - ), - ( - dtypes.float4e2m1, - dtypes.float4e2m1, - 16, - dtypes.float8e4m3, - "dynamic_group_token", - 2, - GemmType.INDEXED, - False, - (128, 128, 128), - ), - ( - dtypes.float4e0m3, - dtypes.float4e0m3, - 16, - dtypes.float8e8m0, - "static_tensor_dynamic_group", - 2, - GemmType.DENSE, - True, - (240, 128, 256), - ), - ), -) -def test_umma_fp4_activation( - a_dtype, - b_dtype, - group_size, - scale_dtype, - quant_mode, - cta_group_size, - gemm_type, - use_tma, - block_shape, - monkeypatch, -): - - def select_config(layer_config, shape_m, gemm_type, **kwargs): - return dict( - mma_type="umma", - block_shape=block_shape, - warp_shape=(block_shape[0], 32, block_shape[2]), - num_stages=3, - num_ctas_per_sm=1, - num_sms=4, - use_warp_spec=True, - use_tma=True, - use_tma_a=use_tma, - use_tma_c=use_tma, - use_tma_b=True, - use_tma_as=use_tma, - use_stream_k=True, - smem_reuse_mode="none", - umma_cta_group_size=cta_group_size, - output_chunk_rows=32, - ) - - monkeypatch.setattr("humming.testing.tuning.get_heuristics_config", select_config) - config = LayerConfig( - shape_n=2 * block_shape[1], - shape_k=1024, - num_experts=0 if gemm_type == GemmType.DENSE else 4, - a_dtype=a_dtype, - b_dtype=b_dtype, - c_dtype=dtypes.bfloat16, - as_dtype=scale_dtype, - bs_dtype=scale_dtype, - input_scale_group_size=group_size, - weight_scale_group_size=group_size, - input_quant_mode=quant_mode, - weight_scale_type="group", - weight_scale_2_type="tensor", - has_bias=True, - ) - case = KernelTestCase( - name="fp4-activation", - layer_config=config, - compute_config=ComputeConfig(gemm_type=gemm_type, use_m_major_input_scale=use_tma), - top_k=2, - seed=2026, - ) - _assert_results(case, (17, 833)) diff --git a/tests/kernels/humming/test_zero_point.py b/tests/kernels/humming/test_zero_point.py index 7e4cb5ed..48076a92 100644 --- a/tests/kernels/humming/test_zero_point.py +++ b/tests/kernels/humming/test_zero_point.py @@ -1,3 +1,5 @@ +import dataclasses + import pytest from humming import dtypes @@ -85,6 +87,19 @@ def _make_cases() -> list[KernelTestCase]: ZERO_POINT_CASES = _make_cases() +ZERO_POINT_CASES += [ + dataclasses.replace( + case, + name=f"{case.name}-channel", + layer_config=dataclasses.replace( + case.layer_config, weight_scale_group_size=0, weight_scale_type="channel" + ), + ) + for case in ZERO_POINT_CASES + if case.layer_config.a_dtype == dtypes.int8 + and case.layer_config.b_dtype in (dtypes.uint4, dtypes.uint7) + and case.layer_config.c_dtype == dtypes.bfloat16 +] @pytest.mark.parametrize("test_case", ZERO_POINT_CASES, ids=str) @@ -107,4 +122,4 @@ def test_zero_point_case_coverage(): assert {str(case.layer_config.a_dtype) for case in floating_cases} == {"float16", "bfloat16"} assert {case.layer_config.b_dtype.num_bits for case in floating_cases} == set(range(1, 9)) assert all(case.layer_config.has_zero_point for case in ZERO_POINT_CASES) - assert all(case.layer_config.weight_scale_group_size == 64 for case in ZERO_POINT_CASES) + assert {case.layer_config.weight_scale_group_size for case in ZERO_POINT_CASES} == {0, 64} diff --git a/tests/test_jit_infrastructure.py b/tests/test_jit_infrastructure.py index 881f8fa5..09d033f6 100644 --- a/tests/test_jit_infrastructure.py +++ b/tests/test_jit_infrastructure.py @@ -4,11 +4,9 @@ import pytest import torch -os.environ.setdefault("HUMMING_DISABLE_PARALLEL_BUILD", "1") - -import humming.jit.compiler as compiler_module # noqa: E402 -import humming.utils.jit as jit_utils # noqa: E402 -from humming.jit.compiler import Compiler # noqa: E402 +import humming.jit.compiler as compiler_module +import humming.utils.jit as jit_utils +from humming.jit.compiler import Compiler class _FakeCompiler(Compiler): @@ -107,10 +105,22 @@ def test_precompiled_manifest_uses_content_hashes(tmp_path, monkeypatch): @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires two CUDA devices") -def test_launcher_runs_device_local_kernels(): +@pytest.mark.parametrize("tuning_source", ("heuristic", "batch_invariant", "heuristic+batch_invariant")) +def test_launcher_runs_device_local_kernels(monkeypatch, tuning_source): from humming import dtypes from humming.config import ComputeConfig, GemmType, LayerConfig - from humming.testing import KernelTestCase, KernelTestRunner + from humming.kernel.humming import HummingKernel + from humming.testing import KernelTestCase, KernelTestRunner, assert_kernel_test_shape_coverage + + monkeypatch.setenv("HUMMING_TEST_TUNING_SOURCE", tuning_source) + compile_batches = [] + compile_many = HummingKernel.compile_many + + def record_compile_batch(kernel_specs, device): + compile_batches.append([config["use_batch_invariant"] for _, config in kernel_specs]) + return compile_many(kernel_specs, device) + + monkeypatch.setattr(HummingKernel, "compile_many", record_compile_batch) def make_case(): return KernelTestCase( @@ -130,14 +140,20 @@ def make_case(): for device_index in range(2): with torch.cuda.device(device_index): runner = KernelTestRunner(make_case()) - result = runner.run((1,))[0] + results = runner.run((1,)) torch.cuda.synchronize(device_index) - torch.testing.assert_close( - result.outputs, - result.outputs_ref, - rtol=runner.test_case.rtol, - atol=runner.test_case.atol, - ) + assert_kernel_test_shape_coverage(results, (1,)) + assert len(compile_batches) == device_index + 1 + expected_modes = [source == "batch_invariant" for source in tuning_source.split("+")] + assert compile_batches[-1] == expected_modes + assert not runner.compute_config.use_batch_invariant + for result in results: + torch.testing.assert_close( + result.outputs, + result.outputs_ref, + rtol=runner.test_case.rtol, + atol=runner.test_case.atol, + ) @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires two CUDA devices") @@ -195,3 +211,33 @@ def load_library(path): monkeypatch.setattr(compiler_module.ctypes, "CDLL", load_library) assert compiler_module.NVRTCCompiler.signature() == "nvrtc+13.2" assert loaded_paths == ["/toolkit/libnvrtc.so"] + + +@pytest.mark.parametrize("limit", (1, 3, 32)) +def test_compile_many_obeys_worker_limit(monkeypatch, limit): + from contextlib import nullcontext + + from humming.jit import runtime + + worker_counts = [] + + class RecordingExecutor: + def __init__(self, max_workers): + worker_counts.append(max_workers) + + def __enter__(self): + return self + + def __exit__(self, *args): + pass + + def map(self, function, specs): + return map(function, specs) + + monkeypatch.setattr(runtime, "get_parallel_build_workers", lambda: limit) + monkeypatch.setattr(runtime, "ThreadPoolExecutor", RecordingExecutor) + monkeypatch.setattr(runtime.torch.cuda, "device", lambda device: nullcontext()) + specs = [(lambda value: value, {"value": value}) for value in range(5)] + assert runtime.KernelRuntime.compile_many(specs, device=0) == list(range(5)) + assert worker_counts == ([] if limit == 1 else [min(limit, len(specs))]) + assert runtime.KernelRuntime.compile_many([], device=0) == [] diff --git a/tests/test_tuning_candidates.py b/tests/test_tuning_candidates.py index ff0bcb9b..16b4daca 100644 --- a/tests/test_tuning_candidates.py +++ b/tests/test_tuning_candidates.py @@ -336,22 +336,43 @@ def test_decision_rejects_an_illegal_selection(): @pytest.mark.parametrize( - "sm_version,a_dtype,b_dtype,group_size,scale_dtype,packed_k,use_f16_accum,expected", + "sm_version,a_dtype,b_dtype,group_size,scale_dtype,packed_k,use_f16_accum,has_zero_point,expected", ( - (90, "bfloat16", "uint4", 128, "bfloat16", False, False, {"mma", "wgmma"}), - (90, "float8e4m3", "uint4", 128, "bfloat16", True, False, {"wgmma"}), - (90, "int4", "int4", 0, "bfloat16", False, False, {"mma"}), - (103, "bfloat16", "uint4", 128, "bfloat16", False, False, {"mma", "umma"}), - (103, "float8e4m3", "float8e4m3", 0, "bfloat16", False, False, {"mma", "umma"}), - (103, "float8e4m3", "float4e2m1", 32, "float8e8m0", False, False, {"umma"}), - (100, "int8", "uint4", 0, "bfloat16", False, False, {"mma", "umma"}), - (103, "int8", "uint4", 0, "bfloat16", False, False, {"mma"}), - (103, "float16", "uint4", 128, "float16", False, True, {"mma"}), - (120, "float4e2m1", "float4e2m1", 32, "float8e8m0", False, False, {"mxmma"}), + (90, "bfloat16", "uint4", 128, "bfloat16", False, False, False, {"mma", "wgmma"}), + (90, "float8e4m3", "uint4", 128, "bfloat16", True, False, False, {"wgmma"}), + (90, "int4", "int4", 0, "bfloat16", False, False, False, {"mma"}), + (103, "bfloat16", "uint4", 128, "bfloat16", False, False, False, {"mma", "umma"}), + (103, "float8e4m3", "float8e4m3", 0, "bfloat16", False, False, False, {"mma", "umma"}), + (103, "float8e4m3", "float4e2m1", 32, "float8e8m0", False, False, False, {"umma"}), + (100, "float8e3m4", "float3e1m1", 0, "bfloat16", False, False, False, {"umma"}), + (103, "float8e3m4", "float3e2m0", 0, "bfloat16", False, False, False, {"umma"}), + (103, "float8e4m3", "float3e1m1", 0, "bfloat16", False, False, False, {"mma", "umma"}), + (100, "int8", "uint4", 0, "bfloat16", False, False, False, {"mma", "umma"}), + (100, "int8", "uint4", 0, "bfloat16", False, False, True, {"mma", "umma"}), + (110, "int8", "uint7", 0, "bfloat16", False, False, True, {"mma", "umma"}), + (103, "int8", "uint4", 0, "bfloat16", False, False, False, {"mma"}), + (103, "float16", "uint4", 128, "float16", False, True, False, {"mma"}), + (120, "float4e2m1", "float4e2m1", 32, "float8e8m0", False, False, False, {"mxmma"}), + (100, "float8e3m4", "uint5", 32, "float8e8m0", False, False, True, {"umma"}), + (103, "float8e4m3", "uint4", 32, "float8e8m0", False, False, True, {"umma"}), + (103, "float8e5m2", "uint3", 32, "float8e8m0", False, False, True, {"umma"}), + (103, "float8e4m3", "uint4", 0, "bfloat16", False, False, True, {"mma", "umma"}), + (100, "float4e2m1", "uint2", 16, "float8e4m3", False, False, True, {"umma"}), + (103, "float4e2m1", "uint2", 32, "float8e8m0", False, False, True, {"umma"}), + (120, "float4e0m3", "uint3", 16, "float8e4m3", False, False, True, {"mxmma"}), ), ) def test_sampled_backends_match_fixed_layout( - sm_version, a_dtype, b_dtype, group_size, scale_dtype, packed_k, use_f16_accum, expected, monkeypatch + sm_version, + a_dtype, + b_dtype, + group_size, + scale_dtype, + packed_k, + use_f16_accum, + has_zero_point, + expected, + monkeypatch, ): from humming.config import ComputeConfig from humming.device import DeviceInfo @@ -365,7 +386,6 @@ def test_sampled_backends_match_fixed_layout( monkeypatch.setattr(HummingKernel, "_instances", {}) monkeypatch.setattr(HummingKernel, "prepare", lambda self: None) monkeypatch.setattr(HummingKernel, "register_kernel", lambda self: None) - monkeypatch.setattr(tuning, "NUM_SAMPLED_TUNING_CONFIGS", 20) layer = LayerConfig( sm_version=sm_version, shape_n=512, @@ -376,12 +396,16 @@ def test_sampled_backends_match_fixed_layout( bs_dtype=scale_dtype, weight_scale_group_size=group_size, use_packed_k_layout=packed_k, + has_zero_point=has_zero_point, ) compute = ComputeConfig(gemm_type=GemmType.DENSE, use_f16_accum=use_f16_accum) configs = tuning.sample_test_tuning_configs(layer, compute, sample_size=20) + assert len(configs) == 20 assert {config["mma_type"] for config in configs} == expected assert "mma_type" not in layer.to_dict() for config in configs: + requires_cp_async = config["use_warp_spec"] or config["mma_type"] in ("wgmma", "umma") + assert not requires_cp_async or config["use_cp_async"] # Exercise the real kernel validators and code generation without compiling. tuning_values = tuning.create_tuning_config(config).to_dict() kernel = HummingKernel(**(layer.to_dict() | compute.to_dict() | tuning_values)) @@ -389,7 +413,8 @@ def test_sampled_backends_match_fixed_layout( assert kernel.use_raw_weight == layer.use_raw_weight -def test_sampled_umma_covers_cooperative_and_dequant_options(monkeypatch): +@pytest.mark.parametrize("sample_size", (None, 32)) +def test_sampled_umma_covers_cooperative_and_dequant_options(monkeypatch, sample_size): from humming.config import ComputeConfig from humming.device import DeviceInfo from humming.testing import tuning @@ -406,16 +431,21 @@ def test_sampled_umma_covers_cooperative_and_dequant_options(monkeypatch): ) compute = ComputeConfig(gemm_type=GemmType.INDEXED, use_batch_invariant=True) layer = dataclasses.replace(layer, num_experts=4) - configs = tuning.sample_test_tuning_configs(layer, compute) + sample_options = {} if sample_size is None else {"sample_size": sample_size} + configs = tuning.sample_test_tuning_configs(layer, compute, **sample_options) + assert tuning.NUM_SAMPLED_TUNING_CONFIGS == 100 + assert len(configs) == (sample_size or tuning.NUM_SAMPLED_TUNING_CONFIGS) umma_configs = [config for config in configs if config["mma_type"] == "umma"] for name in ("umma_cta_group_size", "umma_num_dequant_warpgroups", "output_chunk_rows"): assert {config[name] for config in umma_configs} == set(tuning.SAMPLED_TUNING_VALUES[name]) assert {config["use_tma_b"] for config in umma_configs} == {False, True} + assert any(config["block_shape"][2] == 32 for config in umma_configs) + assert {3, 5} <= {config["num_stages"] for config in umma_configs} for config in configs: assert not config["use_stream_k"] assert not config["use_tma_a"] and not config["use_tma_c"] assert config["block_shape"][2] == config["warp_shape"][2] - assert configs == tuning.sample_test_tuning_configs(layer, compute) + assert configs == tuning.sample_test_tuning_configs(layer, compute, **sample_options) @pytest.mark.parametrize("output_chunk_rows", (-32, 8, 16, 24, 48, 288)) @@ -432,18 +462,21 @@ def test_output_chunk_rows_rejects_invalid_heights(output_chunk_rows): @pytest.mark.parametrize("sm_version", (100, 103, 107, 110)) @pytest.mark.parametrize( - "a_dtype,b_dtype", + "a_dtype,b_dtype,small_m_backend", ( - (dtypes.bfloat16, dtypes.uint4), - (dtypes.int8, dtypes.int8), - (dtypes.float8e4m3, dtypes.float8e4m3), - (dtypes.float8e4m3, dtypes.float4e2m1), - (dtypes.float4e2m1, dtypes.float4e2m1), + (dtypes.bfloat16, dtypes.uint4, "mma"), + (dtypes.int8, dtypes.int8, "mma"), + (dtypes.float8e4m3, dtypes.float8e4m3, "mma"), + (dtypes.float8e4m3, dtypes.float4e2m1, "umma"), + (dtypes.float4e2m1, dtypes.float4e2m1, "umma"), + (dtypes.float8e3m4, dtypes.float8e3m4, "umma"), + (dtypes.float8e3m4, dtypes.float6e2m3, "umma"), ), ) -def test_umma_architecture_selection(sm_version, a_dtype, b_dtype): +def test_umma_architecture_selection(sm_version, a_dtype, b_dtype, small_m_backend, monkeypatch): from humming.config import MmaType from humming.config.mma import get_default_mma_type + from humming.tune.sm100 import Sm100Heuristics scale_config = {} if a_dtype == dtypes.float4e2m1: @@ -466,6 +499,12 @@ def test_umma_architecture_selection(sm_version, a_dtype, b_dtype): expected = MmaType.MMA if a_dtype == dtypes.int8 and sm_version in (103, 107) else MmaType.UMMA assert get_default_mma_type(config) == expected + monkeypatch.delenv("HUMMING_TEST_TUNING_SOURCE", raising=False) + monkeypatch.delenv("PYTEST_CURRENT_TEST", raising=False) + for shape_m, backend in ((17, small_m_backend), (257, expected.value)): + tuning = Sm100Heuristics.get_config(config, shape_m) + assert tuning["mma_type"] == backend + @pytest.mark.parametrize( "sm_version,a_dtype,shape_n,use_f16_accum,expected", @@ -478,13 +517,15 @@ def test_umma_architecture_selection(sm_version, a_dtype, b_dtype): (80, "bfloat16", 256, False, "mma"), ), ) +@pytest.mark.parametrize("tuning_source", ("heuristic", "heuristic+batch_invariant")) def test_heuristic_tests_prefer_available_backend( - sm_version, a_dtype, shape_n, use_f16_accum, expected, monkeypatch + sm_version, a_dtype, shape_n, use_f16_accum, expected, monkeypatch, tuning_source ): from humming.config import ComputeConfig from humming.device import DeviceInfo from humming.testing import tuning + monkeypatch.setenv("HUMMING_TEST_TUNING_SOURCE", tuning_source) monkeypatch.setattr(DeviceInfo, "sm_version", property(lambda self: sm_version)) monkeypatch.setattr(DeviceInfo, "is_ppu", property(lambda self: False)) layer = LayerConfig( @@ -501,63 +542,6 @@ def test_heuristic_tests_prefer_available_backend( assert {config["mma_type"] for config in configs} == {expected} -@pytest.mark.parametrize( - "warp_m,num_ctas,expected", - ( - (16, 3, True), - (32, 3, False), - (32, 2, True), - (64, 3, False), - (48, 4, False), - (64, 4, False), - (128, 1, False), - (128, 3, False), - ), -) -def test_warp_specialization_register_budget(warp_m, num_ctas, expected, monkeypatch): - from humming.config import ComputeConfig, TuningConfig - from humming.config.mma import get_register_budget_error - from humming.device import DeviceInfo - from humming.kernel.humming import HummingKernel - from humming.testing import tuning - - monkeypatch.setattr(DeviceInfo, "max_registers_per_sm", property(lambda self: 65536)) - - layer = dataclasses.replace(_layer(), sm_version=103) - compute = ComputeConfig(gemm_type=GemmType.INDEXED) - config = TuningConfig( - mma_type="mma", - block_shape=(warp_m, 128, 256), - warp_shape=(warp_m, 64, 64), - num_stages=2, - num_ctas_per_sm=num_ctas, - use_warp_spec=True, - use_tma=True, - use_tma_a=False, - use_tma_c=False, - use_stream_k=False, - output_chunk_rows=32, - smem_reuse_mode="last_stage", - ) - assert (get_register_budget_error(layer, config) is None) == expected - schedule_fields = {field.name for field in dataclasses.fields(ScheduleCandidate)} - schedule_config = {name: value for name, value in config.to_dict().items() if name in schedule_fields} - schedule = ScheduleCandidate.from_config("register-budget", schedule_config) - analysis = analyze_candidate(_problem(layer_config=layer), schedule) - has_register_error = any( - reason.startswith("register budget exceeded:") for reason in analysis.rejection_reasons - ) - assert has_register_error == (not expected) - if not expected: - assert not tuning._fits_device_resources(layer, compute, (config.to_dict(), {})) - monkeypatch.setattr(HummingKernel, "_instances", {}) - monkeypatch.setattr( - HummingKernel, "prepare", lambda self: pytest.fail("invalid config reached compilation") - ) - with pytest.raises(ValueError, match="register budget exceeded"): - HummingKernel(**(layer.to_dict() | compute.to_dict() | config.to_dict())) - - @pytest.mark.parametrize("a_dtype,expected", [("int4", "mma"), ("int8", "wgmma"), ("bfloat16", "wgmma")]) def test_sm90_default_backend_supports_activation(a_dtype, expected): from humming.config.mma import get_default_mma_type @@ -611,57 +595,6 @@ def test_sampled_wgmma_accounts_for_all_accumulators( assert tuning._fits_device_resources(layer, compute, (config, {})) == expected -@pytest.mark.parametrize( - "mma_type,raw_weight,group_size,packed,warp_shape,budget,expected", - [ - # MMA: 192 accumulator + 24 A + 16 B = 232; equality must reject. - ("mma", False, 0, False, (96, 64, 64), 240, False), - ("mma", False, 0, False, (96, 64, 64), 248, True), - # WGMMA excludes A in smem, and excludes B as well for raw-weight SS. - ("wgmma", False, 0, False, (96, 64, 64), 216, False), - ("wgmma", False, 0, False, (96, 64, 64), 224, True), - ("wgmma", True, 0, False, (96, 64, 64), 200, False), - ("wgmma", True, 0, False, (96, 64, 64), 208, True), - # FP8 group scale: 128 * 1.25 + 16 A + 16 B = 192, not 288. - ("mma", False, 128, False, (64, 64, 128), 200, False), - ("mma", False, 128, False, (64, 64, 128), 208, True), - ("wgmma", True, 128, False, (64, 64, 128), 168, False), - ("wgmma", True, 128, False, (64, 64, 128), 176, True), - # WGMMA packed B covers all K=128 slabs (16 registers for FP8). - ("wgmma", False, 128, True, (64, 16, 128), 64, False), - ("wgmma", False, 128, True, (64, 16, 128), 72, True), - ], -) -def test_sampled_operand_register_budget( - mma_type, raw_weight, group_size, packed, warp_shape, budget, expected -): - from humming.config import TuningConfig - from humming.config.mma import get_register_budget_error - - a_dtype = "float8e4m3" if group_size else "float16" - layer = LayerConfig( - sm_version=90, - shape_n=1024, - shape_k=1024, - a_dtype=a_dtype, - b_dtype=a_dtype if raw_weight else "uint4", - c_dtype="float16", - use_packed_k_layout=packed, - input_scale_group_size=group_size, - weight_scale_group_size=group_size, - use_fused_e8m0_scale=False, - ) - config = TuningConfig( - mma_type=mma_type, - warp_shape=warp_shape, - block_shape=(warp_shape[0], warp_shape[1] * 4, warp_shape[2]), - use_warp_spec=False, - num_ctas_per_sm=1, - ) - error = get_register_budget_error(layer, config, registers_per_sm=budget * 128) - assert (error is None) == expected - - @pytest.mark.parametrize("registers_per_sm,expected", [(65536, False), (131072, True)]) def test_device_register_budget(registers_per_sm, expected, monkeypatch): from humming.config import ComputeConfig, TuningConfig @@ -789,7 +722,7 @@ def test_mxmma_compiler_version_matches_scale_format( monkeypatch.setattr(KernelTestRunner, "prepare_weight", lambda self: None) monkeypatch.setenv("HUMMING_TEST_TUNING_SOURCE", "sampled") - def stop_before_sampling(*args): + def stop_before_sampling(*args, **kwargs): raise RuntimeError("reached tuning generation") monkeypatch.setattr(runner_module, "sample_test_tuning_configs", stop_before_sampling)