Skip to content

perf(mhc): plan source-split post/pre schedules #310

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
lukealonso wants to merge 2 commits into
base: master
Choose a base branch
from
Open

perf(mhc): plan source-split post/pre schedules #310

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
c66ad79
perf(mhc): plan source-split post/pre schedules
MadeBy561 Sep 2, 2026
1833b50
fix(policy): declare unqualified profile components
lukealonso Sep 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 5 additions & 3 deletions AGENTS.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -67,9 +67,11 @@ integrator-facing sequence and `b12x/policy/` for the contracts.
unconditional `leaf`, scalar `exact`, or disjoint inclusive `range` nodes,
each with an optional default. The tree determines dispatch coverage;
`coverage`, `evidence`, and `source_revision` are audit metadata only.
- The embedded registry is immutable and must contain exactly the component set
registered in the catalog. Provider IDs and schema versions must match their
runtime policies. Keep package-embedded profiles compact and validated.
- The embedded registry is immutable and must account for exactly the component
set registered in the catalog. A component is accounted for by either one
preplanned entry or an explicit `heuristic_components` declaration for an
unqualified device. Provider IDs and schema versions must match their runtime
policies. Keep package-embedded profiles compact and validated.
Generator checkpoints, full evidence artifacts, probes, and service A/B logs
are local working data unless a reviewed change explicitly requires them; do
not leave large untracked validation directories in the repository.
Expand Down
122 changes: 72 additions & 50 deletions b12x/norm/mhc/_impl.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1160,65 +1160,76 @@ def _b12x_mhc_post_pre_impl(
norm_eps=float(norm_eps),
)

prefill_min_tokens = int(
os.environ.get("B12X_MHC_PREFILL_MIN_TOKENS", "96")
)
if planned_config is not None:
tf32_enabled = planned_config.backend == "tf32_tma"
tf32_override = os.environ.get("B12X_MHC_PREFILL_TF32_MMA")
if tf32_override is None:
tf32_override = os.environ.get("B12X_MHC_PREFILL_BF16_MMA")
if tf32_override is not None:
tf32_enabled = tf32_override != "0"
use_prefill_tf32_mma = norm_weight is not None and tf32_enabled
use_prefill_tf32_mma = (
norm_weight is not None and planned_config.backend == "tf32_tma"
)
use_prefill_bf16_mma = False
use_prefill_block_m = (
norm_weight is not None
and planned_config.backend == "native"
and planned_config.native_post_pre_backend == "prefill_block_m"
)
use_prefill_compact = False
prefill_block_m_size = planned_config.prefill_block_m
prefill_tile_n = planned_config.prefill_tile_n
decode_source_splits = (
planned_config.decode_source_splits
if planned_config.backend == "native"
and planned_config.native_post_pre_backend == "decode"
else 0
)
else:
prefill_min_tokens = int(
os.environ.get("B12X_MHC_PREFILL_MIN_TOKENS", "96")
)
use_prefill_tf32_mma = _use_mhc_prefill_tf32_project(
norm_weight=norm_weight,
policy_m=policy_m,
)
use_prefill_bf16_mma = (
_use_mhc_prefill_bf16_project(
norm_weight=norm_weight,
policy_m=policy_m,
fn_bf16=fn_bf16,
use_prefill_bf16_mma = (
_use_mhc_prefill_bf16_project(
norm_weight=norm_weight,
policy_m=policy_m,
fn_bf16=fn_bf16,
)
and not use_prefill_tf32_mma
)
and not use_prefill_tf32_mma
)
use_prefill_block_m = (
not use_prefill_tf32_mma
and not use_prefill_bf16_mma
and norm_weight is not None
and policy_m >= prefill_min_tokens
and os.environ.get("B12X_MHC_PREFILL_BLOCK_M", "1") != "0"
)
prefill_block_m_size = int(
os.environ.get("B12X_MHC_PREFILL_BLOCK_M_SIZE", "2")
)
prefill_tile_n_default = 12 if hidden_size == 7168 else 24
prefill_tile_n = int(
os.environ.get(
"B12X_MHC_PREFILL_TILE_N", str(prefill_tile_n_default)
use_prefill_block_m = (
not use_prefill_tf32_mma
and not use_prefill_bf16_mma
and norm_weight is not None
and policy_m >= prefill_min_tokens
and os.environ.get("B12X_MHC_PREFILL_BLOCK_M", "1") != "0"
)
)
use_prefill_compact = (
not use_prefill_tf32_mma
and not use_prefill_bf16_mma
and not use_prefill_block_m
and norm_weight is not None
and policy_m >= prefill_min_tokens
and os.environ.get("B12X_MHC_PREFILL_COMPACT", "1") != "0"
)
decode_source_splits = 0
if not (
use_prefill_tf32_mma
or use_prefill_bf16_mma
or use_prefill_block_m
or use_prefill_compact
):
decode_source_splits, _ = _selected_post_pre_decode_split_n(
num_tokens=tokens,
hidden_size=hidden_size,
prefill_block_m_size = int(
os.environ.get("B12X_MHC_PREFILL_BLOCK_M_SIZE", "2")
)
prefill_tile_n_default = 12 if hidden_size == 7168 else 24
prefill_tile_n = int(
os.environ.get(
"B12X_MHC_PREFILL_TILE_N", str(prefill_tile_n_default)
)
)
use_prefill_compact = (
not use_prefill_tf32_mma
and not use_prefill_bf16_mma
and not use_prefill_block_m
and norm_weight is not None
and policy_m >= prefill_min_tokens
and os.environ.get("B12X_MHC_PREFILL_COMPACT", "1") != "0"
)
decode_source_splits = 0
if not (
use_prefill_tf32_mma
or use_prefill_bf16_mma
or use_prefill_block_m
or use_prefill_compact
):
decode_source_splits, _ = _selected_post_pre_decode_split_n(
num_tokens=tokens,
hidden_size=hidden_size,
)
if use_prefill_tf32_mma:
run_mhc_post_pre_prefill_gram(
x=x,
Expand Down Expand Up @@ -1282,6 +1293,7 @@ def _b12x_mhc_post_pre_impl(
partials=partials,
out=residual_out,
compute_gram=norm_weight is not None,
config=planned_config,
)
run_mhc_finalize_gram(
residual=residual_out,
Expand Down Expand Up @@ -1312,6 +1324,16 @@ def _b12x_mhc_post_pre_impl(
else 1
),
active_source_splits=decode_source_splits,
decode_finalize_threads=(
planned_config.decode_finalize_threads
if planned_config is not None
else None
),
decode_finalize_ctas=(
planned_config.decode_finalize_ctas
if planned_config is not None
else None
),
)
return residual_out, post_out, comb_out, y_out

Expand Down
Loading