Clean up: archive diagnostics and superseded tests
Kept: - example10 (CUTLASS LLM, O rescale + final normalize) - example9 (SSA kv_coord version) - working_softmax_maybe.py (working softmax snapshot from before the nuke) - test_fmha_v3_stage_c.py (identity softmax baseline, n=128 cos 0.999998) - test_fmha_v3.py (Stage A+B baseline) - layertest.py, cudagraph_test.py (required) - test_cutedsl.py, test_fp4_roundtrip.py (NVFP4 tests) Archived: diag_tma_*, example8, test_diag_multitile, test_reference_fmha, test_ref_minimal, test_tma_coord, test_fmha_v3_diag*, test_fmha_v3_12w, test_dense_router, test_interleave*, test_fused_step1, test_router, test_cache, test_compile_custom_op, test_custom_op, test_layer_schedule
This commit is contained in:
@@ -1,252 +0,0 @@
|
||||
"""Tests for KV cache: schema, allocator, pools, manager lifecycle."""
|
||||
|
||||
import torch
|
||||
import pytest
|
||||
|
||||
from dsv4.model.config import DSV4Config
|
||||
from dsv4.model.layer_schedule import build_schedule, AttentionType, RouterMode, LayerSpec
|
||||
from dsv4.cache.schema import build_schema, compute_block_budget, BLOCK_SIZE_ORIGINAL_TOKENS
|
||||
from dsv4.cache.allocator import BlockAllocator
|
||||
from dsv4.cache.paged_cache import PagedKVPool
|
||||
from dsv4.cache.state_cache import StateCachePool
|
||||
from dsv4.cache.manager import KVCacheManager
|
||||
|
||||
|
||||
# ---- Schema tests ----
|
||||
|
||||
def test_csa_schema():
|
||||
config = DSV4Config.pro()
|
||||
spec = LayerSpec(layer_idx=2, attn=AttentionType.CSA,
|
||||
ffn=__import__('dsv4.model.layer_schedule', fromlist=['FFNType']).FFNType.MOE,
|
||||
router_mode=RouterMode.HASH)
|
||||
schema = build_schema(config, spec)
|
||||
assert schema.entries_per_block == 32 # 128 / 4
|
||||
assert schema.indexer_entries_per_block == 32
|
||||
assert schema.tail_buffer_size == 3 # m - 1
|
||||
assert schema.swa_window_size == 128
|
||||
|
||||
|
||||
def test_hca_schema():
|
||||
config = DSV4Config.pro()
|
||||
spec = LayerSpec(layer_idx=3, attn=AttentionType.HCA,
|
||||
ffn=__import__('dsv4.model.layer_schedule', fromlist=['FFNType']).FFNType.MOE,
|
||||
router_mode=RouterMode.DENSE)
|
||||
schema = build_schema(config, spec)
|
||||
assert schema.entries_per_block == 1 # 128 / 128
|
||||
assert schema.indexer_entries_per_block == 0
|
||||
assert schema.tail_buffer_size == 127 # m' - 1
|
||||
|
||||
|
||||
def test_swa_schema():
|
||||
config = DSV4Config.flash()
|
||||
spec = LayerSpec(layer_idx=0, attn=AttentionType.SWA,
|
||||
ffn=__import__('dsv4.model.layer_schedule', fromlist=['FFNType']).FFNType.MOE,
|
||||
router_mode=RouterMode.HASH)
|
||||
schema = build_schema(config, spec)
|
||||
assert schema.entries_per_block == 0
|
||||
assert schema.indexer_entries_per_block == 0
|
||||
assert schema.tail_buffer_size == 0
|
||||
assert schema.swa_window_size == 128
|
||||
|
||||
|
||||
def test_schema_from_schedule():
|
||||
"""Every layer in a full schedule produces a valid schema."""
|
||||
config = DSV4Config.flash()
|
||||
schedule = build_schedule(config)
|
||||
for spec in schedule:
|
||||
schema = build_schema(config, spec)
|
||||
assert schema.swa_window_size > 0
|
||||
assert schema.entry_head_dim == config.head_dim
|
||||
if spec.attn == AttentionType.SWA:
|
||||
assert schema.entries_per_block == 0
|
||||
else:
|
||||
assert schema.entries_per_block > 0
|
||||
|
||||
|
||||
def test_block_budget():
|
||||
config = DSV4Config.pro()
|
||||
schedule = build_schedule(config)
|
||||
budget = compute_block_budget(config, schedule, 1_000_000, 16)
|
||||
assert "csa" in budget
|
||||
assert "hca" in budget
|
||||
assert budget["csa"] > budget["hca"] # CSA uses more blocks per request
|
||||
|
||||
|
||||
# ---- Allocator tests ----
|
||||
|
||||
def test_acquire_release_roundtrip():
|
||||
alloc = BlockAllocator(num_total_blocks=1024)
|
||||
a = alloc.acquire(10)
|
||||
assert alloc.num_free == 1014
|
||||
b = alloc.acquire(5)
|
||||
assert alloc.num_free == 1009
|
||||
alloc.release(a)
|
||||
assert alloc.num_free == 1019
|
||||
alloc.release(b)
|
||||
assert alloc.num_free == 1024
|
||||
# Re-acquire works after release.
|
||||
c = alloc.acquire(20)
|
||||
assert alloc.num_free == 1004
|
||||
|
||||
|
||||
def test_oom_raises():
|
||||
alloc = BlockAllocator(num_total_blocks=4)
|
||||
alloc.acquire(4)
|
||||
with pytest.raises(RuntimeError, match="OOM"):
|
||||
alloc.acquire(1)
|
||||
|
||||
|
||||
def test_acquire_returns_unique_ids():
|
||||
alloc = BlockAllocator(num_total_blocks=100)
|
||||
a = alloc.acquire(50)
|
||||
b = alloc.acquire(50)
|
||||
assert len(torch.intersect1d(a, b)) == 0
|
||||
|
||||
|
||||
# ---- Pool shape tests ----
|
||||
|
||||
def test_paged_pool_shapes_csa():
|
||||
from dsv4.model.layer_schedule import FFNType
|
||||
config = DSV4Config.pro()
|
||||
spec = LayerSpec(layer_idx=2, attn=AttentionType.CSA,
|
||||
ffn=FFNType.MOE, router_mode=RouterMode.HASH)
|
||||
schema = build_schema(config, spec)
|
||||
pool = PagedKVPool(schema, num_blocks=16)
|
||||
assert pool.entries_fp8.shape == (16, 32, config.head_dim - config.rope_dim)
|
||||
assert pool.entries_rope.shape == (16, 32, config.rope_dim)
|
||||
assert pool.inv_scale.shape == (16, 32)
|
||||
assert pool.indexer_keys_fp4 is not None
|
||||
assert pool.indexer_keys_fp4.shape[1] == 32
|
||||
|
||||
|
||||
def test_paged_pool_shapes_hca():
|
||||
from dsv4.model.layer_schedule import FFNType
|
||||
config = DSV4Config.pro()
|
||||
spec = LayerSpec(layer_idx=3, attn=AttentionType.HCA,
|
||||
ffn=FFNType.MOE, router_mode=RouterMode.DENSE)
|
||||
schema = build_schema(config, spec)
|
||||
pool = PagedKVPool(schema, num_blocks=256)
|
||||
assert pool.entries_fp8.shape == (256, 1, config.head_dim - config.rope_dim)
|
||||
assert pool.entries_rope.shape == (256, 1, config.rope_dim)
|
||||
assert pool.indexer_keys_fp4 is None
|
||||
|
||||
|
||||
def test_state_pool_shapes_csa():
|
||||
from dsv4.model.layer_schedule import FFNType
|
||||
config = DSV4Config.pro()
|
||||
spec = LayerSpec(layer_idx=2, attn=AttentionType.CSA,
|
||||
ffn=FFNType.MOE, router_mode=RouterMode.HASH)
|
||||
schema = build_schema(config, spec)
|
||||
pool = StateCachePool(schema, max_requests=8)
|
||||
assert pool.swa_fp8.shape == (8, 128, config.head_dim - config.rope_dim)
|
||||
assert pool.swa_rope.shape == (8, 128, config.rope_dim)
|
||||
assert pool.tail_ka is not None
|
||||
assert pool.tail_kb is not None # CSA has both streams
|
||||
assert pool.tail_len is not None
|
||||
|
||||
|
||||
def test_state_pool_shapes_hca():
|
||||
from dsv4.model.layer_schedule import FFNType
|
||||
config = DSV4Config.pro()
|
||||
spec = LayerSpec(layer_idx=3, attn=AttentionType.HCA,
|
||||
ffn=FFNType.MOE, router_mode=RouterMode.DENSE)
|
||||
schema = build_schema(config, spec)
|
||||
pool = StateCachePool(schema, max_requests=8)
|
||||
assert pool.tail_ka is not None
|
||||
assert pool.tail_kb is None # HCA only one stream
|
||||
assert pool.tail_za is not None
|
||||
assert pool.tail_len is not None
|
||||
|
||||
|
||||
def test_state_pool_shapes_swa():
|
||||
from dsv4.model.layer_schedule import FFNType
|
||||
config = DSV4Config.flash()
|
||||
spec = LayerSpec(layer_idx=0, attn=AttentionType.SWA,
|
||||
ffn=FFNType.MOE, router_mode=RouterMode.HASH)
|
||||
schema = build_schema(config, spec)
|
||||
pool = StateCachePool(schema, max_requests=8)
|
||||
assert pool.swa_fp8.shape == (8, 128, config.head_dim - config.rope_dim)
|
||||
assert pool.tail_ka is None # No tail for SWA-only
|
||||
assert pool.tail_len is None
|
||||
|
||||
|
||||
# ---- Manager lifecycle tests ----
|
||||
|
||||
def test_admit_release_recycles_slot():
|
||||
config = DSV4Config.flash()
|
||||
schedule = build_schedule(config)
|
||||
mgr = KVCacheManager(config, schedule, max_concurrent_requests=4,
|
||||
num_blocks_per_csa_layer=64, num_blocks_per_hca_layer=64)
|
||||
s1 = mgr.admit_request()
|
||||
mgr.release_request(s1)
|
||||
s2 = mgr.admit_request()
|
||||
assert s1 == s2 # slot was recycled
|
||||
|
||||
|
||||
def test_admit_exhaustion():
|
||||
config = DSV4Config.flash()
|
||||
schedule = build_schedule(config)
|
||||
mgr = KVCacheManager(config, schedule, max_concurrent_requests=2,
|
||||
num_blocks_per_csa_layer=64, num_blocks_per_hca_layer=64)
|
||||
mgr.admit_request()
|
||||
mgr.admit_request()
|
||||
with pytest.raises(RuntimeError, match="concurrent"):
|
||||
mgr.admit_request()
|
||||
|
||||
|
||||
def test_handle_construction_no_alloc():
|
||||
"""acquire() should not allocate GPU memory — critical for cudagraph."""
|
||||
config = DSV4Config.flash()
|
||||
schedule = build_schedule(config)
|
||||
mgr = KVCacheManager(config, schedule, max_concurrent_requests=4,
|
||||
num_blocks_per_csa_layer=64, num_blocks_per_hca_layer=64)
|
||||
slot = mgr.admit_request()
|
||||
torch.cuda.synchronize()
|
||||
before = torch.cuda.memory_allocated()
|
||||
handle = mgr.acquire(
|
||||
layer_idx=0,
|
||||
request_slots=torch.tensor([slot], dtype=torch.int32, device="cuda"),
|
||||
positions=torch.tensor([0], dtype=torch.int32, device="cuda"),
|
||||
request_ids=torch.tensor([0], dtype=torch.int32, device="cuda"),
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
after = torch.cuda.memory_allocated()
|
||||
# Allow small variance from tensor view creation, but no large alloc
|
||||
assert after - before < 1024, f"acquire() allocated {after - before} bytes — breaks cudagraph"
|
||||
assert handle.paged is None # layer 0 is SWA
|
||||
mgr.release_request(slot)
|
||||
|
||||
|
||||
def test_manager_memory_tracking():
|
||||
config = DSV4Config.flash()
|
||||
schedule = build_schedule(config)
|
||||
mgr = KVCacheManager(config, schedule, max_concurrent_requests=4,
|
||||
num_blocks_per_csa_layer=64, num_blocks_per_hca_layer=64)
|
||||
total = mgr.memory_bytes()
|
||||
assert total > 0
|
||||
# Rough sanity: should be in the MB range for this config
|
||||
assert total > 1_000_000 # at least 1 MB
|
||||
assert total < 10_000_000_000 # less than 10 GB
|
||||
|
||||
|
||||
def test_full_flash_stack_construction():
|
||||
"""Construct manager with all 43 Flash layers — pools for every layer."""
|
||||
config = DSV4Config.flash()
|
||||
schedule = build_schedule(config)
|
||||
mgr = KVCacheManager(config, schedule, max_concurrent_requests=4,
|
||||
num_blocks_per_csa_layer=64, num_blocks_per_hca_layer=64)
|
||||
# 2 SWA layers (no paged pool) + 41 compressed layers
|
||||
assert len(mgr.paged_pools) == 43
|
||||
assert mgr.paged_pools[0] is None # SWA
|
||||
assert mgr.paged_pools[1] is None # SWA
|
||||
assert mgr.paged_pools[2] is not None # CSA
|
||||
assert mgr.paged_pools[3] is not None # HCA
|
||||
|
||||
# All state pools present
|
||||
assert len(mgr.state_pools) == 43
|
||||
for i in range(43):
|
||||
assert mgr.state_pools[i] is not None
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
@@ -1,189 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Test torch.compile + CuTeDSL NVFP4 runner via custom_op.
|
||||
|
||||
Critical test: does torch.compile (fullgraph mode) accept the
|
||||
nvfp4::moe_gemm custom op and produce a working compiled graph?
|
||||
|
||||
Run on the B200:
|
||||
docker run --rm --gpus all --entrypoint python3 \
|
||||
-v /root/nvfp4-megamoe-kernel:/root/nvfp4-megamoe-kernel \
|
||||
-v /root/nvidia-meeting:/root/nvidia-meeting:ro \
|
||||
nvfp4-megamoe-kernel-vllm:latest \
|
||||
/root/nvfp4-megamoe-kernel/tests/test_compile_custom_op.py
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
import glob
|
||||
import torch
|
||||
from safetensors import safe_open
|
||||
|
||||
REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
sys.path.insert(0, REPO_ROOT)
|
||||
|
||||
from dsv4.layers.moe import Nvfp4MoE
|
||||
from dsv4.ops.custom_ops import register_runner, nvfp4_moe_gemm
|
||||
|
||||
NVFP4_MODEL_DIR = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4"
|
||||
DEVICE = "cuda"
|
||||
|
||||
|
||||
def find_shards(model_dir):
|
||||
index_path = os.path.join(model_dir, "model.safetensors.index.json")
|
||||
key_to_shard = {}
|
||||
if os.path.exists(index_path):
|
||||
with open(index_path) as f:
|
||||
index = json.load(f)
|
||||
for key, shard in index["weight_map"].items():
|
||||
key_to_shard[key] = os.path.join(model_dir, shard)
|
||||
else:
|
||||
for sf in glob.glob(os.path.join(model_dir, "*.safetensors")):
|
||||
with safe_open(sf, framework="pt") as f:
|
||||
for key in f.keys():
|
||||
key_to_shard[key] = sf
|
||||
return key_to_shard
|
||||
|
||||
|
||||
def load_layer_tensors(model_dir, layer_idx):
|
||||
key_to_shard = find_shards(model_dir)
|
||||
layer_prefix = f"layers.{layer_idx}."
|
||||
shard_to_keys = {}
|
||||
for key, shard in key_to_shard.items():
|
||||
norm_key = key.removeprefix("model.")
|
||||
if not norm_key.startswith(layer_prefix):
|
||||
continue
|
||||
shard_to_keys.setdefault(shard, []).append((key, norm_key))
|
||||
tensors = {}
|
||||
for shard, keys in shard_to_keys.items():
|
||||
with safe_open(shard, framework="pt") as f:
|
||||
for orig_key, norm_key in keys:
|
||||
tensors[norm_key] = f.get_tensor(orig_key)
|
||||
return tensors
|
||||
|
||||
|
||||
def prepare_nvfp4_weights_direct(nvfp4_tensors, layer_idx, expert_indices, intermediate_size):
|
||||
from dsv4.ops.quantize import (
|
||||
quantize_activation_nvfp4,
|
||||
quantize_weight_to_nvfp4,
|
||||
)
|
||||
l1_fp4, l1_sf, l1_gs = [], [], []
|
||||
l2_fp4, l2_sf, l2_gs = [], [], []
|
||||
|
||||
for e in expert_indices:
|
||||
gate_w = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.gate_proj.weight"].to(DEVICE)
|
||||
up_w = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.up_proj.weight"].to(DEVICE)
|
||||
gate_sf = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.gate_proj.weight_scale"].to(DEVICE)
|
||||
up_sf = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.up_proj.weight_scale"].to(DEVICE)
|
||||
gate_gs = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.gate_proj.weight_scale_2"].item()
|
||||
up_gs = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.up_proj.weight_scale_2"].item()
|
||||
|
||||
fused_w = torch.cat([gate_w, up_w], dim=0)
|
||||
fused_w_fp4 = fused_w.view(torch.float4_e2m1fn_x2).permute(1, 0).contiguous()
|
||||
fused_sf = torch.cat([gate_sf, up_sf], dim=0).permute(1, 0).contiguous()
|
||||
l1_max_gs = max(gate_gs, up_gs)
|
||||
if gate_gs != up_gs:
|
||||
fused_sf_f32 = fused_sf.float()
|
||||
fused_sf_f32[:, :intermediate_size] *= (gate_gs / l1_max_gs)
|
||||
fused_sf_f32[:, intermediate_size:] *= (up_gs / l1_max_gs)
|
||||
fused_sf = fused_sf_f32.to(torch.float8_e4m3fn)
|
||||
l1_fp4.append(fused_w_fp4)
|
||||
l1_sf.append(fused_sf)
|
||||
l1_gs.append(l1_max_gs)
|
||||
|
||||
down_key = f"layers.{layer_idx}.mlp.experts.{e}.down_proj.weight"
|
||||
if down_key in nvfp4_tensors:
|
||||
down_w = nvfp4_tensors[down_key].to(DEVICE)
|
||||
down_sf = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.down_proj.weight_scale"].to(DEVICE)
|
||||
down_gs = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.down_proj.weight_scale_2"].item()
|
||||
l2_fp4.append(down_w.view(torch.float4_e2m1fn_x2).permute(1, 0).contiguous())
|
||||
l2_sf.append(down_sf.permute(1, 0).contiguous())
|
||||
l2_gs.append(down_gs)
|
||||
|
||||
return {
|
||||
'l1_fp4': l1_fp4, 'l1_sf': l1_sf, 'l1_gs': l1_gs,
|
||||
'l2_fp4': l2_fp4, 'l2_sf': l2_sf, 'l2_gs': l2_gs,
|
||||
}
|
||||
|
||||
|
||||
def main():
|
||||
torch.manual_seed(42)
|
||||
expert_indices = [0, 1, 2]
|
||||
hidden_size = 7168
|
||||
intermediate_size = 3072
|
||||
|
||||
print("=" * 70)
|
||||
print(" torch.compile + CuTeDSL Custom Op Test")
|
||||
print("=" * 70)
|
||||
|
||||
# Load weights
|
||||
nvfp4_tensors = load_layer_tensors(NVFP4_MODEL_DIR, 0)
|
||||
weights = prepare_nvfp4_weights_direct(nvfp4_tensors, 0, expert_indices, intermediate_size)
|
||||
|
||||
# Create runner
|
||||
runner = Nvfp4MoE(
|
||||
num_experts=len(expert_indices),
|
||||
hidden_size=hidden_size,
|
||||
intermediate_size=intermediate_size,
|
||||
max_num_tokens=8,
|
||||
top_k=2,
|
||||
device="cuda",
|
||||
)
|
||||
runner.prepare_weights_direct(
|
||||
weights['l1_fp4'], weights['l1_sf'], weights['l1_gs'],
|
||||
weights['l2_fp4'], weights['l2_sf'], weights['l2_gs'],
|
||||
)
|
||||
runner_id = register_runner(runner)
|
||||
|
||||
# Test input
|
||||
hidden_states = torch.randn(4, hidden_size, dtype=torch.bfloat16, device=DEVICE) * 2.0
|
||||
topk_ids = torch.tensor([[0, 1]] * 4, dtype=torch.int32, device=DEVICE)
|
||||
topk_weights = torch.tensor([[0.6, 0.4]] * 4, dtype=torch.float32, device=DEVICE)
|
||||
|
||||
# 1. Warmup: compute activation global scales
|
||||
print("\n[0] Computing activation global scales (warmup)...")
|
||||
runner.compute_activation_global_scales(hidden_states, topk_weights, topk_ids)
|
||||
print(f" L1 gs: {runner._l1_activation_global_scale:.6f}")
|
||||
print(f" L2 gs: {runner._l2_activation_global_scale:.6f}")
|
||||
|
||||
# 1. Eager mode (baseline)
|
||||
print("\n[1/2] Running eager mode (baseline)...")
|
||||
runner._ensure_stacked()
|
||||
eager_out = nvfp4_moe_gemm(hidden_states, topk_weights, topk_ids, runner_id, hidden_size)
|
||||
print(f" Eager output: amax={eager_out.abs().max():.4f} mean={eager_out.float().mean():.6f}")
|
||||
|
||||
# 2. torch.compile fullgraph
|
||||
print("\n[2/2] Running torch.compile(fullgraph=True)...")
|
||||
try:
|
||||
@torch.compile(fullgraph=True)
|
||||
def compiled_fn(hs, tw, ti):
|
||||
return nvfp4_moe_gemm(hs, tw, ti, runner_id, hidden_size)
|
||||
|
||||
compiled_out = compiled_fn(hidden_states, topk_weights, topk_ids)
|
||||
print(f" Compiled output: amax={compiled_out.abs().max():.4f} mean={compiled_out.float().mean():.6f}")
|
||||
|
||||
# Compare
|
||||
if eager_out.shape == compiled_out.shape:
|
||||
cos = torch.nn.functional.cosine_similarity(
|
||||
eager_out.flatten().unsqueeze(0).float(),
|
||||
compiled_out.flatten().unsqueeze(0).float(),
|
||||
).item()
|
||||
print(f"\n Eager vs Compiled: cosine={cos:.6f}")
|
||||
if cos > 0.99:
|
||||
print(" ✅ torch.compile produces matching output!")
|
||||
else:
|
||||
print(f" ⚠️ Cosine {cos:.4f} < 0.99 — check for numerical issues")
|
||||
else:
|
||||
print(f" ❌ Shape mismatch: eager={eager_out.shape} compiled={compiled_out.shape}")
|
||||
except Exception as e:
|
||||
print(f" ❌ torch.compile FAILED: {type(e).__name__}: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
sys.exit(1)
|
||||
|
||||
print("\n" + "=" * 70)
|
||||
print(" Test complete ✅")
|
||||
print("=" * 70)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,138 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Test that torch.library.custom_op wrapping works with torch.compile.
|
||||
|
||||
This tests the Dynamo opaqueness without needing a GPU — we just verify:
|
||||
1. The custom_op is registered correctly
|
||||
2. torch.compile treats it as opaque (doesn't try to trace through it)
|
||||
3. FakeTensor shape inference works
|
||||
4. The runner registry works
|
||||
|
||||
Does NOT test actual GEMM output — that needs the B200.
|
||||
"""
|
||||
import sys
|
||||
import os
|
||||
import torch
|
||||
|
||||
REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
sys.path.insert(0, REPO_ROOT)
|
||||
|
||||
|
||||
def test_custom_op_registered():
|
||||
"""Verify nvfp4::linear_gemm and nvfp4::moe_gemm are registered."""
|
||||
from dsv4.ops.custom_ops import nvfp4_linear_gemm, nvfp4_moe_gemm
|
||||
|
||||
# Check they exist as custom ops
|
||||
assert hasattr(nvfp4_linear_gemm, '_name')
|
||||
assert hasattr(nvfp4_moe_gemm, '_name')
|
||||
print("✅ Custom ops registered")
|
||||
|
||||
|
||||
def test_runner_registry():
|
||||
"""Test the runner registry."""
|
||||
from dsv4.ops.custom_ops import register_runner, get_runner
|
||||
|
||||
class FakeRunner:
|
||||
def _run_impl(self, x):
|
||||
return x * 2
|
||||
|
||||
runner = FakeRunner()
|
||||
rid = register_runner(runner)
|
||||
assert rid >= 0
|
||||
|
||||
retrieved = get_runner(rid)
|
||||
assert retrieved is runner
|
||||
print(f"✅ Runner registry works (id={rid})")
|
||||
|
||||
|
||||
def test_fake_tensor_shape_inference():
|
||||
"""Test that FakeTensor impl returns correct shapes."""
|
||||
from dsv4.ops.custom_ops import nvfp4_linear_gemm, nvfp4_moe_gemm
|
||||
|
||||
# linear_gemm fake impl
|
||||
x_fake = torch.empty(4, 7168, dtype=torch.bfloat16, device='meta')
|
||||
out_fake = nvfp4_linear_gemm(x_fake, runner_id=0, out_features=3072)
|
||||
assert out_fake.shape == (4, 3072), f"Expected (4, 3072), got {out_fake.shape}"
|
||||
print(f"✅ linear_gemm fake impl: {x_fake.shape} → {out_fake.shape}")
|
||||
|
||||
# moe_gemm fake impl
|
||||
hs_fake = torch.empty(4, 7168, dtype=torch.bfloat16, device='meta')
|
||||
tw_fake = torch.empty(4, 8, dtype=torch.float32, device='meta')
|
||||
ti_fake = torch.empty(4, 8, dtype=torch.int32, device='meta')
|
||||
out_fake = nvfp4_moe_gemm(hs_fake, tw_fake, ti_fake, runner_id=0, hidden_size=7168)
|
||||
assert out_fake.shape == (4, 7168), f"Expected (4, 7168), got {out_fake.shape}"
|
||||
print(f"✅ moe_gemm fake impl: {hs_fake.shape} → {out_fake.shape}")
|
||||
|
||||
|
||||
def test_torch_compile_skips_custom_op():
|
||||
"""Test that torch.compile doesn't try to trace through the custom op.
|
||||
|
||||
This is the critical test — if compile tries to inline the op, it will
|
||||
fail because the runner's _run_impl uses CuTeDSL internals.
|
||||
|
||||
We use a fake runner that would crash if traced (raises on first call).
|
||||
If torch.compile correctly treats it as opaque, it won't call it during
|
||||
compilation — only the fake impl runs.
|
||||
"""
|
||||
from dsv4.ops.custom_ops import register_runner, nvfp4_linear_gemm
|
||||
|
||||
class ExplodingRunner:
|
||||
"""Runner that explodes if _run_impl is ever called."""
|
||||
call_count = 0
|
||||
def _run_impl(self, x):
|
||||
self.call_count += 1
|
||||
return x # This should never be called during compilation
|
||||
|
||||
runner = ExplodingRunner()
|
||||
rid = register_runner(runner)
|
||||
|
||||
# Compile a function that uses our custom op
|
||||
@torch.compile(fullgraph=True)
|
||||
def forward(x):
|
||||
return nvfp4_linear_gemm(x, runner_id=rid, out_features=3072)
|
||||
|
||||
# With CPU tensors, compile should trace through using FakeTensors
|
||||
# and never call _run_impl
|
||||
x = torch.randn(4, 7168, dtype=torch.bfloat16)
|
||||
# This will fail on CPU because _run_impl needs CUDA, but the point
|
||||
# is that Dynamo should accept the custom op without error.
|
||||
# If it tries to trace through it, we'd get a different error.
|
||||
|
||||
# Instead, just verify Dynamo can handle the graph with custom ops
|
||||
# by checking that the op shows up in the graph
|
||||
try:
|
||||
# Use torch._dynamo to trace without executing
|
||||
import torch._dynamo as dynamo
|
||||
gm, guards = dynamo.export(forward)(x)
|
||||
graph_str = str(gm.graph)
|
||||
assert "nvfp4_linear_gemm" in graph_str, \
|
||||
f"Custom op not found in compiled graph. Graph:\n{graph_str}"
|
||||
print("✅ torch.compile treats custom op as opaque (not inlined)")
|
||||
print(f" Graph contains: ...nvfp4_linear_gemm...")
|
||||
except Exception as e:
|
||||
# On CPU without CUDA, _run_impl can't run. That's fine —
|
||||
# the important thing is Dynamo didn't try to INLINE the op.
|
||||
# If Dynamo tried to trace through it, the error would mention
|
||||
# CuTeDSL/cute.compile, not CUDA.
|
||||
error_str = str(e)
|
||||
if "CuTeDSL" in error_str or "cute" in error_str:
|
||||
print(f"❌ Dynamo tried to trace through the custom op!")
|
||||
print(f" Error: {e}")
|
||||
sys.exit(1)
|
||||
else:
|
||||
print(f"⚠️ Execution error (expected on CPU): {type(e).__name__}")
|
||||
print(f" Dynamo accepted the custom op as opaque ✅")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("=" * 60)
|
||||
print(" Custom Op Dynamo Compatibility Tests")
|
||||
print("=" * 60)
|
||||
|
||||
test_custom_op_registered()
|
||||
test_runner_registry()
|
||||
test_fake_tensor_shape_inference()
|
||||
test_torch_compile_skips_custom_op()
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print(" All tests passed ✅")
|
||||
print("=" * 60)
|
||||
@@ -1,27 +0,0 @@
|
||||
# tests/unit/test_dense_router.py
|
||||
import torch
|
||||
from dsv4.layers.router import Router
|
||||
|
||||
def test_dense_router_matches_spec(N=64, H=4096, E=256, k=6):
|
||||
X = torch.randn(N, H, dtype=torch.bfloat16, device='cuda')
|
||||
W = torch.randn(H, E, dtype=torch.bfloat16, device='cuda')
|
||||
bias = torch.randn(E, dtype=torch.float32, device='cuda') * 0.01
|
||||
scaling = 2.5
|
||||
|
||||
# Oracle: directly compute the spec, in one expression, in FP32.
|
||||
# This is not "a PyTorch reference implementation" — it's the math.
|
||||
logits = (X.float() @ W.float())
|
||||
act = torch.sqrt(torch.nn.functional.softplus(logits))
|
||||
score = act + bias
|
||||
ids = score.topk(k, dim=-1).indices
|
||||
w = act.gather(-1, ids)
|
||||
w = w / w.sum(-1, keepdim=True) * scaling
|
||||
|
||||
# Kernel under test:
|
||||
router = Router(H, E, k, scaling, mode='dense')
|
||||
router.W_gate.copy_(W)
|
||||
router.e_bias.copy_(bias)
|
||||
out_w, out_ids = router(X, layer_idx=5)
|
||||
|
||||
assert (out_ids == ids).all() # ids must be exact match
|
||||
torch.testing.assert_close(out_w, w, atol=1e-4, rtol=1e-3)
|
||||
@@ -1,317 +0,0 @@
|
||||
"""
|
||||
FMHA Stage-C multi-tile diagnostic.
|
||||
Tests whether the combined K+V pipeline actually loads different KV tiles.
|
||||
Runs identity softmax (no rescale, no online softmax) to isolate the
|
||||
TMA + MMA pipeline from softmax logic.
|
||||
"""
|
||||
import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline
|
||||
from cutlass.cute.nvgpu import cpasync, tcgen05
|
||||
from cutlass import Float32, BFloat16, Int32, Boolean, const_expr
|
||||
from cutlass.utils import LayoutEnum
|
||||
from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset
|
||||
import cuda.bindings.driver as cuda
|
||||
import cutlass.torch as ct
|
||||
import math
|
||||
|
||||
HEAD_DIM = 64
|
||||
|
||||
|
||||
class FmhaV3Diag:
|
||||
"""Identity-softmax multi-tile with combined K+V barrier.
|
||||
If this fails, the problem is in TMA/MMA, not softmax."""
|
||||
def __init__(self, s_k=256):
|
||||
self.s_k = s_k
|
||||
self.n_kv_tiles = s_k // 128
|
||||
self.acc_dtype = Float32; self.qk_acc_dtype = Float32
|
||||
self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16
|
||||
self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1
|
||||
self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE
|
||||
self.epilogue_warp_id = (0,1,2,3); self.mma_warp_id = 4; self.tma_warp_id = 5
|
||||
self.threads_per_cta = 192; self.num_c_stage = 2
|
||||
self.kv_stage = 2; self.q_stage = 1; self.num_c_stage = 2
|
||||
|
||||
def _setup(self, qk_mma, pv_mma):
|
||||
qk_ik = cute.size(qk_mma.shape_mnk, mode=[2])
|
||||
self.qk_mma_tiler = (128, 128, qk_ik * 4)
|
||||
pv_ik = cute.size(pv_mma.shape_mnk, mode=[2])
|
||||
self.pv_mma_tiler = (128, HEAD_DIM, pv_ik * (128 // pv_ik))
|
||||
self.mma_tiler = self.qk_mma_tiler
|
||||
self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,))
|
||||
self.cta_tile_shape_mnk = (self.qk_mma_tiler[0]//cute.size(qk_mma.thr_id.shape), HEAD_DIM, self.qk_mma_tiler[2])
|
||||
self.c_layout = LayoutEnum.ROW_MAJOR
|
||||
self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype)
|
||||
self.num_ab_stage = 1; self.num_acc_stage = 1
|
||||
self.q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, self.q_dtype, self.q_stage)
|
||||
self.k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, self.q_dtype, self.kv_stage)
|
||||
self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, self.kv_stage)
|
||||
self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2)
|
||||
self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1)
|
||||
qk_thr = qk_mma.get_slice(0); qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2])
|
||||
tStS = qk_thr.make_fragment_C(qk_as)
|
||||
pv_thr = pv_mma.get_slice(0); pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2])
|
||||
tOtO = pv_thr.make_fragment_C(pv_as)
|
||||
self.tmem_s0_offset = 0; self.tmem_p0_offset = 32
|
||||
p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width
|
||||
p_end = self.tmem_p0_offset + p_cols_fp32
|
||||
s_cols = self.qk_mma_tiler[1]
|
||||
o_after = max(s_cols, p_end)
|
||||
self.tmem_o0_offset = ((o_after + 31) // 32) * 32
|
||||
o_cols = find_tmem_tensor_col_offset(tOtO)
|
||||
total = self.tmem_o0_offset + o_cols
|
||||
self.num_tmem_alloc_cols = 1
|
||||
while self.num_tmem_alloc_cols < total:
|
||||
self.num_tmem_alloc_cols *= 2
|
||||
cta = cute.size(qk_mma.thr_id.shape)
|
||||
q_s = cute.slice_(self.q_smem_s,(None,None,None,0))
|
||||
k_s = cute.slice_(self.k_smem_s,(None,None,None,0))
|
||||
v_s = cute.slice_(self.v_smem_s,(None,None,None,0))
|
||||
self.q_tx_bytes = cute.size_in_bytes(self.q_dtype, q_s) * cta
|
||||
self.kv_tx_bytes = (cute.size_in_bytes(self.q_dtype, k_s) +
|
||||
cute.size_in_bytes(self.q_dtype, v_s)) * cta
|
||||
|
||||
@cute.jit
|
||||
def __call__(self, q, k, v, c, stream):
|
||||
self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype
|
||||
self.a_major = LayoutEnum.from_tensor(q).mma_major_mode()
|
||||
self.b_major = LayoutEnum.from_tensor(k).mma_major_mode()
|
||||
v_fmha = cute.make_tensor(
|
||||
v.iterator,
|
||||
cute.make_layout(
|
||||
(HEAD_DIM, self.s_k, 1),
|
||||
stride=(1, HEAD_DIM, HEAD_DIM * self.s_k),
|
||||
),
|
||||
)
|
||||
self.v_major = LayoutEnum.from_tensor(v_fmha).mma_major_mode()
|
||||
self.c_layout = LayoutEnum.from_tensor(c)
|
||||
qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, self.a_major, self.b_major, self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM)
|
||||
pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, self.qk_acc_dtype, self.cta_group, (128,HEAD_DIM), tcgen05.OperandSource.TMEM)
|
||||
self._setup(qk_mma, pv_mma)
|
||||
q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)); v_s = cute.slice_(self.v_smem_s,(None,None,None,0))
|
||||
tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id),q,q_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape)
|
||||
tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id),k,k_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape)
|
||||
tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v_fmha,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape)
|
||||
epi_s = cute.select(self.c_smem_s,mode=[0,1])
|
||||
tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile)
|
||||
self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.v_smem_s,self.p_tmem_s,self.c_smem_s,self.epi_tile).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream)
|
||||
|
||||
@cute.kernel
|
||||
def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, q_smem_s, k_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile):
|
||||
warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
|
||||
tidx,_,_ = cute.arch.thread_idx()
|
||||
if warp_idx == self.tma_warp_id:
|
||||
cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k); cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c)
|
||||
|
||||
@cute.struct
|
||||
class SS:
|
||||
q_bar: cute.struct.MemRange[cutlass.Int64, self.q_stage*2]
|
||||
kv_bar: cute.struct.MemRange[cutlass.Int64, self.kv_stage*2]
|
||||
s_bar: cute.struct.MemRange[cutlass.Int64, 2]
|
||||
acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage*2]
|
||||
tmem_dealloc: cutlass.Int64; holding: cutlass.Int32
|
||||
smem = utils.SmemAllocator(); st = smem.allocate(SS)
|
||||
|
||||
qp,qc = pipeline.PipelineTmaUmma.create(barrier_storage=st.q_bar.data_ptr(),num_stages=self.q_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.q_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants()
|
||||
kvp,kvc = pipeline.PipelineTmaUmma.create(barrier_storage=st.kv_bar.data_ptr(),num_stages=self.kv_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.kv_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants()
|
||||
s_prod,s_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.s_bar.data_ptr(),num_stages=1,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.epilogue_warp_id))).make_participants()
|
||||
softmax_done_bar = pipeline.NamedBarrier(barrier_id=3, num_threads=32 + 32*len(self.epilogue_warp_id))
|
||||
acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(),num_stages=self.num_acc_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,len(self.epilogue_warp_id)),cta_layout_vmnk=cl_vmnk,defer_sync=True)
|
||||
tmem_bar = pipeline.NamedBarrier(barrier_id=2,num_threads=32*len((self.mma_warp_id,*self.epilogue_warp_id)))
|
||||
tmem = utils.TmemAllocator(st.holding.ptr,barrier_for_retrieve=tmem_bar,allocator_warp_id=self.epilogue_warp_id[0],is_two_cta=cute.size(qk_mma.thr_id.shape)==2,two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr)
|
||||
pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk,is_relaxed=True)
|
||||
|
||||
sQ = smem.allocate_tensor(element_type=self.q_dtype,layout=q_smem_s.outer,byte_alignment=128,swizzle=q_smem_s.inner)
|
||||
sK = smem.allocate_tensor(element_type=self.q_dtype,layout=k_smem_s.outer,byte_alignment=128,swizzle=k_smem_s.inner)
|
||||
sV = smem.allocate_tensor(element_type=self.q_dtype,layout=v_smem_s.outer,byte_alignment=128,swizzle=v_smem_s.inner)
|
||||
sC = smem.allocate_tensor(element_type=self.o_dtype,layout=c_smem_s.outer,byte_alignment=128,swizzle=c_smem_s.inner)
|
||||
|
||||
gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None))
|
||||
gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None))
|
||||
gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None))
|
||||
gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None))
|
||||
n_kv_tiles = self.n_kv_tiles
|
||||
|
||||
qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0)
|
||||
tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK)
|
||||
tCgV = pv_thr.partition_B(gV); tCgC = pv_thr.partition_C(gC)
|
||||
a_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,0,None,0)).shape)
|
||||
tAsQ,tAgQ = cpasync.tma_partition(tma_q,0,a_lay,cute.group_modes(sQ,0,3),cute.group_modes(tCgQ,0,3))
|
||||
b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape)
|
||||
tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3))
|
||||
tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3))
|
||||
# =====================================================================
|
||||
# ⚠️⚠️⚠️ CRITICAL: TMA PARTITION TENSOR MODE ORDERING ⚠️⚠️⚠️
|
||||
# After tma_partition, tBgK/tVgV have 4 modes: (V_grouped, ?, KV_tiles, ?)
|
||||
# Mode 4 is the GMEM tile dimension. DO NOT pre-slice to fewer modes!
|
||||
# See README.md for details.
|
||||
# =====================================================================
|
||||
tAgQ = tAgQ[(None,0,None,0)]
|
||||
tBgK = tBgK[(None,0,None,0)]
|
||||
tVgV = tVgV[(None,0,None,0)]
|
||||
|
||||
tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK)
|
||||
tCrV = pv_mma.make_fragment_B(sV)
|
||||
|
||||
qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2])
|
||||
tStS = qk_thr.make_fragment_C(qk_as)
|
||||
tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout)
|
||||
pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2])
|
||||
tOtO = pv_thr.make_fragment_C(pv_as)
|
||||
tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout)
|
||||
|
||||
tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer)
|
||||
tOrP_base = pv_thr.make_fragment_A(tP)
|
||||
tOrP = tOrP_base[(None,None,None,0)]
|
||||
tOrP0 = cute.make_tensor(
|
||||
tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset,
|
||||
tOrP.layout)
|
||||
|
||||
tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage))
|
||||
tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage))
|
||||
pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk)
|
||||
|
||||
# TMA LOAD (combined K+V barrier, Int32(kt) for GMEM coord)
|
||||
if warp_idx == self.tma_warp_id:
|
||||
qp.reset(); qh = qp.acquire_and_advance()
|
||||
cute.copy(tma_q, tAgQ[(None, Int32(0))], tAsQ[(None, qh.index)], tma_bar_ptr=qh.barrier)
|
||||
qp.tail()
|
||||
kvp.reset(); pk = kvp.try_acquire()
|
||||
# Force runtime Int32 (not literal) — option 3 from CUTLASS LLM
|
||||
kv_coord = Int32(0 + 0)
|
||||
for kt in cutlass.range(self.n_kv_tiles, unroll=1):
|
||||
kvh = kvp.acquire_and_advance(pk)
|
||||
# kv_coord indexes the surviving mode 1 (from original mode 2) of 2D tBgK/tVgV.
|
||||
# Using (None, kv_coord) on a pre-sliced 4-mode tensor SILENTLY BREAKS multi-tile!
|
||||
cute.copy(tma_k, tBgK[None, kv_coord], tBsK[(None, kvh.index)], tma_bar_ptr=kvh.barrier)
|
||||
cute.copy(tma_v, tVgV[None, kv_coord], tVsV[(None, kvh.index)], tma_bar_ptr=kvh.barrier)
|
||||
kv_coord += 1
|
||||
pk = cutlass.Boolean(1)
|
||||
kvp.tail()
|
||||
|
||||
# MMA (combined K+V slot)
|
||||
if warp_idx == self.mma_warp_id:
|
||||
tmem.wait_for_alloc()
|
||||
qc.reset(); qh = qc.wait_and_advance(); qh.release()
|
||||
kvc.reset(); pk = kvc.try_wait()
|
||||
acc_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_stage)
|
||||
acc_pipe.producer_acquire(acc_st)
|
||||
for kt in range(self.n_kv_tiles):
|
||||
kvh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1)
|
||||
sh = s_prod.acquire_and_advance()
|
||||
qk_mma.set(tcgen05.Field.ACCUMULATE, False)
|
||||
for kb in cutlass.range(cute.size(tCrQ, mode=[2]), unroll_full=True):
|
||||
cute.gemm(qk_mma, tStS0, tCrQ[(None,None,kb,0)], tCrK[(None,None,kb,kvh.index)], tStS0)
|
||||
qk_mma.set(tcgen05.Field.ACCUMULATE, True)
|
||||
cute.arch.fence_view_async_tmem_store()
|
||||
sh.commit()
|
||||
softmax_done_bar.arrive_and_wait()
|
||||
pv_mma.set(tcgen05.Field.ACCUMULATE, kt != 0)
|
||||
for kb in cutlass.range(cute.size(tOrP0, mode=[2]), unroll_full=True):
|
||||
cute.gemm(pv_mma, tOtO0, tOrP0[(None,None,kb)], tCrV[(None,None,kb,kvh.index)], tOtO0)
|
||||
pv_mma.set(tcgen05.Field.ACCUMULATE, True)
|
||||
cute.arch.fence_view_async_tmem_store()
|
||||
kvh.release()
|
||||
acc_pipe.producer_commit(acc_st); acc_st.advance()
|
||||
acc_pipe.producer_tail(acc_st)
|
||||
|
||||
# SOFTMAX (identity: just load S, store as P, no scaling)
|
||||
if warp_idx < self.mma_warp_id:
|
||||
tmem.allocate(self.num_tmem_alloc_cols)
|
||||
tmem.wait_for_alloc()
|
||||
tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype)
|
||||
sfw_idx = tidx % (32 * len(self.epilogue_warp_id))
|
||||
|
||||
# S load
|
||||
tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype)
|
||||
tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS0)
|
||||
thr_load = tiled_tmem_load.get_slice(sfw_idx)
|
||||
tTMEM_LOADtS = thr_load.partition_S(tStS0)
|
||||
cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1]))
|
||||
tScS = qk_thr.partition_C(cS)
|
||||
tTMEM_LOADcS = thr_load.partition_D(tScS)
|
||||
|
||||
# P store
|
||||
p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width
|
||||
tStP_layout = cute.composition(tStS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32)))
|
||||
tStP0 = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStP_layout)
|
||||
tmem_store_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype)
|
||||
tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStP0)
|
||||
thr_store = tiled_tmem_store.get_slice(sfw_idx)
|
||||
tTMEM_STOREtP = thr_store.partition_D(tStP0)
|
||||
tScP_layout = cute.composition(tScS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32)))
|
||||
tScP = cute.make_tensor(tScS.iterator, tScP_layout)
|
||||
tTMEM_STOREcP = thr_store.partition_S(tScP)
|
||||
|
||||
for kt in range(self.n_kv_tiles):
|
||||
si_handle = s_cons.wait_and_advance()
|
||||
|
||||
# Load S from TMEM
|
||||
tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype)
|
||||
cute.copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS)
|
||||
cute.arch.fence_view_async_tmem_load()
|
||||
|
||||
# Identity: convert S to BF16 and store as P
|
||||
rP_words = cute.make_rmem_tensor(tTMEM_STOREcP.shape, self.qk_acc_dtype)
|
||||
rP_bf16 = cute.make_tensor(cute.recast_ptr(rP_words.iterator, dtype=self.q_dtype), tTMEM_LOADrS.layout)
|
||||
frg_cnt = 4
|
||||
frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt
|
||||
tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile))
|
||||
rP_bf16_frg = cute.logical_divide(rP_bf16, cute.make_layout(frg_tile))
|
||||
for j in range(frg_cnt):
|
||||
s_vec = tTMEM_LOADrS_frg[None, j].load()
|
||||
rP_bf16_frg[None, j].store(s_vec.to(self.q_dtype))
|
||||
|
||||
cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP)
|
||||
cute.arch.fence_view_async_tmem_store()
|
||||
si_handle.release()
|
||||
softmax_done_bar.arrive()
|
||||
|
||||
tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout)
|
||||
acc_cons_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage)
|
||||
c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id))
|
||||
c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp)
|
||||
acc_cons_st = utils.gemm.sm100.epilogue_tma_store(self, tidx, warp_idx, tma_c, tCtO_base, sC, tCgC, epi_tile, 0, const_expr(lambda x: x), (0,0,0), acc_cons_st, acc_pipe, c_pipe)
|
||||
c_pipe.producer_tail()
|
||||
tmem.relinquish_alloc_permit()
|
||||
tmem.free(tmem_ptr)
|
||||
|
||||
|
||||
def test():
|
||||
for n in [128, 256, 384]:
|
||||
torch.manual_seed(42)
|
||||
m, hd = 128, HEAD_DIM
|
||||
q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda')
|
||||
k = torch.randn(n, hd, 1, dtype=torch.bfloat16, device='cuda')
|
||||
v = torch.randn(n, hd, dtype=torch.bfloat16, device='cuda')
|
||||
v_kernel = v.unsqueeze(-1)
|
||||
c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda')
|
||||
|
||||
qf = q[:, :, 0].float()
|
||||
kf = k[:, :, 0].float()
|
||||
# Identity softmax reference: QK (no scale, no softmax) @ V
|
||||
ref = (qf @ kf.T).bfloat16().float() @ v.float()
|
||||
|
||||
mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q))
|
||||
mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k))
|
||||
mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel))
|
||||
mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c))
|
||||
stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
|
||||
|
||||
kernel = FmhaV3Diag(s_k=n)
|
||||
print(f'n={n}: Compiling...', flush=True)
|
||||
compiled = cute.compile(kernel, mQ, mK, mV, mC, stream)
|
||||
compiled(mQ, mK, mV, mC, stream)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
out = c[:, :, 0].float()
|
||||
cos = torch.nn.functional.cosine_similarity(
|
||||
out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)
|
||||
).item()
|
||||
n_tiles = n // 128
|
||||
print(f'DIAG n={n} ({n_tiles} tiles): cos {cos:.6f} {"PASS" if cos >= 0.99 else "FAIL"}')
|
||||
if cos < 0.99:
|
||||
print(f' out[0,:4]={out[0,:4].tolist()}')
|
||||
print(f' ref[0,:4]={ref[0,:4].tolist()}')
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
test()
|
||||
@@ -1,307 +0,0 @@
|
||||
"""
|
||||
FMHA Stage-C multi-tile diagnostic.
|
||||
Tests whether the combined K+V pipeline actually loads different KV tiles.
|
||||
Runs identity softmax (no rescale, no online softmax) to isolate the
|
||||
TMA + MMA pipeline from softmax logic.
|
||||
"""
|
||||
import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline
|
||||
from cutlass.cute.nvgpu import cpasync, tcgen05
|
||||
from cutlass import Float32, BFloat16, Int32, Boolean, const_expr
|
||||
from cutlass.utils import LayoutEnum
|
||||
from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset
|
||||
import cuda.bindings.driver as cuda
|
||||
import cutlass.torch as ct
|
||||
import math
|
||||
|
||||
HEAD_DIM = 64
|
||||
|
||||
|
||||
class FmhaV3Diag:
|
||||
"""Identity-softmax multi-tile with combined K+V barrier.
|
||||
If this fails, the problem is in TMA/MMA, not softmax."""
|
||||
def __init__(self, s_k=256):
|
||||
self.s_k = s_k
|
||||
self.n_kv_tiles = s_k // 128
|
||||
self.acc_dtype = Float32; self.qk_acc_dtype = Float32
|
||||
self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16
|
||||
self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1
|
||||
self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE
|
||||
self.epilogue_warp_id = (0,1,2,3); self.mma_warp_id = 4; self.tma_warp_id = 5
|
||||
self.threads_per_cta = 192; self.num_c_stage = 2
|
||||
self.kv_stage = 2; self.q_stage = 1; self.num_c_stage = 2
|
||||
|
||||
def _setup(self, qk_mma, pv_mma):
|
||||
qk_ik = cute.size(qk_mma.shape_mnk, mode=[2])
|
||||
self.qk_mma_tiler = (128, 128, qk_ik * 4)
|
||||
pv_ik = cute.size(pv_mma.shape_mnk, mode=[2])
|
||||
self.pv_mma_tiler = (128, HEAD_DIM, pv_ik * (128 // pv_ik))
|
||||
self.mma_tiler = self.qk_mma_tiler
|
||||
self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,))
|
||||
self.cta_tile_shape_mnk = (self.qk_mma_tiler[0]//cute.size(qk_mma.thr_id.shape), HEAD_DIM, self.qk_mma_tiler[2])
|
||||
self.c_layout = LayoutEnum.ROW_MAJOR
|
||||
self.epi_tile = utils.sm100.compute_epilogue_tile_shape(self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype)
|
||||
self.num_ab_stage = 1; self.num_acc_stage = 1
|
||||
self.q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, self.q_dtype, self.q_stage)
|
||||
self.k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, self.q_dtype, self.kv_stage)
|
||||
self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, self.kv_stage)
|
||||
self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2)
|
||||
self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1)
|
||||
qk_thr = qk_mma.get_slice(0); qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2])
|
||||
tStS = qk_thr.make_fragment_C(qk_as)
|
||||
pv_thr = pv_mma.get_slice(0); pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2])
|
||||
tOtO = pv_thr.make_fragment_C(pv_as)
|
||||
self.tmem_s0_offset = 0; self.tmem_p0_offset = 32
|
||||
p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width
|
||||
p_end = self.tmem_p0_offset + p_cols_fp32
|
||||
s_cols = self.qk_mma_tiler[1]
|
||||
o_after = max(s_cols, p_end)
|
||||
self.tmem_o0_offset = ((o_after + 31) // 32) * 32
|
||||
o_cols = find_tmem_tensor_col_offset(tOtO)
|
||||
total = self.tmem_o0_offset + o_cols
|
||||
self.num_tmem_alloc_cols = 1
|
||||
while self.num_tmem_alloc_cols < total:
|
||||
self.num_tmem_alloc_cols *= 2
|
||||
cta = cute.size(qk_mma.thr_id.shape)
|
||||
q_s = cute.slice_(self.q_smem_s,(None,None,None,0))
|
||||
k_s = cute.slice_(self.k_smem_s,(None,None,None,0))
|
||||
v_s = cute.slice_(self.v_smem_s,(None,None,None,0))
|
||||
self.q_tx_bytes = cute.size_in_bytes(self.q_dtype, q_s) * cta
|
||||
self.kv_tx_bytes = (cute.size_in_bytes(self.q_dtype, k_s) +
|
||||
cute.size_in_bytes(self.q_dtype, v_s)) * cta
|
||||
|
||||
@cute.jit
|
||||
def __call__(self, q, k, v, c, stream):
|
||||
self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype
|
||||
self.a_major = LayoutEnum.from_tensor(q).mma_major_mode()
|
||||
self.b_major = LayoutEnum.from_tensor(k).mma_major_mode()
|
||||
v_fmha = cute.make_tensor(
|
||||
v.iterator,
|
||||
cute.make_layout(
|
||||
(HEAD_DIM, self.s_k, 1),
|
||||
stride=(1, HEAD_DIM, HEAD_DIM * self.s_k),
|
||||
),
|
||||
)
|
||||
self.v_major = LayoutEnum.from_tensor(v_fmha).mma_major_mode()
|
||||
self.c_layout = LayoutEnum.from_tensor(c)
|
||||
qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, self.a_major, self.b_major, self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM)
|
||||
pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, self.qk_acc_dtype, self.cta_group, (128,HEAD_DIM), tcgen05.OperandSource.TMEM)
|
||||
self._setup(qk_mma, pv_mma)
|
||||
q_s = cute.slice_(self.q_smem_s,(None,None,None,0)); k_s = cute.slice_(self.k_smem_s,(None,None,None,0)); v_s = cute.slice_(self.v_smem_s,(None,None,None,0))
|
||||
tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id),q,q_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape)
|
||||
tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id),k,k_s,self.qk_mma_tiler,qk_mma,self.cluster_layout_vmnk.shape)
|
||||
tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v_fmha,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape)
|
||||
epi_s = cute.select(self.c_smem_s,mode=[0,1])
|
||||
tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile)
|
||||
self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.v_smem_s,self.p_tmem_s,self.c_smem_s,self.epi_tile).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream)
|
||||
|
||||
@cute.kernel
|
||||
def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, q_smem_s, k_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile):
|
||||
warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
|
||||
tidx,_,_ = cute.arch.thread_idx()
|
||||
if warp_idx == self.tma_warp_id:
|
||||
cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k); cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c)
|
||||
|
||||
@cute.struct
|
||||
class SS:
|
||||
q_bar: cute.struct.MemRange[cutlass.Int64, self.q_stage*2]
|
||||
kv_bar: cute.struct.MemRange[cutlass.Int64, self.kv_stage*2]
|
||||
s_bar: cute.struct.MemRange[cutlass.Int64, 2]
|
||||
acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage*2]
|
||||
tmem_dealloc: cutlass.Int64; holding: cutlass.Int32
|
||||
smem = utils.SmemAllocator(); st = smem.allocate(SS)
|
||||
|
||||
qp,qc = pipeline.PipelineTmaUmma.create(barrier_storage=st.q_bar.data_ptr(),num_stages=self.q_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.q_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants()
|
||||
kvp,kvc = pipeline.PipelineTmaUmma.create(barrier_storage=st.kv_bar.data_ptr(),num_stages=self.kv_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,1),tx_count=self.kv_tx_bytes,cta_layout_vmnk=cl_vmnk,defer_sync=True).make_participants()
|
||||
s_prod,s_cons = pipeline.PipelineUmmaAsync.create(barrier_storage=st.s_bar.data_ptr(),num_stages=1,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,32*len(self.epilogue_warp_id))).make_participants()
|
||||
softmax_done_bar = pipeline.NamedBarrier(barrier_id=3, num_threads=32 + 32*len(self.epilogue_warp_id))
|
||||
acc_pipe = pipeline.PipelineUmmaAsync.create(barrier_storage=st.acc_bar.data_ptr(),num_stages=self.num_acc_stage,producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread,len(self.epilogue_warp_id)),cta_layout_vmnk=cl_vmnk,defer_sync=True)
|
||||
tmem_bar = pipeline.NamedBarrier(barrier_id=2,num_threads=32*len((self.mma_warp_id,*self.epilogue_warp_id)))
|
||||
tmem = utils.TmemAllocator(st.holding.ptr,barrier_for_retrieve=tmem_bar,allocator_warp_id=self.epilogue_warp_id[0],is_two_cta=cute.size(qk_mma.thr_id.shape)==2,two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr)
|
||||
pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk,is_relaxed=True)
|
||||
|
||||
sQ = smem.allocate_tensor(element_type=self.q_dtype,layout=q_smem_s.outer,byte_alignment=128,swizzle=q_smem_s.inner)
|
||||
sK = smem.allocate_tensor(element_type=self.q_dtype,layout=k_smem_s.outer,byte_alignment=128,swizzle=k_smem_s.inner)
|
||||
sV = smem.allocate_tensor(element_type=self.q_dtype,layout=v_smem_s.outer,byte_alignment=128,swizzle=v_smem_s.inner)
|
||||
sC = smem.allocate_tensor(element_type=self.o_dtype,layout=c_smem_s.outer,byte_alignment=128,swizzle=c_smem_s.inner)
|
||||
|
||||
gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None))
|
||||
gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None))
|
||||
gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None))
|
||||
gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None))
|
||||
n_kv_tiles = self.n_kv_tiles
|
||||
|
||||
qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0)
|
||||
tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK)
|
||||
tCgV = pv_thr.partition_B(gV); tCgC = pv_thr.partition_C(gC)
|
||||
a_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,0,None,0)).shape)
|
||||
tAsQ,tAgQ = cpasync.tma_partition(tma_q,0,a_lay,cute.group_modes(sQ,0,3),cute.group_modes(tCgQ,0,3))
|
||||
b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape)
|
||||
tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3))
|
||||
tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3))
|
||||
tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,None,None,None,None,None,None,None)]; tVgV = tVgV[(None,None,None,None,None,None,None,None)]
|
||||
|
||||
tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK)
|
||||
tCrV = pv_mma.make_fragment_B(sV)
|
||||
|
||||
qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2])
|
||||
tStS = qk_thr.make_fragment_C(qk_as)
|
||||
tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout)
|
||||
pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2])
|
||||
tOtO = pv_thr.make_fragment_C(pv_as)
|
||||
tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout)
|
||||
|
||||
tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer)
|
||||
tOrP_base = pv_thr.make_fragment_A(tP)
|
||||
tOrP = tOrP_base[(None,None,None,0)]
|
||||
tOrP0 = cute.make_tensor(
|
||||
tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset,
|
||||
tOrP.layout)
|
||||
|
||||
tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage))
|
||||
tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage))
|
||||
pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk)
|
||||
|
||||
# TMA LOAD (combined K+V barrier, Int32(kt) for GMEM coord)
|
||||
if warp_idx == self.tma_warp_id:
|
||||
qp.reset(); qh = qp.acquire_and_advance()
|
||||
cute.copy(tma_q, tAgQ[(None, Int32(0))], tAsQ[(None, qh.index)], tma_bar_ptr=qh.barrier)
|
||||
qp.tail()
|
||||
kvp.reset(); pk = kvp.try_acquire()
|
||||
# Force runtime Int32 (not literal) — option 3 from CUTLASS LLM
|
||||
kv_coord = Int32(0 + 0)
|
||||
for kt in cutlass.range(self.n_kv_tiles, unroll=1):
|
||||
kvh = kvp.acquire_and_advance(pk)
|
||||
cute.copy(tma_k, tBgK[None, None, None, None, kv_coord, None, None, None], tBsK[(None, kvh.index)], tma_bar_ptr=kvh.barrier)
|
||||
cute.copy(tma_v, tVgV[None, None, None, None, kv_coord, None, None, None], tVsV[(None, kvh.index)], tma_bar_ptr=kvh.barrier)
|
||||
kv_coord += 1
|
||||
pk = cutlass.Boolean(1)
|
||||
kvp.tail()
|
||||
|
||||
# MMA (combined K+V slot)
|
||||
if warp_idx == self.mma_warp_id:
|
||||
tmem.wait_for_alloc()
|
||||
qc.reset(); qh = qc.wait_and_advance(); qh.release()
|
||||
kvc.reset(); pk = kvc.try_wait()
|
||||
acc_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_stage)
|
||||
acc_pipe.producer_acquire(acc_st)
|
||||
for kt in range(self.n_kv_tiles):
|
||||
kvh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1)
|
||||
sh = s_prod.acquire_and_advance()
|
||||
qk_mma.set(tcgen05.Field.ACCUMULATE, False)
|
||||
for kb in cutlass.range(cute.size(tCrQ, mode=[2]), unroll_full=True):
|
||||
cute.gemm(qk_mma, tStS0, tCrQ[(None,None,kb,0)], tCrK[(None,None,kb,kvh.index)], tStS0)
|
||||
qk_mma.set(tcgen05.Field.ACCUMULATE, True)
|
||||
cute.arch.fence_view_async_tmem_store()
|
||||
sh.commit()
|
||||
softmax_done_bar.arrive_and_wait()
|
||||
pv_mma.set(tcgen05.Field.ACCUMULATE, kt != 0)
|
||||
for kb in cutlass.range(cute.size(tOrP0, mode=[2]), unroll_full=True):
|
||||
cute.gemm(pv_mma, tOtO0, tOrP0[(None,None,kb)], tCrV[(None,None,kb,kvh.index)], tOtO0)
|
||||
pv_mma.set(tcgen05.Field.ACCUMULATE, True)
|
||||
cute.arch.fence_view_async_tmem_store()
|
||||
kvh.release()
|
||||
acc_pipe.producer_commit(acc_st); acc_st.advance()
|
||||
acc_pipe.producer_tail(acc_st)
|
||||
|
||||
# SOFTMAX (identity: just load S, store as P, no scaling)
|
||||
if warp_idx < self.mma_warp_id:
|
||||
tmem.allocate(self.num_tmem_alloc_cols)
|
||||
tmem.wait_for_alloc()
|
||||
tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype)
|
||||
sfw_idx = tidx % (32 * len(self.epilogue_warp_id))
|
||||
|
||||
# S load
|
||||
tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype)
|
||||
tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS0)
|
||||
thr_load = tiled_tmem_load.get_slice(sfw_idx)
|
||||
tTMEM_LOADtS = thr_load.partition_S(tStS0)
|
||||
cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1]))
|
||||
tScS = qk_thr.partition_C(cS)
|
||||
tTMEM_LOADcS = thr_load.partition_D(tScS)
|
||||
|
||||
# P store
|
||||
p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width
|
||||
tStP_layout = cute.composition(tStS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32)))
|
||||
tStP0 = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStP_layout)
|
||||
tmem_store_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype)
|
||||
tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStP0)
|
||||
thr_store = tiled_tmem_store.get_slice(sfw_idx)
|
||||
tTMEM_STOREtP = thr_store.partition_D(tStP0)
|
||||
tScP_layout = cute.composition(tScS.layout, cute.make_layout((self.pv_mma_tiler[0], p_cols_fp32)))
|
||||
tScP = cute.make_tensor(tScS.iterator, tScP_layout)
|
||||
tTMEM_STOREcP = thr_store.partition_S(tScP)
|
||||
|
||||
for kt in range(self.n_kv_tiles):
|
||||
si_handle = s_cons.wait_and_advance()
|
||||
|
||||
# Load S from TMEM
|
||||
tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype)
|
||||
cute.copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS)
|
||||
cute.arch.fence_view_async_tmem_load()
|
||||
|
||||
# Identity: convert S to BF16 and store as P
|
||||
rP_words = cute.make_rmem_tensor(tTMEM_STOREcP.shape, self.qk_acc_dtype)
|
||||
rP_bf16 = cute.make_tensor(cute.recast_ptr(rP_words.iterator, dtype=self.q_dtype), tTMEM_LOADrS.layout)
|
||||
frg_cnt = 4
|
||||
frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt
|
||||
tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile))
|
||||
rP_bf16_frg = cute.logical_divide(rP_bf16, cute.make_layout(frg_tile))
|
||||
for j in range(frg_cnt):
|
||||
s_vec = tTMEM_LOADrS_frg[None, j].load()
|
||||
rP_bf16_frg[None, j].store(s_vec.to(self.q_dtype))
|
||||
|
||||
cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP)
|
||||
cute.arch.fence_view_async_tmem_store()
|
||||
si_handle.release()
|
||||
softmax_done_bar.arrive()
|
||||
|
||||
tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout)
|
||||
acc_cons_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage)
|
||||
c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id))
|
||||
c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp)
|
||||
acc_cons_st = utils.gemm.sm100.epilogue_tma_store(self, tidx, warp_idx, tma_c, tCtO_base, sC, tCgC, epi_tile, 0, const_expr(lambda x: x), (0,0,0), acc_cons_st, acc_pipe, c_pipe)
|
||||
c_pipe.producer_tail()
|
||||
tmem.relinquish_alloc_permit()
|
||||
tmem.free(tmem_ptr)
|
||||
|
||||
|
||||
def test():
|
||||
for n in [128, 256, 384]:
|
||||
torch.manual_seed(42)
|
||||
m, hd = 128, HEAD_DIM
|
||||
q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda')
|
||||
k = torch.randn(n, hd, 1, dtype=torch.bfloat16, device='cuda')
|
||||
v = torch.randn(n, hd, dtype=torch.bfloat16, device='cuda')
|
||||
v_kernel = v.unsqueeze(-1)
|
||||
c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda')
|
||||
|
||||
qf = q[:, :, 0].float()
|
||||
kf = k[:, :, 0].float()
|
||||
# Identity softmax reference: QK (no scale, no softmax) @ V
|
||||
ref = (qf @ kf.T).bfloat16().float() @ v.float()
|
||||
|
||||
mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q))
|
||||
mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k))
|
||||
mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel))
|
||||
mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c))
|
||||
stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
|
||||
|
||||
kernel = FmhaV3Diag(s_k=n)
|
||||
print(f'n={n}: Compiling...', flush=True)
|
||||
compiled = cute.compile(kernel, mQ, mK, mV, mC, stream)
|
||||
compiled(mQ, mK, mV, mC, stream)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
out = c[:, :, 0].float()
|
||||
cos = torch.nn.functional.cosine_similarity(
|
||||
out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)
|
||||
).item()
|
||||
n_tiles = n // 128
|
||||
print(f'DIAG n={n} ({n_tiles} tiles): cos {cos:.6f} {"PASS" if cos >= 0.99 else "FAIL"}')
|
||||
if cos < 0.99:
|
||||
print(f' out[0,:4]={out[0,:4].tolist()}')
|
||||
print(f' ref[0,:4]={ref[0,:4].tolist()}')
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
test()
|
||||
@@ -1,92 +0,0 @@
|
||||
"""Test: Validate SiLU in registers (Step 1 of fused SwiGLU).
|
||||
|
||||
Compiles the fused kernel with fused_swiglu=True, runs it, and compares
|
||||
the BF16 output with PyTorch SiLU applied to the standard L1 GEMM output.
|
||||
"""
|
||||
import torch
|
||||
import sys
|
||||
sys.path.insert(0, '/root/dsv4-nvfp4-workspace/kernel')
|
||||
|
||||
from dsv4.ops.quantize import (
|
||||
quantize_weight_to_nvfp4,
|
||||
quantize_activation_nvfp4,
|
||||
)
|
||||
from dsv4.ops.layouts import (
|
||||
make_b_k_major,
|
||||
assemble_scales_2d_side,
|
||||
assemble_scales_3d_side,
|
||||
)
|
||||
from dsv4.ops.gemm_runner import (
|
||||
run_nvfp4_grouped_gemm,
|
||||
run_fused_swiglu_grouped_gemm,
|
||||
warmup_compilation,
|
||||
)
|
||||
|
||||
|
||||
def test_silu_step1():
|
||||
device = "cuda"
|
||||
num_experts = 4
|
||||
hidden = 512
|
||||
intermediate = 256
|
||||
num_tokens = 32
|
||||
|
||||
torch.manual_seed(42)
|
||||
x = torch.randn(num_tokens, hidden, dtype=torch.bfloat16, device=device)
|
||||
l1_w = torch.randn(num_experts, 2 * intermediate, hidden, dtype=torch.bfloat16, device=device)
|
||||
|
||||
l1_fp4_list, l1_sf_list, l1_gs_list = [], [], []
|
||||
for e in range(num_experts):
|
||||
w_fp4, w_sf, w_gs = quantize_weight_to_nvfp4(l1_w[e].T)
|
||||
l1_fp4_list.append(w_fp4)
|
||||
l1_sf_list.append(w_sf)
|
||||
l1_gs_list.append(w_gs)
|
||||
|
||||
l1_mat_b = make_b_k_major(torch.stack(l1_fp4_list))
|
||||
l1_scale_b = assemble_scales_3d_side(l1_sf_list)
|
||||
l1_gs = torch.tensor(l1_gs_list, dtype=torch.float32, device=device)
|
||||
|
||||
gs_val = x.abs().max().item() / (6.0 * 448.0)
|
||||
x_fp4, x_sf = quantize_activation_nvfp4(x, gs_val)
|
||||
tokens_per_expert = [num_tokens // num_experts] * num_experts
|
||||
scale_a = assemble_scales_2d_side([x_sf[i*tpe:(i+1)*tpe] for i, tpe in enumerate(tokens_per_expert)])
|
||||
expert_offsets = torch.tensor(
|
||||
[sum(tokens_per_expert[:e+1]) for e in range(num_experts)],
|
||||
dtype=torch.int32, device=device,
|
||||
)
|
||||
global_scale_a = torch.full((num_experts,), gs_val, dtype=torch.float32, device=device)
|
||||
|
||||
warmup_compilation(num_experts, hidden // 2, (2 * intermediate) // 2, device)
|
||||
|
||||
# 1. Standard L1 GEMM (no SiLU)
|
||||
out_bf16 = run_nvfp4_grouped_gemm(
|
||||
mat_a=x_fp4, mat_b=l1_mat_b,
|
||||
scale_a=scale_a, scale_b=l1_scale_b,
|
||||
expert_offsets=expert_offsets,
|
||||
global_scale_a=global_scale_a, global_scale_b=l1_gs,
|
||||
)
|
||||
silu_ref = torch.nn.functional.silu(out_bf16)
|
||||
print(f"Standard L1 output: shape={out_bf16.shape}, amax={out_bf16.abs().amax().item():.4f}")
|
||||
print(f"PyTorch SiLU ref: amax={silu_ref.abs().amax().item():.4f}")
|
||||
|
||||
# 2. Fused kernel with SiLU in registers
|
||||
print("\nCompiling fused kernel (first time, may take a while)...")
|
||||
out_fused = run_fused_swiglu_grouped_gemm(
|
||||
mat_a=x_fp4, mat_b=l1_mat_b,
|
||||
scale_a=scale_a, scale_b=l1_scale_b,
|
||||
expert_offsets=expert_offsets,
|
||||
global_scale_a=global_scale_a, global_scale_b=l1_gs,
|
||||
)
|
||||
print(f"Fused SiLU output: shape={out_fused.shape}, amax={out_fused.abs().amax().item():.4f}")
|
||||
|
||||
# 3. Compare
|
||||
diff = (out_fused - silu_ref).float()
|
||||
rel_err = diff.norm() / silu_ref.float().norm()
|
||||
max_err = diff.abs().max()
|
||||
print(f"\n=== Results ===")
|
||||
print(f"Relative error: {rel_err.item():.6f}")
|
||||
print(f"Max abs error: {max_err.item():.6f}")
|
||||
print(f"PASS" if rel_err.item() < 0.1 else "FAIL (tolerance: 0.1 for NVFP4 quant noise)")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_silu_step1()
|
||||
@@ -1,144 +0,0 @@
|
||||
"""Test: Verify weight interleave produces correct gate/up pairs in GEMM output.
|
||||
|
||||
Stage 1 validation: If interleaved weights produce the same GEMM result
|
||||
as non-interleaved weights (after de-interleaving the output), the
|
||||
interleave is correct and the fused epilogue can safely assume gate/up
|
||||
pairs are adjacent in registers.
|
||||
"""
|
||||
import torch
|
||||
import sys
|
||||
sys.path.insert(0 = '/root/dsv4-nvfp4-workspace/kernel') # FIXME
|
||||
|
||||
from dsv4.ops.quantize import (
|
||||
quantize_to_nvfp4,
|
||||
quantize_activation_nvfp4,
|
||||
quantize_weight_to_nvfp4,
|
||||
)
|
||||
from dsv4.ops.layouts import (
|
||||
interleave_l1_weights,
|
||||
deinterleave_l1_weights,
|
||||
make_b_k_major,
|
||||
assemble_scales_2d_side,
|
||||
assemble_scales_3d_side,
|
||||
)
|
||||
from dsv4.ops.gemm_runner import (
|
||||
run_nvfp4_grouped_gemm,
|
||||
)
|
||||
|
||||
|
||||
def test_interleave_correctness():
|
||||
"""Verify that interleaving weights and de-interleaving the GEMM output
|
||||
gives the same result as non-interleaved weights.
|
||||
"""
|
||||
device = "cuda"
|
||||
num_experts = 4
|
||||
hidden = 512
|
||||
intermediate = 256
|
||||
num_tokens = 32
|
||||
|
||||
# Create random BF16 input
|
||||
x = torch.randn(num_tokens, hidden, dtype=torch.bfloat16, device=device)
|
||||
|
||||
# Create random BF16 weights for gate and up
|
||||
gate_w = torch.randn(num_experts, intermediate, hidden, dtype=torch.bfloat16, device=device)
|
||||
up_w = torch.randn(num_experts, intermediate, hidden, dtype=torch.bfloat16, device=device)
|
||||
|
||||
# === Path A: Non-interleaved (current production path) ===
|
||||
# Fuse gate+up: (E, 2*intermediate, hidden)
|
||||
l1_bf16 = torch.cat([gate_w, up_w], dim=1) # (E, 6144, 7168) → (E, 2*inter, hidden)
|
||||
|
||||
# Quantize weights
|
||||
l1_fp4_list, l1_sf_list, l1_gs_list = [], [], []
|
||||
for e in range(num_experts):
|
||||
w_fp4, w_sf, w_gs = quantize_weight_to_nvfp4(l1_bf16[e].T) # (K, N)
|
||||
l1_fp4_list.append(w_fp4)
|
||||
l1_sf_list.append(w_sf)
|
||||
l1_gs_list.append(w_gs)
|
||||
|
||||
# Stack and convert
|
||||
l1_mat_b = make_b_k_major(torch.stack(l1_fp4_list))
|
||||
l1_scale_b = assemble_scales_3d_side(l1_sf_list)
|
||||
l1_gs = torch.tensor(l1_gs_list, dtype=torch.float32, device=device)
|
||||
|
||||
# Quantize activation
|
||||
gs_val = x.abs().max().item() / (6.0 * 448.0)
|
||||
x_fp4, x_sf = quantize_activation_nvfp4(x, gs_val)
|
||||
|
||||
# Assemble scales
|
||||
tokens_per_expert = [num_tokens // num_experts] * num_experts
|
||||
scale_a = assemble_scales_2d_side([x_sf[i*tpe:(i+1)*tpe] for i, tpe in enumerate(tokens_per_expert)])
|
||||
|
||||
expert_offsets = torch.tensor(
|
||||
[sum(tokens_per_expert[:e+1]) for e in range(num_experts)],
|
||||
dtype=torch.int32, device=device,
|
||||
)
|
||||
global_scale_a = torch.full((num_experts,), gs_val, dtype=torch.float32, device=device)
|
||||
|
||||
# Run GEMM
|
||||
out_a = run_nvfp4_grouped_gemm(
|
||||
mat_a=x_fp4, mat_b=l1_mat_b,
|
||||
scale_a=scale_a, scale_b=l1_scale_b,
|
||||
expert_offsets=expert_offsets,
|
||||
global_scale_a=global_scale_a, global_scale_b=l1_gs,
|
||||
)
|
||||
# out_a: (num_tokens, 2*intermediate) BF16
|
||||
# gate = out_a[:, :intermediate], up = out_a[:, intermediate:]
|
||||
gate_a = out_a[:, :intermediate]
|
||||
up_a = out_a[:, intermediate:]
|
||||
result_a = torch.nn.functional.silu(gate_a) * up_a # SwiGLU result
|
||||
|
||||
# === Path B: Interleaved weights ===
|
||||
# Quantize gate and up separately, then interleave
|
||||
gate_fp4, gate_sf, gate_gs = [], [], []
|
||||
up_fp4, up_sf, up_gs = [], [], []
|
||||
for e in range(num_experts):
|
||||
g4, gs4, gg4 = quantize_weight_to_nvfp4(gate_w[e].T)
|
||||
u4, us4, ug4 = quantize_weight_to_nvfp4(up_w[e].T)
|
||||
gate_fp4.append(g4)
|
||||
gate_sf.append(gs4)
|
||||
gate_gs.append(gg4)
|
||||
up_fp4.append(u4)
|
||||
up_sf.append(us4)
|
||||
up_gs.append(ug4)
|
||||
|
||||
# Fuse and interleave
|
||||
gate_stacked = torch.stack(gate_fp4) # (E, K_packed, N/2)
|
||||
up_stacked = torch.stack(up_fp4) # (E, K_packed, N/2)
|
||||
l1_bf16_fp4 = torch.cat([gate_stacked, up_stacked], dim=2) # (E, K, N) non-interleaved
|
||||
l1_interleaved = interleave_l1_weights(l1_bf16_fp4) # interleaved
|
||||
|
||||
# Make K-major
|
||||
l1_mat_b_int = make_b_k_major(l1_interleaved)
|
||||
|
||||
# Scale assembly: gate and up scales combined
|
||||
l1_scale_b_int = assemble_scales_3d_side(gate_sf + up_sf) # interleave scales too?
|
||||
# Actually, the scale interleaving needs to match the weight interleaving.
|
||||
# This is more complex. For Stage 1, let's use a simpler approach.
|
||||
|
||||
# Actually, for the interleaved path to produce the same GEMM output,
|
||||
# we need the SFB to also be interleaved to match.
|
||||
# The GEMM is: A (M, K) x B (E, K, N) = C (M, N)
|
||||
# If we permute the N dimension of B, we permute the N dimension of C.
|
||||
# So the output columns are also interleaved.
|
||||
|
||||
# For this test, we just verify that the interleaved GEMM output,
|
||||
# when de-interleaved, matches the non-interleaved output.
|
||||
|
||||
# But the SFB (scale_b) must match the interleaved B.
|
||||
# The B tensor has its N columns interleaved, so the SFB must be
|
||||
# interleaved in the same way.
|
||||
|
||||
# SFB for interleaved B: we need to interleave the scales too.
|
||||
# Since scales are per-(K_sf, N) and we're interleaving N at granularity 4 FP4 cols,
|
||||
# the scales need to be interleaved at the same granularity.
|
||||
|
||||
# This is getting complex. Let me simplify: just test the interleave
|
||||
# function itself, not the full GEMM.
|
||||
|
||||
print("Interleave/deinterleave round-trip: PASSED (tested in bridge.py)")
|
||||
print("Full GEMM interleave test: SKIPPED (requires SFB interleaving)")
|
||||
print("Stage 1 kernel test will validate the full pipeline")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_interleave_correctness()
|
||||
@@ -1,137 +0,0 @@
|
||||
"""Test: Verify that interleaved L1 weights produce the same GEMM result.
|
||||
|
||||
The key insight: we quantize gate+up TOGETHER (same as non-interleaved),
|
||||
then interleave the ALREADY-QUANTIZED FP4 bytes and scales in the N dimension.
|
||||
This preserves quantization fidelity.
|
||||
"""
|
||||
import torch
|
||||
import sys
|
||||
sys.path.insert(0, '/root/dsv4-nvfp4-workspace/kernel')
|
||||
|
||||
from dsv4.ops.quantize import (
|
||||
quantize_weight_to_nvfp4,
|
||||
quantize_activation_nvfp4,
|
||||
)
|
||||
from dsv4.ops.layouts import (
|
||||
interleave_l1_weights,
|
||||
make_b_k_major,
|
||||
assemble_scales_2d_side,
|
||||
assemble_scales_3d_side,
|
||||
)
|
||||
from dsv4.ops.gemm_runner import (
|
||||
run_nvfp4_grouped_gemm,
|
||||
warmup_compilation,
|
||||
)
|
||||
|
||||
|
||||
def interleave_sfb(raw_scales, granularity_bf16=8):
|
||||
"""Interleave gate/up scales at the same granularity as the FP4 weights.
|
||||
|
||||
raw_scales: list of (K_sf, N) float8_e4m3fn tensors where N = 2*intermediate_sf
|
||||
Returns: list of (K_sf, N) float8_e4m3fn with interleaved gate/up
|
||||
"""
|
||||
g = granularity_bf16 // 2 # 4 FP4 scale columns per group
|
||||
result = []
|
||||
for sf in raw_scales:
|
||||
K_sf, N = sf.shape
|
||||
N_half = N // 2
|
||||
gate = sf[:, :N_half].reshape(K_sf, N_half // g, g)
|
||||
up = sf[:, N_half:].reshape(K_sf, N_half // g, g)
|
||||
interleaved = torch.stack([gate, up], dim=2).reshape(K_sf, N)
|
||||
result.append(interleaved)
|
||||
return result
|
||||
|
||||
|
||||
def test_interleave_gemm():
|
||||
device = "cuda"
|
||||
num_experts = 4
|
||||
hidden = 512
|
||||
intermediate = 256
|
||||
num_tokens = 32
|
||||
|
||||
torch.manual_seed(42)
|
||||
x = torch.randn(num_tokens, hidden, dtype=torch.bfloat16, device=device)
|
||||
gate_w = torch.randn(num_experts, intermediate, hidden, dtype=torch.bfloat16, device=device)
|
||||
up_w = torch.randn(num_experts, intermediate, hidden, dtype=torch.bfloat16, device=device)
|
||||
|
||||
# === Path A: Non-interleaved ===
|
||||
l1_bf16 = torch.cat([gate_w, up_w], dim=1) # (E, 2*inter, hidden)
|
||||
l1_fp4_list, l1_sf_list, l1_gs_list = [], [], []
|
||||
for e in range(num_experts):
|
||||
w_fp4, w_sf, w_gs = quantize_weight_to_nvfp4(l1_bf16[e].T)
|
||||
l1_fp4_list.append(w_fp4)
|
||||
l1_sf_list.append(w_sf)
|
||||
l1_gs_list.append(w_gs)
|
||||
|
||||
l1_mat_b = make_b_k_major(torch.stack(l1_fp4_list))
|
||||
l1_scale_b = assemble_scales_3d_side(l1_sf_list)
|
||||
l1_gs = torch.tensor(l1_gs_list, dtype=torch.float32, device=device)
|
||||
|
||||
gs_val = x.abs().max().item() / (6.0 * 448.0)
|
||||
x_fp4, x_sf = quantize_activation_nvfp4(x, gs_val)
|
||||
tokens_per_expert = [num_tokens // num_experts] * num_experts
|
||||
scale_a = assemble_scales_2d_side([x_sf[i*tpe:(i+1)*tpe] for i, tpe in enumerate(tokens_per_expert)])
|
||||
expert_offsets = torch.tensor(
|
||||
[sum(tokens_per_expert[:e+1]) for e in range(num_experts)],
|
||||
dtype=torch.int32, device=device,
|
||||
)
|
||||
global_scale_a = torch.full((num_experts,), gs_val, dtype=torch.float32, device=device)
|
||||
|
||||
warmup_compilation(num_experts, hidden // 2, (2 * intermediate) // 2, device)
|
||||
out_a = run_nvfp4_grouped_gemm(
|
||||
mat_a=x_fp4, mat_b=l1_mat_b,
|
||||
scale_a=scale_a, scale_b=l1_scale_b,
|
||||
expert_offsets=expert_offsets,
|
||||
global_scale_a=global_scale_a, global_scale_b=l1_gs,
|
||||
)
|
||||
|
||||
# === Path B: Interleaved (quantize together, interleave after) ===
|
||||
# Use the SAME quantized weights, just interleave the N dimension
|
||||
l1_stacked = torch.stack(l1_fp4_list) # (E, K, N)
|
||||
l1_interleaved = interleave_l1_weights(l1_stacked)
|
||||
l1_mat_b_int = make_b_k_major(l1_interleaved)
|
||||
|
||||
# Interleave scales to match
|
||||
l1_sf_interleaved = interleave_sfb(l1_sf_list)
|
||||
l1_scale_b_int = assemble_scales_3d_side(l1_sf_interleaved)
|
||||
# Global scales are the same (quantized together)
|
||||
|
||||
out_b = run_nvfp4_grouped_gemm(
|
||||
mat_a=x_fp4, mat_b=l1_mat_b_int,
|
||||
scale_a=scale_a, scale_b=l1_scale_b_int,
|
||||
expert_offsets=expert_offsets,
|
||||
global_scale_a=global_scale_a, global_scale_b=l1_gs,
|
||||
)
|
||||
|
||||
# De-interleave out_b BF16 to match out_a layout
|
||||
N = out_b.shape[1]
|
||||
N_half = N // 2
|
||||
g = 8 # granularity in BF16
|
||||
out_b_reshaped = out_b.reshape(num_tokens, N // (2 * g), 2, g)
|
||||
gate_b = out_b_reshaped[:, :, 0, :].reshape(num_tokens, N_half)
|
||||
up_b = out_b_reshaped[:, :, 1, :].reshape(num_tokens, N_half)
|
||||
out_b_deint = torch.cat([gate_b, up_b], dim=1)
|
||||
|
||||
diff = (out_a - out_b_deint).float()
|
||||
rel_err = diff.norm() / out_a.float().norm()
|
||||
max_err = diff.abs().max()
|
||||
|
||||
print(f"Non-interleaved vs interleaved+deinterleaved:")
|
||||
print(f" Relative error: {rel_err.item():.6f}")
|
||||
print(f" Max abs error: {max_err.item():.6f}")
|
||||
print(f" PASS" if rel_err.item() < 0.01 else " FAIL")
|
||||
|
||||
# Apply SiLU and compare
|
||||
gate_a = out_a[:, :intermediate]
|
||||
up_a = out_a[:, intermediate:]
|
||||
result_a = torch.nn.functional.silu(gate_a) * up_a
|
||||
result_b = torch.nn.functional.silu(gate_b) * up_b
|
||||
|
||||
diff2 = (result_a - result_b).float()
|
||||
rel_err2 = diff2.norm() / result_a.float().norm()
|
||||
print(f" SiLU result error: {rel_err2.item():.6f}")
|
||||
print(f" SiLU PASS" if rel_err2.item() < 0.01 else " SiLU FAIL")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_interleave_gemm()
|
||||
@@ -1,85 +0,0 @@
|
||||
"""Tests for layer schedule — pure data, no kernels, no tensors."""
|
||||
|
||||
from dsv4.model.config import DSV4Config
|
||||
from dsv4.model.layer_schedule import (
|
||||
AttentionType, FFNType, RouterMode,
|
||||
LayerSpec, build_schedule, validate_schedule,
|
||||
)
|
||||
|
||||
|
||||
def test_flash_schedule():
|
||||
config = DSV4Config.flash()
|
||||
schedule = build_schedule(config)
|
||||
validate_schedule(schedule, config)
|
||||
|
||||
assert len(schedule) == 43
|
||||
|
||||
# First two layers: SWA + hash routing
|
||||
assert schedule[0].attn == AttentionType.SWA
|
||||
assert schedule[1].attn == AttentionType.SWA
|
||||
assert schedule[0].router_mode == RouterMode.HASH
|
||||
assert schedule[1].router_mode == RouterMode.HASH
|
||||
|
||||
# Layer 2: CSA + hash routing (last hash layer)
|
||||
assert schedule[2].attn == AttentionType.CSA
|
||||
assert schedule[2].router_mode == RouterMode.HASH
|
||||
|
||||
# Layer 3: HCA + dense routing (first dense layer)
|
||||
assert schedule[3].attn == AttentionType.HCA
|
||||
assert schedule[3].router_mode == RouterMode.DENSE
|
||||
|
||||
# Alternation continues
|
||||
assert schedule[4].attn == AttentionType.CSA
|
||||
assert schedule[5].attn == AttentionType.HCA
|
||||
|
||||
# All layers are MoE
|
||||
for spec in schedule:
|
||||
assert spec.ffn == FFNType.MOE
|
||||
|
||||
|
||||
def test_pro_schedule():
|
||||
config = DSV4Config.pro()
|
||||
schedule = build_schedule(config)
|
||||
validate_schedule(schedule, config)
|
||||
|
||||
assert len(schedule) == 61
|
||||
|
||||
# First two layers: HCA + hash routing
|
||||
assert schedule[0].attn == AttentionType.HCA
|
||||
assert schedule[1].attn == AttentionType.HCA
|
||||
assert schedule[0].router_mode == RouterMode.HASH
|
||||
|
||||
# Layer 2: CSA + hash routing
|
||||
assert schedule[2].attn == AttentionType.CSA
|
||||
assert schedule[2].router_mode == RouterMode.HASH
|
||||
|
||||
# Layer 3: HCA + dense routing
|
||||
assert schedule[3].attn == AttentionType.HCA
|
||||
assert schedule[3].router_mode == RouterMode.DENSE
|
||||
|
||||
|
||||
def test_layer_spec_frozen():
|
||||
"""LayerSpec is frozen — mutation should raise."""
|
||||
config = DSV4Config.flash()
|
||||
spec = build_schedule(config)[0]
|
||||
try:
|
||||
spec.attn = AttentionType.HCA
|
||||
assert False, "should have raised"
|
||||
except AttributeError:
|
||||
pass
|
||||
|
||||
|
||||
def test_schedule_indices_match():
|
||||
"""Each LayerSpec.layer_idx matches its position in the list."""
|
||||
config = DSV4Config.flash()
|
||||
schedule = build_schedule(config)
|
||||
for i, spec in enumerate(schedule):
|
||||
assert spec.layer_idx == i
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_flash_schedule()
|
||||
test_pro_schedule()
|
||||
test_layer_spec_frozen()
|
||||
test_schedule_indices_match()
|
||||
print("All schedule tests passed")
|
||||
@@ -1,217 +0,0 @@
|
||||
"""Unit tests for DSV4 Router — dense and hash modes.
|
||||
|
||||
Test strategy:
|
||||
Each kernel has a closed-form mathematical spec. The unit test computes
|
||||
the spec in one expression in FP32 (PyTorch) and compares against the
|
||||
kernel output. This is not "a PyTorch reference implementation" — it's
|
||||
the math. Compare against that. No "ref/" file, no second implementation
|
||||
drift, no two debug streams.
|
||||
|
||||
The oracle is the same five lines of math as the kernel spec, written
|
||||
declaratively. Compare against that.
|
||||
|
||||
DO NOT RUN THESE TESTS — Carmine is actively testing Stage C.
|
||||
Write the tests, commit them, they'll be run later.
|
||||
|
||||
Tie-breaking: When two scores are exactly equal, torch.topk and the kernel
|
||||
may pick different indices. Use the same tie-break rule: lower index wins
|
||||
on ties. If the test fails on tie-breaking, fix the kernel, not the test.
|
||||
"""
|
||||
|
||||
import torch
|
||||
import math
|
||||
|
||||
|
||||
def test_fused_activation_topk(N=64, E=256, k=6, seed=42):
|
||||
"""Test the fused activation + top-k kernel against the math spec.
|
||||
|
||||
Oracle:
|
||||
logits = X @ W (FP32)
|
||||
act = sqrt(softplus(logits))
|
||||
score = act + bias
|
||||
ids = argtopk(score, k) with lower-index tie-break
|
||||
raw_w = gather(act, ids)
|
||||
topk_w = raw_w / sum(raw_w) * scaling
|
||||
"""
|
||||
torch.manual_seed(seed)
|
||||
scaling = 2.5
|
||||
|
||||
logits = torch.randn(N, E, dtype=torch.float32, device='cuda')
|
||||
e_bias = torch.randn(E, dtype=torch.float32, device='cuda') * 0.01
|
||||
|
||||
# Oracle — the math, one expression at a time
|
||||
act = torch.sqrt(torch.nn.functional.softplus(logits))
|
||||
score = act + e_bias
|
||||
# torch.topk tie-breaking: picks lower index on ties (matches our kernel)
|
||||
topk_result = score.topk(k, dim=-1)
|
||||
ids = topk_result.indices
|
||||
raw_w = act.gather(-1, ids)
|
||||
w = raw_w / raw_w.sum(-1, keepdim=True) * scaling
|
||||
|
||||
# Kernel under test:
|
||||
from dsv4.kernels.router._activation_topk import run_fused_activation_topk
|
||||
out_w = torch.empty(N, k, dtype=torch.float32, device='cuda')
|
||||
out_ids = torch.empty(N, k, dtype=torch.int32, device='cuda')
|
||||
run_fused_activation_topk(logits, e_bias, scaling, k, out_w, out_ids)
|
||||
|
||||
# Verify
|
||||
assert (out_ids == ids).all(), f"top-k indices mismatch"
|
||||
torch.testing.assert_close(out_w, w, atol=1e-4, rtol=1e-3)
|
||||
|
||||
|
||||
def test_fused_activation_topk_decode_shapes():
|
||||
"""Test the activation+topk kernel at decode-relevant N values."""
|
||||
for N in [1, 4, 16, 64]:
|
||||
test_fused_activation_topk(N=N, E=256, k=6, seed=N)
|
||||
|
||||
|
||||
def test_fused_activation_topk_pro_experts():
|
||||
"""Test with 384 experts (Pro model)."""
|
||||
test_fused_activation_topk(N=64, E=384, k=6, seed=123)
|
||||
|
||||
|
||||
def test_hash_router(N=128, vocab_size=128000, k=6, num_experts=256, seed=42):
|
||||
"""Test the hash router against the math spec.
|
||||
|
||||
Oracle:
|
||||
topk_ids[n, h] = hash_lut[token_ids[n], h]
|
||||
topk_w[n, h] = 1.0 / k
|
||||
"""
|
||||
torch.manual_seed(seed)
|
||||
|
||||
# Build a random LUT
|
||||
hash_lut = torch.randint(0, num_experts, (vocab_size, k), dtype=torch.int32, device='cuda')
|
||||
token_ids = torch.randint(0, vocab_size, (N,), dtype=torch.int32, device='cuda')
|
||||
|
||||
# Oracle — literally just indexing
|
||||
expected_ids = hash_lut[token_ids] # [N, k]
|
||||
expected_w = torch.full((N, k), 1.0 / k, dtype=torch.float32, device='cuda')
|
||||
|
||||
# Kernel under test:
|
||||
from dsv4.kernels.router import hash_router_dispatch
|
||||
out_w = torch.empty(N, k, dtype=torch.float32, device='cuda')
|
||||
out_ids = torch.empty(N, k, dtype=torch.int32, device='cuda')
|
||||
hash_router_dispatch(token_ids, hash_lut, k, out_w, out_ids)
|
||||
|
||||
assert (out_ids == expected_ids).all(), f"hash router IDs mismatch"
|
||||
torch.testing.assert_close(out_w, expected_w, atol=1e-7, rtol=1e-7)
|
||||
|
||||
|
||||
def test_hash_router_edge_cases():
|
||||
"""Test hash router with N=1 and N=max_num_tokens."""
|
||||
test_hash_router(N=1, vocab_size=128000, k=6)
|
||||
test_hash_router(N=8192, vocab_size=128000, k=6)
|
||||
|
||||
|
||||
def test_topk_select(N=64, E=256, k=6, seed=42):
|
||||
"""Test standalone top-k selection against torch.topk.
|
||||
|
||||
Oracle:
|
||||
(values, indices) = score.topk(k, dim=-1)
|
||||
Lower index wins on ties (torch.topk default).
|
||||
"""
|
||||
torch.manual_seed(seed)
|
||||
scores = torch.randn(N, E, dtype=torch.float32, device='cuda')
|
||||
|
||||
# Oracle
|
||||
expected = scores.topk(k, dim=-1)
|
||||
expected_ids = expected.indices
|
||||
expected_values = expected.values
|
||||
|
||||
# Kernel under test:
|
||||
from dsv4.ops.topk import topk_select
|
||||
out_values, out_ids = topk_select(scores, k)
|
||||
|
||||
assert (out_ids == expected_ids).all(), f"top-k IDs mismatch"
|
||||
torch.testing.assert_close(out_values, expected_values, atol=1e-6, rtol=1e-6)
|
||||
|
||||
|
||||
def test_dense_router_decode(N=64, H=4096, E=256, k=6, seed=42):
|
||||
"""Test the full dense router (GEMM + activation + topk) against the spec.
|
||||
|
||||
Oracle:
|
||||
logits = (X.float() @ W.float())
|
||||
act = sqrt(softplus(logits))
|
||||
score = act + bias
|
||||
ids = score.topk(k).indices
|
||||
w = act.gather(-1, ids)
|
||||
w = w / w.sum(-1, keepdim=True) * scaling
|
||||
"""
|
||||
torch.manual_seed(seed)
|
||||
scaling = 2.5
|
||||
|
||||
X = torch.randn(N, H, dtype=torch.bfloat16, device='cuda')
|
||||
W = torch.randn(H, E, dtype=torch.bfloat16, device='cuda')
|
||||
bias = torch.randn(E, dtype=torch.float32, device='cuda') * 0.01
|
||||
|
||||
# Oracle — the math, in one expression, in FP32
|
||||
logits = (X.float() @ W.float())
|
||||
act = torch.sqrt(torch.nn.functional.softplus(logits))
|
||||
score = act + bias
|
||||
ids = score.topk(k, dim=-1).indices
|
||||
w = act.gather(-1, ids)
|
||||
w = w / w.sum(-1, keepdim=True) * scaling
|
||||
|
||||
# Kernel under test:
|
||||
from dsv4.layers.router import Router
|
||||
router = Router(H, E, k, scaling, mode='dense', max_num_tokens=N)
|
||||
router.load_weights(W_gate=W, e_bias=bias)
|
||||
router.finalize_weights()
|
||||
out_w, out_ids = router(X)
|
||||
|
||||
assert (out_ids == ids).all(), f"router IDs mismatch"
|
||||
torch.testing.assert_close(out_w, w, atol=1e-3, rtol=1e-3)
|
||||
|
||||
|
||||
def test_dense_router_decode_shapes():
|
||||
"""Test dense router at decode-relevant N values."""
|
||||
for N in [1, 4, 16, 64]:
|
||||
test_dense_router_decode(N=N, H=4096, E=256, k=6, seed=N)
|
||||
|
||||
|
||||
def test_hash_router_via_router_class():
|
||||
"""Test the Router class in hash mode."""
|
||||
vocab_size = 128000
|
||||
k = 6
|
||||
num_experts = 256
|
||||
N = 64
|
||||
|
||||
hash_lut = torch.randint(0, num_experts, (vocab_size, k), dtype=torch.int32, device='cuda')
|
||||
token_ids = torch.randint(0, vocab_size, (N,), dtype=torch.int32, device='cuda')
|
||||
|
||||
# Oracle
|
||||
expected_ids = hash_lut[token_ids]
|
||||
expected_w = torch.full((N, k), 1.0 / k, dtype=torch.float32, device='cuda')
|
||||
|
||||
# Router class
|
||||
from dsv4.layers.router import Router
|
||||
router = Router(
|
||||
hidden_size=4096, # not used in hash mode
|
||||
num_experts=num_experts,
|
||||
top_k=k,
|
||||
mode='hash',
|
||||
vocab_size=vocab_size,
|
||||
max_num_tokens=N,
|
||||
)
|
||||
router.load_weights(hash_lut=hash_lut)
|
||||
router.finalize_weights()
|
||||
out_w, out_ids = router(hidden_states=None, token_ids=token_ids)
|
||||
|
||||
assert (out_ids == expected_ids).all(), f"hash router class IDs mismatch"
|
||||
torch.testing.assert_close(out_w, expected_w, atol=1e-7, rtol=1e-7)
|
||||
|
||||
|
||||
def test_softplus_numerical_stability():
|
||||
"""Verify the numerically stable softplus matches the spec.
|
||||
|
||||
For x = -100: softplus(x) ≈ 0, sqrt(softplus(x)) ≈ 0
|
||||
For x = 0: softplus(x) = log(2) ≈ 0.693, sqrt ≈ 0.832
|
||||
For x = 100: softplus(x) ≈ 100, sqrt(softplus(x)) ≈ 10
|
||||
"""
|
||||
# This tests the Python math, not the kernel. It's a sanity check
|
||||
# that the formula max(x,0) + log1p(exp(-|x|)) works correctly.
|
||||
x = torch.tensor([-100.0, 0.0, 100.0], dtype=torch.float32)
|
||||
sp = torch.nn.functional.softplus(x)
|
||||
act = torch.sqrt(sp)
|
||||
expected = torch.tensor([0.0, math.sqrt(math.log(2.0)), 10.0], dtype=torch.float32)
|
||||
torch.testing.assert_close(act, expected, atol=1e-3, rtol=1e-3)
|
||||
Reference in New Issue
Block a user