From c7a9c0db4b7fe2523ee00aaeb27d9f2794b01620 Mon Sep 17 00:00:00 2001 From: Micuks <43397551+Micuks@users.noreply.github.com> Date: Tue, 28 Jul 2026 19:39:15 +0800 Subject: [PATCH 1/2] fix: use global head count for auto-CP planning --- .../ops/gated_delta_rule/chunk/cp_context.py | 21 ++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/flash_qla/ops/gated_delta_rule/chunk/cp_context.py b/flash_qla/ops/gated_delta_rule/chunk/cp_context.py index 5a265e6..363ee30 100644 --- a/flash_qla/ops/gated_delta_rule/chunk/cp_context.py +++ b/flash_qla/ops/gated_delta_rule/chunk/cp_context.py @@ -2,6 +2,7 @@ # Licensed under The MIT License [see LICENSE for details] import math +import os import torch import tilelang @@ -59,7 +60,25 @@ def _calc_cp_seqs( num_chunks = [tilelang.cdiv(x, chunk_size) for x in seqlens] # autocp - H = num_v_heads + planning_heads = os.environ.get("FLASH_QLA_AUTO_CP_PLANNING_HEADS") + if planning_heads is None: + H = num_v_heads + else: + try: + H = int(planning_heads) + except ValueError as exc: + raise ValueError( + "FLASH_QLA_AUTO_CP_PLANNING_HEADS must be an integer" + ) from exc + if H <= 0: + raise ValueError( + "FLASH_QLA_AUTO_CP_PLANNING_HEADS must be positive" + ) + if H < num_v_heads or H % num_v_heads != 0: + raise ValueError( + "FLASH_QLA_AUTO_CP_PLANNING_HEADS must be a multiple of " + f"local num_v_heads ({num_v_heads})" + ) # Latency model: T = a·L_cp + b·(B·H·Lc/P) / L_cp + c # Minimizing T yields the theoretical optimum: L_cp* ∝ √(B·H·Lc / P), where P = MULTI_PROCESSOR_COUNT, L_cp = max_local_chunks # Scaled by empirical factor (3) and aligned to the nearest power of 2 for optimal SM scheduling & memory alignment. From c2870abec74b8a3b133114bc6d9c55ad2d3a3962 Mon Sep 17 00:00:00 2001 From: Micuks <43397551+Micuks@users.noreply.github.com> Date: Wed, 29 Jul 2026 15:58:27 +0800 Subject: [PATCH 2/2] test: add standalone auto-CP regression --- tests/validate_global_head_autocp_planning.py | 194 ++++++++++++++++++ 1 file changed, 194 insertions(+) create mode 100644 tests/validate_global_head_autocp_planning.py diff --git a/tests/validate_global_head_autocp_planning.py b/tests/validate_global_head_autocp_planning.py new file mode 100644 index 0000000..909644d --- /dev/null +++ b/tests/validate_global_head_autocp_planning.py @@ -0,0 +1,194 @@ +#!/usr/bin/env python3 +"""Standalone FlashQLA regression for head-sharding-sensitive auto-CP planning. + +This validation intentionally has no model or CausalViT dependency. It compares +one 16-head FlashQLA launch with four independent 4-head launches using identical +per-head synthetic BF16 inputs and mixed variable-length sequence boundaries. + +Expected behavior: + +* upstream base 6ef4858, no override: + auto-CP chooses different segmentations and the outputs diverge; +* either revision with auto-CP disabled: + the 16-head and 4x4-head outputs are bit-exact; +* fixed revision with FLASH_QLA_AUTO_CP_PLANNING_HEADS=16: + auto-CP chooses the same segmentation and the outputs are bit-exact. +""" + +from __future__ import annotations + +import argparse +import json +import os +from pathlib import Path + +import torch + +from flash_qla import chunk_gated_delta_rule_fwd +from flash_qla.ops.gated_delta_rule.chunk.cp_context import _calc_cp_seqs +from flash_qla.utils import l2norm + + +LENGTHS = (988, 16128, 936, 16128) +NUM_HEADS = 16 +HEADS_PER_SHARD = 4 +HEAD_DIM = 128 +CHUNK_SIZE = 64 +SEED = 20260729 + + +def _stats(lhs: torch.Tensor, rhs: torch.Tensor) -> dict[str, float]: + lhs_f = lhs.float() + rhs_f = rhs.float() + delta = lhs_f - rhs_f + return { + "max_abs": delta.abs().max().item(), + "mean_abs": delta.abs().mean().item(), + "relative_l2": (delta.norm() / rhs_f.norm().clamp_min(1e-12)).item(), + "exact_fraction": (lhs == rhs).float().mean().item(), + } + + +def _make_inputs(device: torch.device) -> dict[str, torch.Tensor]: + torch.manual_seed(SEED) + total_tokens = sum(LENGTHS) + shape = (1, total_tokens, NUM_HEADS, HEAD_DIM) + head_shape = (1, total_tokens, NUM_HEADS) + return { + "q": l2norm(torch.randn(shape, device=device, dtype=torch.bfloat16)), + "k": l2norm(torch.randn(shape, device=device, dtype=torch.bfloat16)), + "v": torch.randn(shape, device=device, dtype=torch.bfloat16), + "g": torch.nn.functional.logsigmoid( + torch.randn(head_shape, device=device, dtype=torch.float32) + ) + / 16, + "beta": torch.randn( + head_shape, device=device, dtype=torch.float32 + ).sigmoid(), + } + + +def _run( + inputs: dict[str, torch.Tensor], + cu_seqlens: torch.Tensor, + *, + heads_per_call: int, + auto_cp: bool, +) -> torch.Tensor: + outputs = [] + for start in range(0, NUM_HEADS, heads_per_call): + head_slice = slice(start, start + heads_per_call) + _, _, output, _, _ = chunk_gated_delta_rule_fwd( + q=inputs["q"][..., head_slice, :].contiguous(), + k=inputs["k"][..., head_slice, :].contiguous(), + v=inputs["v"][..., head_slice, :].contiguous(), + g=inputs["g"][..., head_slice].contiguous(), + beta=inputs["beta"][..., head_slice].contiguous(), + scale=HEAD_DIM**-0.5, + initial_state=None, + cu_seqlens=cu_seqlens, + output_final_state=False, + output_h=False, + auto_cp=auto_cp, + ) + outputs.append(output) + return torch.cat(outputs, dim=-2) + + +def _planning(cu_seqlens: torch.Tensor, heads: int) -> dict[str, object]: + use_cp, cp_cu, *_ = _calc_cp_seqs(cu_seqlens, CHUNK_SIZE, heads) + return { + "use_cp": bool(use_cp), + "segment_count": int(cp_cu.numel() - 1) if cp_cu is not None else None, + "boundaries": cp_cu.tolist() if cp_cu is not None else None, + } + + +def _assert_exact(stats: dict[str, float], label: str) -> None: + assert stats["max_abs"] == 0.0, (label, stats) + assert stats["relative_l2"] == 0.0, (label, stats) + assert stats["exact_fraction"] == 1.0, (label, stats) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument( + "--expect", + required=True, + choices=("upstream-divergence", "fixed-exact"), + ) + parser.add_argument("--json-out", type=Path) + args = parser.parse_args() + + if not torch.cuda.is_available(): + raise RuntimeError("This validation requires an SM90 CUDA GPU") + + device = torch.device("cuda:0") + capability = torch.cuda.get_device_capability(device) + if capability != (9, 0): + raise RuntimeError(f"Expected SM90, found compute capability {capability}") + + lengths = torch.tensor(LENGTHS, device=device, dtype=torch.int32) + cu_seqlens = torch.nn.functional.pad( + lengths.cumsum(0, dtype=torch.int32), (1, 0) + ) + planning = { + "h16": _planning(cu_seqlens, NUM_HEADS), + "h4": _planning(cu_seqlens, HEADS_PER_SHARD), + } + + inputs = _make_inputs(device) + with torch.inference_mode(): + auto_cp_h16 = _run( + inputs, cu_seqlens, heads_per_call=NUM_HEADS, auto_cp=True + ) + auto_cp_4x4 = _run( + inputs, cu_seqlens, heads_per_call=HEADS_PER_SHARD, auto_cp=True + ) + no_auto_cp_h16 = _run( + inputs, cu_seqlens, heads_per_call=NUM_HEADS, auto_cp=False + ) + no_auto_cp_4x4 = _run( + inputs, cu_seqlens, heads_per_call=HEADS_PER_SHARD, auto_cp=False + ) + + auto_cp_stats = _stats(auto_cp_4x4, auto_cp_h16) + no_auto_cp_stats = _stats(no_auto_cp_4x4, no_auto_cp_h16) + _assert_exact(no_auto_cp_stats, "auto_cp_disabled") + + if args.expect == "upstream-divergence": + assert planning["h16"]["boundaries"] != planning["h4"]["boundaries"], planning + assert auto_cp_stats["max_abs"] > 0.0, auto_cp_stats + assert auto_cp_stats["relative_l2"] > 0.0, auto_cp_stats + assert auto_cp_stats["exact_fraction"] < 1.0, auto_cp_stats + else: + assert os.environ.get("FLASH_QLA_AUTO_CP_PLANNING_HEADS") == "16" + assert planning["h16"]["boundaries"] == planning["h4"]["boundaries"], planning + _assert_exact(auto_cp_stats, "fixed_auto_cp") + + result = { + "expectation": args.expect, + "flash_qla_file": __import__("flash_qla").__file__, + "planning_heads_env": os.environ.get( + "FLASH_QLA_AUTO_CP_PLANNING_HEADS" + ), + "torch_version": torch.__version__, + "device": torch.cuda.get_device_name(device), + "compute_capability": list(capability), + "seed": SEED, + "lengths": list(LENGTHS), + "shape": [1, sum(LENGTHS), NUM_HEADS, HEAD_DIM], + "planning": planning, + "auto_cp_h16_vs_4x4": auto_cp_stats, + "auto_cp_disabled_h16_vs_4x4": no_auto_cp_stats, + "status": "PASS", + } + rendered = json.dumps(result, indent=2, sort_keys=True) + print(rendered) + if args.json_out is not None: + args.json_out.parent.mkdir(parents=True, exist_ok=True) + args.json_out.write_text(rendered + "\n") + + +if __name__ == "__main__": + main()