[feat] add CUTLASS kernel backend for HSTU attention

[tiankongdeguiji]

Summary

• Integrate the hstu_attn CUTLASS-based fused attention kernel as a new Kernel.CUTLASS backend for DlrmHSTU
• Rename unused Kernel.CUDA enum to Kernel.CUTLASS in both Python and protobuf definitions
• Add cutlass_hstu_mha wrapper in tzrec/ops/_cuda/ that adapts TorchEasyRec's interface to hstu_attn_varlen_func
• Supports Ampere (A100), Ada (L20), and Hopper (H100, H20) GPUs
• Requires attention_dim == hidden_dim; cached inference falls back to Triton

Test plan

• Unit tests: pytest tzrec/ops/hstu_attention_test.py -k "cutlass" -v
• Integration test: pytest tzrec/tests/rank_integration_test.py -k "dlrm_hstu_cutlass" -v
• Verify forward + backward numerical correctness against PyTorch reference

🤖 Generated with Claude Code

[github-actions]

Nit: bare assert is stripped under python -O

This constraint depends on user model config (attention_dim vs hidden_dim), so it should survive optimized mode. The rest of the codebase uses torch._assert for runtime validation. Consider switching for consistency:

Suggested change

	assert q.shape[2] == v.shape[2], (
	f"CUTLASS hstu_attn requires attention_dim == hidden_dim, "
	f"got q.shape[2]={q.shape[2]} != v.shape[2]={v.shape[2]}"
	)
	torch._assert(q.shape[2] == v.shape[2],
	f"CUTLASS hstu_attn requires attention_dim == hidden_dim, "
	f"got q.shape[2]={q.shape[2]} != v.shape[2]={v.shape[2]}"
	)

Also worth validating this earlier at config parse time so users get a clear error before any GPU work begins.

[github-actions]

Minor: int32 cast could silently wrap on very large cumulative offsets

If seq_offsets[-1] exceeds 2^31 - 1 (~2.1B tokens), the .to(torch.int32) silently wraps, producing negative offsets → out-of-bounds CUDA memory access. Unlikely in typical recommendation workloads, but a defensive check would prevent silent corruption:

if not is_fx_tracing():
    torch._assert(seq_offsets[-1] <= torch.iinfo(torch.int32).max,
                  "seq_offsets values exceed int32 range")

[github-actions]

Test name says "export" but doesn't test export

Other HSTU integration tests (e.g., the AOT variant above) call test_export after train+eval. Since CUTLASS introduces a new @torch.fx.wrap function, verifying export/tracing works correctly seems especially valuable. Either add the export step or rename to test_rank_dlrm_hstu_cutlass_train_eval.

[github-actions]

Bug: delta_hstu_mha bypasses the CUTLASS wrapper entirely

When kernel == Kernel.CUTLASS, this dispatches directly to triton_cached_hstu_mha, completely skipping the cutlass_cached_hstu_mha wrapper defined in cutlass_hstu_attention.py. This means:

• Users selecting CUTLASS get zero feedback that delta attention silently falls back to Triton (the warning in cutlass_cached_hstu_mha is dead code)
• The @torch.fx.wrap on cutlass_cached_hstu_mha is never invoked, which may affect FX tracing correctness
• If CUTLASS later adds delta support, this dispatch won't route to it

Should add a separate CUTLASS branch here (like hstu_mha does) that routes through cutlass_cached_hstu_mha:

if kernel == Kernel.CUTLASS:
    from tzrec.ops._cuda.cutlass_hstu_attention import cutlass_cached_hstu_mha
    return cutlass_cached_hstu_mha(...)
elif kernel == Kernel.TRITON:
    from tzrec.ops._triton.triton_hstu_attention import triton_cached_hstu_mha
    return triton_cached_hstu_mha(...)

[github-actions]

Code Review Summary

Nice work integrating the CUTLASS kernel backend! The wrapper design with graceful Triton fallback is well thought out, and the hypothesis-based tests provide good coverage. A few issues to address:

Must Fix

1. delta_hstu_mha bypasses cutlass_cached_hstu_mha wrapper (hstu_attention.py:160) — When kernel == Kernel.CUTLASS, the delta path dispatches directly to Triton, making the cutlass_cached_hstu_mha wrapper dead code. The fallback warning never fires, and @torch.fx.wrap is skipped. Should route through the CUTLASS wrapper like hstu_mha does.

Should Fix

2. Bare assert for attn_dim == hidden_dim (cutlass_hstu_attention.py:94) — Stripped under python -O. Use torch._assert or raise ValueError for consistency with the rest of the validation in this codebase.
3. Integration test name says "export" but doesn't test it (rank_integration_test.py:962) — Since CUTLASS adds a new @torch.fx.wrap function, testing export/tracing is especially valuable. Either add the export step or rename the test.
4. Docs mention cu126 but only cu129 wheels exist in extra.txt (dlrm_hstu.md:171) — Verify cu126 wheels are actually hosted, and check that the version string format (cu129 vs cu12.9) doesn't cause pip resolution issues.

Nice to Have

5. int32 cast of seq_offsets (cutlass_hstu_attention.py:99) — Could silently wrap on very large cumulative offsets. A defensive bounds check would prevent potential out-of-bounds CUDA memory access.
6. Consider using tzrec.utils.logging_util instead of stdlib logging to match project conventions.

🤖 Generated with Claude Code

[github-actions]

perf/correctness: The .contiguous() on seq_offsets (line 311) is wasted — the .to(torch.int32) on the next line always returns a new contiguous tensor regardless.

More importantly, the int32 cast silently overflows if cumulative offsets exceed INT32_MAX. Consider adding a bounds check:

Suggested change

	seq_offsets = seq_offsets.contiguous()
	cu_seqlens = seq_offsets.to(torch.int32)
	cu_seqlens = seq_offsets.to(torch.int32)

And optionally guard with:

if seq_offsets.max() > torch.iinfo(torch.int32).max:
    raise ValueError(f"cu_seqlens exceed int32 range: max={seq_offsets.max().item()}")

[github-actions]

perf: .reshape() on a non-contiguous input (from the slice) already returns a new contiguous tensor, so the trailing .contiguous() is a no-op. If the shape isn't actually changing (N is already the first dim), this can be simplified:

Suggested change

	return out[:, :, :head_dim].reshape(-1, num_heads, head_dim).contiguous()
	return out[:, :, :head_dim].contiguous()

[github-actions]

correctness: The meta function returns torch.empty_like(q), which happens to be correct today because cutlass_hstu_mha enforces q.shape[2] == v.shape[2]. However, semantically the output shape derives from v's hidden_dim, not q's attn_dim. If the constraint is ever relaxed, this meta would silently return the wrong shape during tracing.

Consider using torch.empty_like(v) (or at minimum mirroring the q.shape[2] == v.shape[2] assertion here) to make the contract explicit.

[github-actions]

nit: The rest of the codebase uses from tzrec.utils.logging_util import logger. Using logging.getLogger(__name__) here creates a separate logger that won't inherit project-level formatting/handlers.

Suggested change

	logger = logging.getLogger(__name__)
	from tzrec.utils.logging_util import logger

[github-actions]

test coverage / naming: delta_hstu_mha unconditionally falls back CUTLASS→Triton here, which means test_delta_attn_cutlass in the test file is actually testing Triton-vs-PyTorch equivalence — no CUTLASS code path is exercised. Consider either:

1. Renaming the test to test_delta_attn_cutlass_fallback to make this clear, or
2. Adding a logger.warning here (matching the pattern in hstu_mha) so the fallback is observable.

[github-actions]

Code Review Summary

Well-structured PR. The low-level torch.library.define/impl approach (instead of @torch.library.custom_op) is well-motivated by the AOTI deadlock, and the rationale comment at the top of cutlass_hstu_attention.py is excellent — it will prevent future contributors from "simplifying" it back and reintroducing the bug. The DenseAutocastWrapper with Final[int] encoding for torch.jit.script is a clean solution.

Inline comments posted (6)

File	Topic
cutlass_hstu_attention.py:311-312	Redundant .contiguous() before .to(int32) + silent int32 overflow risk
cutlass_hstu_attention.py:94	Redundant .contiguous() after .reshape() on non-contiguous input
cutlass_hstu_attention.py:157	Meta function uses empty_like(q) — fragile if attn_dim == hidden_dim constraint is relaxed
hstu_attention.py:29	Logger should use tzrec.utils.logging_util.logger per project convention
hstu_attention.py:165-166	delta_hstu_mha CUTLASS→Triton fallback is silent + test name is misleading
docs/dlrm_hstu.md:171	Docs claim cu126 support but no cu126 wheels exist

Additional observations (not posted inline)

• Maintainability: The if kernel == Kernel.CUTLASS: kernel = Kernel.TRITON guard is repeated 13+ times across layer_norm.py, jagged_tensors.py, mm.py, position.py, hstu_compute.py. Consider centralizing this in Kernel or a utility (e.g. resolve_kernel(kernel, op_name)) so new ops/kernels don't require touching every file.

• Test coverage gaps: No unit tests for DenseAutocastWrapper, mixed_precision_to_dtype(), or resolve_mixed_precision(). Also no tests for the CUTLASS input validation error paths (attn_dim != hidden_dim, unsupported dtype). These are small and cheap to add.

• FP16 test coverage: test_attn_cutlass only samples torch.bfloat16. The production code also supports torch.float16 — adding it to the sampled dtypes would improve coverage at minimal cost.