chore: improve ktransformers maintenance path

[lphuc2250gma]

Summary:

• Add or tighten focused edge-case tests or type assertions in kt-kernel/examples/test_mla.py, kt-kernel/examples/test_mlp.py related to Python typing, tests, CLI ergonomics, observability; avoid docs-only changes and broad refactors.
• Keep the change narrow so it is straightforward to review.

Notes:

• I kept this scoped to the relevant implementation and tests.

[gemini-code-assist]

Code Review

This pull request introduces type annotations to functions in test_mla.py and test_mlp.py, and enhances assertion error messages with formatted difference values. The review feedback points out that formatting 0-dimensional PyTorch tensors directly with float specifiers can cause issues, and suggests calling .item() to retrieve the underlying Python float first.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

The variable diff_relative_mean is a PyTorch Tensor (0-dimensional). Directly formatting a PyTorch tensor with float format specifiers like :.4e can raise a TypeError or produce unexpected formatting in older PyTorch versions. It is safer and more idiomatic to call .item() to retrieve the underlying Python float before formatting.

Suggested change

	assert diff_relative_mean < 2e-1, f"CPU and Torch outputs diverge: relative_mean={diff_relative_mean:.4e} >= 2e-1"
	assert diff_relative_mean < 2e-1, f"CPU and Torch outputs diverge: relative_mean={diff_relative_mean.item():.4e} >= 2e-1"

[gemini-code-assist]

The variable diff is a PyTorch Tensor (0-dimensional). Directly formatting a PyTorch tensor with float format specifiers like :.4e can raise a TypeError or produce unexpected formatting in older PyTorch versions. It is safer and more idiomatic to call .item() to retrieve the underlying Python float before formatting.

Suggested change

	assert diff < 0.001, f"MLP output mismatch: diff={diff:.4e} exceeds threshold 0.001"
	assert diff < 0.001, f"MLP output mismatch: diff={diff.item():.4e} exceeds threshold 0.001"