Spark-H3 / tests /test_sol_layout.py

Publish Spark-H3 code and model card (part 3)

b78342b verified 17 days ago

3.5 kB

import sys,types
from pathlib import Path
import pytest
import torch
from torch import nn
import h3_sparse_attention.processor as proc

@pytest.mark.parametrize('fused',[False,True])
@pytest.mark.parametrize('context',[0,3])
@pytest.mark.parametrize('reblock',[False,True])
def test_sol_layout_roundtrip_exact(monkeypatch,fused,context,reblock):
    torch.manual_seed(42)
    n=128+context;heads=2;hidden=256
    perm=torch.randperm(n);inverse=perm.argsort()
    layout=proc.PackedLayout(perm,inverse,(2,8,8),128,n,torch.zeros(128,3,dtype=torch.long))
    config=proc.H3SparseAttentionConfig.sol(2,warmup_percent=0,sol_dense_layers=0,sol_log_density=False,sol_landmark_preprocess=reblock,sol_landmark_preprocess_version='v2')
    controller=proc._Controller(config);controller.layout=layout;controller.evaluation_index=0
    def fake_sol(q,k,v,**kw):
        assert q.is_contiguous() and k.is_contiguous() and v.is_contiguous()
        return q*.25+k*.5+v
    package=types.ModuleType('sol_attn');package.sol_attn=fake_sol;package.get_sol_attn_backend=lambda device:'test'
    monkeypatch.setitem(sys.modules,'sol_attn',package)
    if reblock:
        # Reblock dispatches through the CUDA-only Spark boundary.  This is a
        # processor layout test, so replace that boundary while retaining its
        # BTHD input/output contract; GPU kernel behavior is covered separately
        # by the Spark integration tests.
        integration=__import__(
            'h3_sparse_attention.spark_integration',
            fromlist=['spark_attention_bthd'],
        )
        def fake_spark_attention_bthd(
            controller, q, k, v, layout, layer, *, return_bthd=True
        ):
            assert q.is_contiguous() and k.is_contiguous() and v.is_contiguous()
            output=q*.25+k*.5+v
            return output if return_bthd else output.permute(0,2,1,3).contiguous()
        monkeypatch.setattr(
            integration, 'spark_attention_bthd', fake_spark_attention_bthd
        )
    linear=lambda size:nn.Linear(hidden,size,bias=False,dtype=torch.bfloat16)
    attn=types.SimpleNamespace(heads=heads,fused_projections=fused,to_qkv=linear(3*hidden),to_q=linear(hidden),to_k=linear(hidden),to_v=linear(hidden),norm_q=nn.Identity(),norm_k=nn.Identity(),to_out=[linear(hidden),nn.Identity()])
    processor=proc._H3SparseProcessor(0,None,controller)
    x=torch.randn(1,n,hidden,dtype=torch.bfloat16)
    native_attention = proc._sol_attention
    def checked_attention(
        controller, q, k, v, layout, layer, *, return_bthd=False,
        inputs_bthd=False,
    ):
        assert return_bthd
        assert not inputs_bthd
        assert all(t.permute(0,2,1,3).is_contiguous() for t in (q,k,v))
        return native_attention(
            controller, q, k, v, layout, layer,
            return_bthd=True, inputs_bthd=inputs_bthd,
        )
    with torch.no_grad():
        query,key,value = attn.to_qkv(x).chunk(3,-1) if fused else (attn.to_q(x),attn.to_k(x),attn.to_v(x))
        q,k,v = [t.unflatten(-1,(heads,-1)).index_select(1,perm).permute(0,2,1,3).contiguous()
                 for t in (query,key,value)]
        expected = native_attention(controller,q,k,v,layout,0).permute(0,2,1,3)
        expected = attn.to_out[1](attn.to_out[0](expected.index_select(1,inverse).flatten(2,3)))
        monkeypatch.setattr(proc,'_sol_attention',checked_attention)
        monkeypatch.setenv('H3_SOL_LAYOUT_FAST','0')
        after=processor(attn,x)
    assert torch.equal(expected,after)

	import sys,types
	from pathlib import Path
	import pytest
	import torch
	from torch import nn
	import h3_sparse_attention.processor as proc

	@pytest.mark.parametrize('fused',[False,True])
	@pytest.mark.parametrize('context',[0,3])
	@pytest.mark.parametrize('reblock',[False,True])
	def test_sol_layout_roundtrip_exact(monkeypatch,fused,context,reblock):
	torch.manual_seed(42)
	n=128+context;heads=2;hidden=256
	perm=torch.randperm(n);inverse=perm.argsort()
	layout=proc.PackedLayout(perm,inverse,(2,8,8),128,n,torch.zeros(128,3,dtype=torch.long))
	config=proc.H3SparseAttentionConfig.sol(2,warmup_percent=0,sol_dense_layers=0,sol_log_density=False,sol_landmark_preprocess=reblock,sol_landmark_preprocess_version='v2')
	controller=proc._Controller(config);controller.layout=layout;controller.evaluation_index=0
	def fake_sol(q,k,v,**kw):
	assert q.is_contiguous() and k.is_contiguous() and v.is_contiguous()
	return q.25+k.5+v
	package=types.ModuleType('sol_attn');package.sol_attn=fake_sol;package.get_sol_attn_backend=lambda device:'test'
	monkeypatch.setitem(sys.modules,'sol_attn',package)
	if reblock:
	# Reblock dispatches through the CUDA-only Spark boundary. This is a
	# processor layout test, so replace that boundary while retaining its
	# BTHD input/output contract; GPU kernel behavior is covered separately
	# by the Spark integration tests.
	integration=__import__(
	'h3_sparse_attention.spark_integration',
	fromlist=['spark_attention_bthd'],
	)
	def fake_spark_attention_bthd(
	controller, q, k, v, layout, layer, *, return_bthd=True
	):
	assert q.is_contiguous() and k.is_contiguous() and v.is_contiguous()
	output=q.25+k.5+v
	return output if return_bthd else output.permute(0,2,1,3).contiguous()
	monkeypatch.setattr(
	integration, 'spark_attention_bthd', fake_spark_attention_bthd
	)
	linear=lambda size:nn.Linear(hidden,size,bias=False,dtype=torch.bfloat16)
	attn=types.SimpleNamespace(heads=heads,fused_projections=fused,to_qkv=linear(3*hidden),to_q=linear(hidden),to_k=linear(hidden),to_v=linear(hidden),norm_q=nn.Identity(),norm_k=nn.Identity(),to_out=[linear(hidden),nn.Identity()])
	processor=proc._H3SparseProcessor(0,None,controller)
	x=torch.randn(1,n,hidden,dtype=torch.bfloat16)
	native_attention = proc._sol_attention
	def checked_attention(
	controller, q, k, v, layout, layer, *, return_bthd=False,
	inputs_bthd=False,
	):
	assert return_bthd
	assert not inputs_bthd
	assert all(t.permute(0,2,1,3).is_contiguous() for t in (q,k,v))
	return native_attention(
	controller, q, k, v, layout, layer,
	return_bthd=True, inputs_bthd=inputs_bthd,
	)
	with torch.no_grad():
	query,key,value = attn.to_qkv(x).chunk(3,-1) if fused else (attn.to_q(x),attn.to_k(x),attn.to_v(x))
	q,k,v = [t.unflatten(-1,(heads,-1)).index_select(1,perm).permute(0,2,1,3).contiguous()
	for t in (query,key,value)]
	expected = native_attention(controller,q,k,v,layout,0).permute(0,2,1,3)
	expected = attn.to_out[1](attn.to_out[0](expected.index_select(1,inverse).flatten(2,3)))
	monkeypatch.setattr(proc,'_sol_attention',checked_attention)
	monkeypatch.setenv('H3_SOL_LAYOUT_FAST','0')
	after=processor(attn,x)
	assert torch.equal(expected,after)