119 lines
4.8 KiB
Diff
119 lines
4.8 KiB
Diff
diff --git a/src/h3_blackwell_runtime/block.py b/src/h3_blackwell_runtime/block.py
|
|
index 6177e73..2c3650a 100644
|
|
--- a/src/h3_blackwell_runtime/block.py
|
|
+++ b/src/h3_blackwell_runtime/block.py
|
|
@@ -1,5 +1,7 @@
|
|
"""Direct MiniMax H3 DiT block over the standalone Sage3 attention unit."""
|
|
|
|
+import os
|
|
+
|
|
import torch
|
|
from torch import nn
|
|
from typing import TYPE_CHECKING
|
|
@@ -45,12 +47,15 @@ class H3SwiGLU(nn.Module):
|
|
self.fc2 = fc2
|
|
self.chunks = 1
|
|
self.chunk_threshold = 4096
|
|
+ self.fused_nvfp4_swiglu = os.getenv("H3_NVFP4_SWIGLU_FUSION", "").lower() in {
|
|
+ "1", "true", "yes", "on",
|
|
+ }
|
|
|
|
@classmethod
|
|
def from_checkpoint(cls, checkpoint: H3Checkpoint, prefix: str, *, output_dtype=torch.bfloat16):
|
|
return cls(
|
|
- checkpoint.nvfp4_linear(f"{prefix}.fc1", output_dtype=output_dtype),
|
|
- checkpoint.nvfp4_linear(f"{prefix}.fc2", output_dtype=output_dtype),
|
|
+ checkpoint.nvfp4_linear(f"{prefix}.fc1", output_dtype=output_dtype, role="h3_mlp_fc1"),
|
|
+ checkpoint.nvfp4_linear(f"{prefix}.fc2", output_dtype=output_dtype, role="h3_mlp_fc2"),
|
|
)
|
|
|
|
def forward(
|
|
@@ -71,7 +76,27 @@ class H3SwiGLU(nn.Module):
|
|
return self._forward_chunk(x)
|
|
|
|
def _forward_chunk(self, x: torch.Tensor) -> torch.Tensor:
|
|
- gate, up = self.fc1(x).chunk(2, dim=-1)
|
|
+ gate_up = self.fc1(x)
|
|
+ if self.fused_nvfp4_swiglu:
|
|
+ return self.fc2.forward_swiglu(gate_up)
|
|
+ gate, up = gate_up.chunk(2, dim=-1)
|
|
+ return self.fc2(torch.nn.functional.silu(gate).mul_(up))
|
|
+
|
|
+ def forward_modulated(
|
|
+ self,
|
|
+ x: torch.Tensor,
|
|
+ shift: torch.Tensor,
|
|
+ scale: torch.Tensor,
|
|
+ row_index: torch.Tensor,
|
|
+ ) -> torch.Tensor:
|
|
+ if self.chunks != 1:
|
|
+ from .h3_fusion import fused_modulate_
|
|
+
|
|
+ return self(fused_modulate_(x, shift, scale, row_index))
|
|
+ gate_up = self.fc1.forward_modulated(x, shift, scale, row_index)
|
|
+ if self.fused_nvfp4_swiglu:
|
|
+ return self.fc2.forward_swiglu(gate_up)
|
|
+ gate, up = gate_up.chunk(2, dim=-1)
|
|
return self.fc2(torch.nn.functional.silu(gate).mul_(up))
|
|
|
|
|
|
@@ -101,6 +126,12 @@ class H3DiTBlock(nn.Module):
|
|
self.attention = attention
|
|
self.mlp = mlp
|
|
self.norm_eps = norm_eps
|
|
+ self.fused_elementwise = os.getenv("H3_FUSED_ELEMENTWISE", "").lower() in {
|
|
+ "1", "true", "yes", "on",
|
|
+ }
|
|
+ self.fused_nvfp4_modulation = os.getenv("H3_NVFP4_MODULATE_FUSION", "").lower() in {
|
|
+ "1", "true", "yes", "on",
|
|
+ }
|
|
self.register_buffer("norm1_weight", norm1_weight, persistent=False)
|
|
self.register_buffer("norm2_weight", norm2_weight, persistent=False)
|
|
|
|
@@ -128,6 +159,46 @@ class H3DiTBlock(nn.Module):
|
|
sequence_parallel: "SequenceParallelContext | None" = None,
|
|
tensor_parallel: "SequenceParallelContext | None" = None,
|
|
) -> torch.Tensor:
|
|
+ if self.fused_elementwise and x.is_cuda and x.dtype == torch.bfloat16 and not x.requires_grad:
|
|
+ from .h3_fusion import fused_gate_add_, fused_modulate_, segment_index
|
|
+
|
|
+ row_index = segment_index(x.shape[0], segments, x.device)
|
|
+ if (
|
|
+ self.fused_nvfp4_modulation
|
|
+ and sequence_parallel is None
|
|
+ and tensor_parallel is None
|
|
+ ):
|
|
+ h = rms_norm(x, self.norm1_weight, self.norm_eps)
|
|
+ x = fused_gate_add_(
|
|
+ x,
|
|
+ self.attention(
|
|
+ h,
|
|
+ rope_rotation,
|
|
+ modulation=(shift_msa, scale_msa, row_index),
|
|
+ ),
|
|
+ gate_msa,
|
|
+ row_index,
|
|
+ )
|
|
+ h = rms_norm(x, self.norm2_weight, self.norm_eps)
|
|
+ return fused_gate_add_(
|
|
+ x,
|
|
+ self.mlp.forward_modulated(h, shift_mlp, scale_mlp, row_index),
|
|
+ gate_mlp,
|
|
+ row_index,
|
|
+ )
|
|
+ h = fused_modulate_(
|
|
+ rms_norm(x, self.norm1_weight, self.norm_eps), shift_msa, scale_msa, row_index,
|
|
+ )
|
|
+ x = fused_gate_add_(
|
|
+ x,
|
|
+ self.attention(h, rope_rotation, sequence_parallel, tensor_parallel),
|
|
+ gate_msa,
|
|
+ row_index,
|
|
+ )
|
|
+ h = fused_modulate_(
|
|
+ rms_norm(x, self.norm2_weight, self.norm_eps), shift_mlp, scale_mlp, row_index,
|
|
+ )
|
|
+ return fused_gate_add_(x, self.mlp(h, tensor_parallel), gate_mlp, row_index)
|
|
h = modulate_segments(rms_norm(x, self.norm1_weight, self.norm_eps), shift_msa, scale_msa, segments)
|
|
x = gate_segments(
|
|
x,
|