[Core] Support FP8 weight storage in unquantized linear and embedding layers

[lucaspirola]

Summary

• Add FP8 (float8_e4m3fn / float8_e5m2) dtype handling to UnquantizedLinearMethod and UnquantizedEmbeddingMethod
• Cast FP8 weights to compute dtype before GEMM operations and after embedding lookup
• Preserve FP8 dtype in VocabParallelEmbedding.weight_loader instead of silently upcasting to BF16

Motivation

This enables post-training FP8 compression of embed_tokens and lm_head layers to save VRAM without requiring a dedicated quantization scheme. For models with large vocabularies (131K+), this frees ~640 MB of VRAM that can be used for KV cache, significantly extending context length on memory-constrained GPUs.

Test plan

• Verified on GLM-4.7-Flash-REAP-23B-A3B with FP8 embed_tokens (saves 464 MB)
• Verified on Devstral-24B with FP8 embed_tokens (saves 640 MB)
• Pre-commit passes (ruff, mypy, typos)
• Unit tests for FP8 dtype path in both methods

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request adds support for FP8 weight storage in unquantized linear and embedding layers. This is achieved by casting FP8 weights to the compute dtype before GEMM operations and after embedding lookups, while preserving the FP8 dtype for weights in memory to save VRAM. The changes look mostly correct, but I've found a critical issue where the compute dtype is hardcoded when casting embedding outputs, which could lead to errors or incorrect behavior for models not using bfloat16.

[gemini-code-assist]

Hardcoding the cast to torch.bfloat16 is incorrect as it assumes the model's compute dtype is always bfloat16. This will cause dtype mismatches or incorrect results if the model is running with a different compute dtype, such as float16 or float32.

To ensure correctness across different model configurations, you should use torch.get_default_dtype() to dynamically get the model's compute dtype, which is set during model loading.

Suggested change

	output = output.to(torch.bfloat16)
	output = output.to(torch.get_default_dtype())